Merge remote-tracking branch 'upstream/main'

Author: dchourasia

tests/kfto/kfto_mnist_training_test.go

@@ -155,7 +155,7 @@ func runKFTOPyTorchMnistJob(t *testing.T, accelerator Accelerator, image string,
 	defer test.Client().Kubeflow().KubeflowV1().PyTorchJobs(namespace.Name).Delete(test.Ctx(), tuningJob.Name, *metav1.NewDeleteOptions(0))
 
 	// Make sure the Workload is created and running
-	test.Eventually(GetKueueWorkloads(test, namespace.Name), TestTimeoutMedium).
+	test.Eventually(KueueWorkloads(test, namespace.Name), TestTimeoutMedium).
 		Should(
 			And(
 				HaveLen(1),
@@ -238,7 +238,7 @@ func createKFTOPyTorchMnistJob(test Test, namespace string, config corev1.Config
 										python3 /mnt/files/download_mnist_datasets.py --dataset_path "/tmp/datasets/mnist" && \
 										echo -e "\n\n Dataset downloaded to /tmp/datasets/mnist" && ls -R /tmp/datasets/mnist && \
 										echo -e "\n\n Starting training..." && \
-										torchrun --nproc_per_node=%d /mnt/files/mnist.py --dataset_path "/tmp/datasets/mnist" --epochs 7 --save_every 2 --batch_size 128 --lr 0.001 --snapshot_path "mnist_snapshot.pt" --backend %s`, numProcPerNode, backend),
+										torchrun --nproc_per_node=%d /mnt/files/mnist.py --dataset_path "/tmp/datasets/mnist" --epochs 3 --save_every 2 --batch_size 128 --lr 0.001 --snapshot_path "mnist_snapshot.pt" --backend %s`, numProcPerNode, backend),
 									},
 									VolumeMounts: []corev1.VolumeMount{
 										{