Merge remote-tracking branch 'upstream/main'

red-hat-data-services · Jan 8, 2025 · 8440348 · 8440348
2 parents 8b7bea5 + a439296
commit 8440348
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)