Merge pull request #274 from jacobtomlinson/gcp-configure-disk

[GCP] Add disk type config option

dask_cloudprovider/cloudprovider.yaml

@@ -86,6 +86,7 @@ cloudprovider:
     filesystem_size: 50 # amount in GBs of hard drive space to allocate
     ngpus: "" # number of GPUs to use
     gpu_type: "" # type of gpus to use: nvidia-tesla-k80, nvidia-tesla-p100, nvidia-tesla-t4
+    disk_type: "pd-standard" # type of disk to use: pd-standard, pd-ssd
     docker_image: "daskdev/dask:latest" # docker image to use
     auto_shutdown: true # Shutdown instances automatically if the scheduler or worker services time out.
     public_ingress: true # configure the scheduler to be externally accessible.  This assumes firefwall rules for 8787 and 8786

dask_cloudprovider/gcp/instances.py

@@ -52,6 +52,7 @@ def __init__(
         projectid=None,
         machine_type=None,
         filesystem_size=None,
+        disk_type=None,
         on_host_maintenance=None,
         source_image=None,
         docker_image=None,
@@ -83,6 +84,7 @@ def __init__(
         self.docker_image = docker_image or self.config.get("docker_image")
         self.env_vars = env_vars
         self.filesystem_size = filesystem_size or self.config.get("filesystem_size")
+        self.disk_type = disk_type or self.config.get("disk_type")
         self.ngpus = ngpus or self.config.get("ngpus")
         self.network = network or self.config.get("network")
         self.gpu_type = gpu_type or self.config.get("gpu_type")
@@ -110,7 +112,7 @@ def create_gcp_config(self):
                     "deviceName": self.name,
                     "initializeParams": {
                         "sourceImage": self.source_image,
-                        "diskType": f"projects/{self.projectid}/zones/{self.zone}/diskTypes/pd-standard",
+                        "diskType": f"projects/{self.projectid}/zones/{self.zone}/diskTypes/{self.disk_type}",
                         "diskSizeGb": f"{self.filesystem_size}",  # nvidia-gpu-cloud cannot be smaller than 32 GB
                         "labels": {},
                         # "source": "projects/nv-ai-infra/zones/us-east1-c/disks/ngc-gpu-dask-rapids-docker-experiment",
@@ -288,6 +290,7 @@ async def start_scheduler(self):
             f"\n  Docker Image: {self.docker_image} "
             f"\n  Machine Type: {self.machine_type} "
             f"\n  Filesytsem Size: {self.filesystem_size} "
+            f"\n  Disk Type: {self.disk_type} "
             f"\n  N-GPU Type: {self.ngpus} {self.gpu_type}"
             f"\n  Zone: {self.zone} "
         )
@@ -429,6 +432,9 @@ class GCPCluster(VMCluster):
         You can see a list of GPUs available in each zone with ``gcloud compute accelerator-types list``.
     filesystem_size: int (optional)
         The VM filesystem size in GB. Defaults to ``50``.
+    disk_type: str (optional)
+        Type of disk to use. Default is ``pd-standard``.
+        You can see a list of disks available in each zone with ``gcloud compute disk-types list``.
     on_host_maintenance: str (optional)
         The Host Maintenance GCP option.  Defaults to ``TERMINATE``.
     n_workers: int (optional)
@@ -543,6 +549,7 @@ def __init__(
         ngpus=None,
         gpu_type=None,
         filesystem_size=None,
+        disk_type=None,
         auto_shutdown=None,
         bootstrap=True,
         preemptible=None,
@@ -571,6 +578,7 @@ def __init__(
             "source_image": source_image or self.config.get("source_image"),
             "docker_image": docker_image or self.config.get("docker_image"),
             "filesystem_size": filesystem_size or self.config.get("filesystem_size"),
+            "disk_type": disk_type or self.config.get("disk_type"),
             "on_host_maintenance": on_host_maintenance
             or self.config.get("on_host_maintenance"),
             "zone": zone or self.config.get("zone"),