Review feedback

Author: Maxusmusti

demo-notebooks/batch-job/batch_mnist_mcad.ipynb

@@ -8,7 +8,8 @@
    "outputs": [],
    "source": [
     "# Import pieces from codeflare-sdk\n",
-    "from codeflare_sdk.cluster.auth import TokenAuthentication"
+    "from codeflare_sdk.cluster.auth import TokenAuthentication\n",
+    "from codeflare_sdk.job.jobs import DDPJobDefinition"
    ]
   },
   {
@@ -45,16 +46,6 @@
     "Now that we are logged in, we can directly submit our batch job (model training on two workers with four gpus each) to MCAD via torchx."
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "id": "3cc6183a-8f6e-4347-af91-d088ed422544",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from codeflare_sdk.job.jobs import DDPJobDefinition"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 6,
@@ -70,7 +61,17 @@
     }
    ],
    "source": [
-    "job = DDPJobDefinition(name=\"mnistjob\", script=\"mnist.py\", scheduler_args={\"namespace\": \"default\"}, j=\"1x1\", gpu=0, cpu=1, memMB=8000, image=\"quay.io/michaelclifford/mnist-test:latest\").submit()"
+    "job = DDPJobDefinition(\n",
+    "    name=\"mnistjob\",\n",
+    "    script=\"mnist.py\",\n",
+    "    scheduler_args={\"namespace\": \"default\"},\n",
+    "    j=\"1x1\",\n",
+    "    gpu=0,\n",
+    "    cpu=1,\n",
+    "    memMB=8000,\n",
+    "    image=\"quay.io/michaelclifford/mnist-test:latest\"\n",
+    ")\n",
+    "job.submit()"
    ]
   },
   {

demo-notebooks/batch-job/batch_mnist_ray.ipynb

@@ -9,7 +9,8 @@
    "source": [
     "# Import pieces from codeflare-sdk\n",
     "from codeflare_sdk.cluster.cluster import Cluster, ClusterConfiguration\n",
-    "from codeflare_sdk.cluster.auth import TokenAuthentication"
+    "from codeflare_sdk.cluster.auth import TokenAuthentication\n",
+    "from codeflare_sdk.job.jobs import DDPJobDefinition"
    ]
   },
   {
@@ -51,8 +52,20 @@
     }
    ],
    "source": [
-    "# Create our cluster and submit appwrapper\n",
-    "cluster = Cluster(ClusterConfiguration(name='mnisttest', min_worker=2, max_worker=2, min_cpus=8, max_cpus=8, min_memory=16, max_memory=16, gpu=4, instascale=True, machine_types=[\"m5.xlarge\", \"p3.8xlarge\"]))"
+    "# Create our cluster and submit appwrapper (reduce specs as desired)\n",
+    "cluster = Cluster(ClusterConfiguration(\n",
+    "    name='mnisttest',\n",
+    "    namespace='default',\n",
+    "    min_worker=2,\n",
+    "    max_worker=2,\n",
+    "    min_cpus=8,\n",
+    "    max_cpus=8,\n",
+    "    min_memory=16,\n",
+    "    max_memory=16,\n",
+    "    gpu=4,\n",
+    "    instascale=True, # Can be set to false if scaling not needed\n",
+    "    machine_types=[\"m5.xlarge\", \"g4dn.xlarge\"] # Can be removed if above is false\n",
+    "))"
    ]
   },
   {
@@ -293,16 +306,6 @@
     "Now that our resource cluster is ready, we can directly submit our batch job (model training on two workers with four gpus each) to the cluster via torchx."
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "id": "3cc6183a-8f6e-4347-af91-d088ed422544",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from codeflare_sdk.job.jobs import DDPJobDefinition"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 10,
@@ -320,7 +323,11 @@
     }
    ],
    "source": [
-    "job = DDPJobDefinition(script=\"mnist.py\", scheduler_args={\"requirements\": \"requirements.txt\"}).submit(cluster)"
+    "job = DDPJobDefinition(\n",
+    "    script=\"mnist.py\",\n",
+    "    scheduler_args={\"requirements\": \"requirements.txt\"}\n",
+    ")\n",
+    "job.submit(cluster)"
    ]
   },
   {