Download MNIST dataset from specific location

Author: sutaakar

test/e2e/mnist.py

@@ -15,6 +15,7 @@
 import os
 
 import torch
+import requests
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.callbacks.progress import TQDMProgressBar
 from torch import nn
@@ -32,6 +33,8 @@
 print("MASTER_ADDR: is ", os.getenv("MASTER_ADDR"))
 print("MASTER_PORT: is ", os.getenv("MASTER_PORT"))
 
+print("MNIST_DATASET_URL: is ", os.getenv("MNIST_DATASET_URL"))
+MNIST_DATASET_URL = os.getenv("MNIST_DATASET_URL")
 
 class LitMNIST(LightningModule):
     def __init__(self, data_dir=PATH_DATASETS, hidden_size=64, learning_rate=2e-4):
@@ -110,8 +113,34 @@ def configure_optimizers(self):
     ####################
 
     def prepare_data(self):
-        # download
-        print("Downloading MNIST dataset...")
+        datasetFiles = [
+            "t10k-images-idx3-ubyte",
+            "t10k-labels-idx1-ubyte",
+            "train-images-idx3-ubyte",
+            "train-labels-idx1-ubyte"
+        ]
+
+        # Create required folder structure
+        downloadLocation = os.path.join(PATH_DATASETS, "MNIST", "raw")
+        os.makedirs(downloadLocation, exist_ok=True)
+        print(f"{downloadLocation} folder_path created!")
+
+        for file in datasetFiles:
+            print(f"Downloading MNIST dataset {file}... to path : {downloadLocation}")
+            response = requests.get(f"{MNIST_DATASET_URL}{file}", stream=True)
+            filePath = os.path.join(downloadLocation, file)
+
+            #to download dataset file
+            try:
+                if response.status_code == 200:
+                    open(filePath, 'wb').write(response.content)
+                    print(f"{file}: Downloaded and saved zipped file to path - {filePath}")
+                else:
+                    print(f"Failed to download file {file}")
+            except Exception as e:
+                print(e)
+        print(f"Downloaded MNIST dataset to... {downloadLocation}")
+
         MNIST(self.data_dir, train=True, download=True)
         MNIST(self.data_dir, train=False, download=True)
 

test/e2e/mnist_pytorch_mcad_job_test.go

@@ -79,7 +79,8 @@ func TestMNISTPyTorchMCAD(t *testing.T) {
 							Name:  "job",
 							Image: GetPyTorchImage(),
 							Env: []corev1.EnvVar{
-								corev1.EnvVar{Name: "PYTHONUSERBASE", Value: "/workdir"},
+								{Name: "PYTHONUSERBASE", Value: "/workdir"},
+								{Name: "MNIST_DATASET_URL", Value: GetMnistDatasetURL()},
 							},
 							Command: []string{"/bin/sh", "-c", "pip install -r /test/requirements.txt && torchrun /test/mnist.py"},
 							VolumeMounts: []corev1.VolumeMount{

test/e2e/mnist_rayjob_mcad_raycluster_test.go

@@ -17,7 +17,6 @@ limitations under the License.
 package e2e
 
 import (
-	"encoding/base64"
 	"testing"
 
 	. "github.com/onsi/gomega"
@@ -143,13 +142,6 @@ func TestMNISTRayJobMCADRayCluster(t *testing.T) {
 					RayStartParams: map[string]string{},
 					Template: corev1.PodTemplateSpec{
 						Spec: corev1.PodSpec{
-							InitContainers: []corev1.Container{
-								{
-									Name:    "init-myservice",
-									Image:   "busybox:1.28",
-									Command: []string{"sh", "-c", "until nslookup $RAY_IP.$(cat /var/run/secrets/kubernetes.io/serviceaccount/namespace).svc.cluster.local; do echo waiting for myservice; sleep 2; done"},
-								},
-							},
 							Containers: []corev1.Container{
 								{
 									Name:  "ray-worker",
@@ -230,21 +222,29 @@ func TestMNISTRayJobMCADRayCluster(t *testing.T) {
 		},
 		Spec: rayv1.RayJobSpec{
 			Entrypoint: "python /home/ray/jobs/mnist.py",
-			RuntimeEnv: base64.StdEncoding.EncodeToString([]byte(`
-{
-  "pip": [
-    "pytorch_lightning==1.5.10",
-    "torchmetrics==0.9.1",
-    "torchvision==0.12.0"
-  ],
-  "env_vars": {
-  }
-}
-`)),
+			RuntimeEnvYAML: `
+  pip:
+    - pytorch_lightning==1.5.10
+    - torchmetrics==0.9.1
+    - torchvision==0.12.0
+  env_vars:
+    MNIST_DATASET_URL: "` + GetMnistDatasetURL() + `"
+`,
 			ClusterSelector: map[string]string{
 				RayJobDefaultClusterSelectorKey: rayCluster.Name,
 			},
 			ShutdownAfterJobFinishes: false,
+			SubmitterPodTemplate: &corev1.PodTemplateSpec{
+				Spec: corev1.PodSpec{
+					RestartPolicy: corev1.RestartPolicyNever,
+					Containers: []corev1.Container{
+						{
+							Image: GetRayImage(),
+							Name:  "rayjob-submitter-pod",
+						},
+					},
+				},
+			},
 		},
 	}
 	rayJob, err = test.Client().Ray().RayV1().RayJobs(namespace.Name).Create(test.Ctx(), rayJob, metav1.CreateOptions{})
@@ -256,6 +256,10 @@ func TestMNISTRayJobMCADRayCluster(t *testing.T) {
 	test.T().Logf("Connecting to Ray cluster at: %s", rayDashboardURL.String())
 	rayClient := NewRayClusterClient(rayDashboardURL)
 
+	// Wait for Ray job id to be available, this value is needed for writing logs in defer
+	test.Eventually(RayJob(test, rayJob.Namespace, rayJob.Name), TestTimeoutShort).
+		Should(WithTransform(RayJobId, Not(BeEmpty())))
+
 	// Retrieving the job logs once it has completed or timed out
 	defer WriteRayJobAPILogs(test, rayClient, GetRayJobId(test, rayJob.Namespace, rayJob.Name))
 

test/upgrade/olm_upgrade_test.go

@@ -96,7 +96,8 @@ func TestMNISTCreateAppWrapper(t *testing.T) {
 							Name:  "job",
 							Image: GetPyTorchImage(),
 							Env: []corev1.EnvVar{
-								corev1.EnvVar{Name: "PYTHONUSERBASE", Value: "/workdir"},
+								{Name: "PYTHONUSERBASE", Value: "/workdir"},
+								{Name: "MNIST_DATASET_URL", Value: GetMnistDatasetURL()},
 							},
 							Command: []string{"/bin/sh", "-c", "pip install -r /test/requirements.txt && torchrun /test/mnist.py"},
 							VolumeMounts: []corev1.VolumeMount{