Add tests to run KFTO pytorch MNIST training using multi-node/multi-g…

…pu usecases
opendatahub-io · Jan 6, 2025 · 48c546c · 48c546c
1 parent 0dcc478
commit 48c546c
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/tests/kfto/kfto_mnist_training_test.go b/tests/kfto/kfto_mnist_training_test.go
@@ -30,17 +30,26 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
-func TestPyTorchJobMnistCpu(t *testing.T) {
+func TestPyTorchJobMnistMultiNodeCpu(t *testing.T) {
 	runKFTOPyTorchMnistJob(t, 0, 2, "", GetCudaTrainingImage(), "resources/requirements.txt")
 }
-func TestPyTorchJobMnistWithCuda(t *testing.T) {
+
+func TestPyTorchJobMnistMultiNodeWithCuda(t *testing.T) {
 	runKFTOPyTorchMnistJob(t, 1, 1, "nvidia.com/gpu", GetCudaTrainingImage(), "resources/requirements.txt")
 }
 
-func TestPyTorchJobMnistWithROCm(t *testing.T) {
+func TestPyTorchJobMnistMultiNodeWithROCm(t *testing.T) {
 	runKFTOPyTorchMnistJob(t, 1, 1, "amd.com/gpu", GetROCmTrainingImage(), "resources/requirements-rocm.txt")
 }
 
+func TestPyTorchJobMnistMultiNodeMultiGpuWithCuda(t *testing.T) {
+	runKFTOPyTorchMnistJob(t, 2, 1, "nvidia.com/gpu", GetCudaTrainingImage(), "resources/requirements.txt")
+}
+
+func TestPyTorchJobMnistMultiNodeMultiGpuWithROCm(t *testing.T) {
+	runKFTOPyTorchMnistJob(t, 2, 1, "amd.com/gpu", GetROCmTrainingImage(), "resources/requirements-rocm.txt")
+}
+
 func runKFTOPyTorchMnistJob(t *testing.T, numGpus int, workerReplicas int, gpuLabel string, image string, requirementsFile string) {
 	test := With(t)