Add GPU to ClusterConfiguration

oksanabaza · May 2, 2024 · 039bb29 · 039bb29
1 parent a9b314e
commit 039bb29
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/demo-notebooks/guided-demos/0_basic_ray.ipynb b/demo-notebooks/guided-demos/0_basic_ray.ipynb
@@ -61,7 +61,8 @@
     "# Create and configure our cluster object\n",
     "# The SDK will try to find the name of your default local queue based on the annotation \"kueue.x-k8s.io/default-queue\": \"true\" unless you specify the local queue manually below\n",
     "cluster = Cluster(ClusterConfiguration(\n",
-    "    name='raytest',\n",
+    "    name='raytest', \n",
+    "    head_gpus=1, # If you are using GPU, this needs to be specified \n",
     "    namespace='default', # Update to your namespace\n",
     "    num_workers=2,\n",
     "    min_cpus=1,\n",
@@ -71,7 +72,7 @@
     "    num_gpus=0,\n",
     "    image=\"quay.io/project-codeflare/ray:latest-py39-cu118\",\n",
     "    write_to_file=False, # When enabled Ray Cluster yaml files are written to /HOME/.codeflare/resources \n",
-    "    # local_queue=\"local-queue-name\" # Specify the local queue manually\n",
+    "    local_queue=\"local-queue-name\" # Specify the local queue manually\n",
     "))"
    ]
   },

diff --git a/demo-notebooks/guided-demos/2_basic_interactive.ipynb b/demo-notebooks/guided-demos/2_basic_interactive.ipynb
@@ -61,6 +61,7 @@
     "cluster_name = \"interactivetest\"\n",
     "cluster = Cluster(ClusterConfiguration(\n",
     "    name=cluster_name,\n",
+    "    head_gpus=1, # If you are using GPU, this needs to be specified \n",
     "    namespace=namespace,\n",
     "    num_workers=2,\n",
     "    min_cpus=2,\n",