Fix issue 390 support different machine types on gcp

dask_cloudprovider/cloudprovider.yaml

@@ -101,7 +101,9 @@ cloudprovider:
     network_projectid: null # GCP project id where the network exists
     projectid: "" # name of the google cloud project
     on_host_maintenance: "TERMINATE"
-    machine_type: "n1-standard-1" # size of the machine type to use
+    machine_type: "n1-standard-1" # size of the machine type to use for the scheduler and all workers
+    scheduler_machine_type: "n1-standard-1" # size of the machine type to use for the scheduler
+    worker_machine_type: "n1-standard-1" # size of the machine type to use for all workers
     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

dask_cloudprovider/gcp/instances.py

@@ -417,7 +417,15 @@ class GCPCluster(VMCluster):
         be cases (i.e. Shared VPC) when network configurations from a different GCP project are used.
     machine_type: str
         The VM machine_type. You can get a full list with ``gcloud compute machine-types list``.
-        The default is ``n1-standard-1`` which is 3.75GB RAM and 1 vCPU
+        The default is ``n1-standard-1`` which is 3.75GB RAM and 1 vCPU.
+        This will determine the resources available to both the sceduler and all workers.
+        If supplied, you may not specify ``scheduler_machine_type`` or ``worker_machine_type``.
+    scheduler_machine_type: str
+        The VM machine_type. This will determine the resources available to the scheduler.
+        The default is ``n1-standard-1`` which is 3.75GB RAM and 1 vCPU.
+    worker_machine_type: str
+        The VM machine_type. This will determine the resources available to all workers.
+        The default is ``n1-standard-1`` which is 3.75GB RAM and 1 vCPU.
     source_image: str
         The OS image to use for the VM. Dask Cloudprovider will boostrap Ubuntu based images automatically.
         Other images require Docker and for GPUs the NVIDIA Drivers and NVIDIA Docker.
@@ -445,10 +453,11 @@ class GCPCluster(VMCluster):
     extra_bootstrap: list[str] (optional)
         Extra commands to be run during the bootstrap phase.
     ngpus: int (optional)
-        The number of GPUs to atatch to the instance.
+        The number of GPUs to atatch to the worker instance. No work is expected to be done on scheduler, so no
+        GPU there.
         Default is ``0``.
     gpu_type: str (optional)
-        The name of the GPU to use. This must be set if ``ngpus>0``.
+        The name of the GPU to use on worker. This must be set if ``ngpus>0``.
         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``.
@@ -573,6 +582,8 @@ def __init__(
         network=None,
         network_projectid=None,
         machine_type=None,
+        scheduler_machine_type=None,
+        worker_machine_type=None,
         on_host_maintenance=None,
         source_image=None,
         docker_image=None,
@@ -603,7 +614,16 @@ def __init__(
             bootstrap if bootstrap is not None else self.config.get("bootstrap")
         )
         self.machine_type = machine_type or self.config.get("machine_type")
-        self.gpu_instance = "gpu" in self.machine_type or bool(ngpus)
+        if machine_type is None:
+            self.scheduler_machine_type = scheduler_machine_type or self.config.get("scheduler_machine_type")
+            self.worker_machine_type = worker_machine_type or self.config.get("worker_machine_type")
+            if self.scheduler_machine_type is None or self.worker_machine_type is None:
+                raise ValueError("machine_type and scheduler_machine_type must be set")
+        else:
+            if scheduler_machine_type is not None or worker_machine_type is not None:
+                raise ValueError("If you specify machine_type, you may not specify scheduler_machine_type or worker_machine_type")
+            self.scheduler_machine_type = machine_type
+            self.worker_machine_type = machine_type
         self.debug = debug
         self.options = {
             "cluster": self,
@@ -617,6 +637,8 @@ def __init__(
             or self.config.get("on_host_maintenance"),
             "zone": zone or self.config.get("zone"),
             "machine_type": self.machine_type,
+            "scheduler_machine_type": self.scheduler_machine_type,
+            "worker_machine_type": self.worker_machine_type,
             "ngpus": ngpus or self.config.get("ngpus"),
             "network": network or self.config.get("network"),
             "network_projectid": network_projectid
@@ -635,6 +657,18 @@ def __init__(
         }
         self.scheduler_options = {**self.options}
         self.worker_options = {**self.options}
+        self.scheduler_options["machine_type"] = self.scheduler_machine_type
+        self.worker_options["machine_type"] = self.worker_machine_type
+
+        # Scheduler always does not have GPUs as no work is expected to be done there
+        self.scheduler_options["ngpus"] = 0
+        self.scheduler_options["gpu_type"] = None
+        self.scheduler_options["gpu_instance"] = False
+
+        if ngpus or self.config.get("ngpus"):
+            self.worker_options["ngpus"] = ngpus or self.config.get("ngpus")
+            self.worker_options["gpu_type"] = gpu_type or self.config.get("gpu_type")
+            self.worker_options["gpu_instance"] = True
 
         if "extra_bootstrap" not in kwargs:
             kwargs["extra_bootstrap"] = self.config.get("extra_bootstrap")