Fix bootstrapping on GCP (#172)

* Fix bootstrapping on GCP

* Also bootstrap GPU drivers

* Fix GPU bootstrapping

* Add ubuntu-drivers package

* Install drivers non-interactive

dask_cloudprovider/gcp/instances.py

@@ -58,6 +58,7 @@ def __init__(
         ngpus=None,
         gpu_type=None,
         bootstrap=None,
+        gpu_instance=None,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -76,7 +77,7 @@ def __init__(
         self.filesystem_size = filesystem_size or self.config.get("filesystem_size")
         self.ngpus = ngpus or self.config.get("ngpus")
         self.gpu_type = gpu_type or self.config.get("gpu_type")
-        self.gpu_instance = "gpu" in self.machine_type or bool(self.ngpus)
+        self.gpu_instance = gpu_instance
         self.bootstrap = bootstrap
 
         self.general_zone = "-".join(self.zone.split("-")[:2])  # us-east1-c -> us-east1
@@ -491,7 +492,7 @@ def __init__(
         gpu_type=None,
         filesystem_size=None,
         auto_shutdown=None,
-        boostrap=True,
+        bootstrap=True,
         **kwargs,
     ):
 
@@ -505,6 +506,11 @@ def __init__(
         )
         self.scheduler_class = GCPScheduler
         self.worker_class = GCPWorker
+        self.bootstrap = (
+            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)
         self.options = {
             "cluster": self,
             "config": self.config,
@@ -513,10 +519,11 @@ def __init__(
             "docker_image": docker_image or self.config.get("docker_image"),
             "filesystem_size": filesystem_size or self.config.get("filesystem_size"),
             "zone": zone or self.config.get("zone"),
-            "machine_type": machine_type or self.config.get("machine_type"),
+            "machine_type": self.machine_type,
             "ngpus": ngpus or self.config.get("ngpus"),
             "gpu_type": gpu_type or self.config.get("gpu_type"),
-            "bootstrap": boostrap,
+            "gpu_instance": self.gpu_instance,
+            "bootstrap": self.bootstrap,
         }
         self.scheduler_options = {**self.options}
         self.worker_options = {**self.options}

dask_cloudprovider/gcp/tests/test_gcp.py

@@ -59,6 +59,7 @@ async def test_get_cloud_init():
 
     cloud_init = GCPCluster.get_cloud_init()
     assert "dask-scheduler" in cloud_init
+    assert "# Bootstrap" in cloud_init
 
 
 @pytest.mark.asyncio

dask_cloudprovider/generic/cloud-init.yaml.j2

@@ -1,12 +1,14 @@
 #cloud-config
 
 {% if bootstrap %}
+# Bootstrap
 packages:
   - apt-transport-https
   - ca-certificates
   - curl
   - gnupg-agent
   - software-properties-common
+  - ubuntu-drivers-common
 
 # Enable ipv4 forwarding, required on CIS hardened machines
 write_files:
@@ -24,13 +26,6 @@ system_info:
     groups: [docker]
 {% endif %}
 
-{% if bootstrap and gpu_instance %}
-# Install NVIDIA driver if GPU instance
-drivers:
-  nvidia:
-    license-accepted: true
-{% endif %}
-
 runcmd:
   {% if bootstrap %}
   # Install Docker
@@ -42,6 +37,18 @@ runcmd:
   - systemctl enable docker
   {% endif %}
 
+  {% if bootstrap and gpu_instance %}
+  # Install NVIDIA driver
+  - DEBIAN_FRONTEND=noninteractive ubuntu-drivers install
+
+  # Install NVIDIA docker
+  - curl -fsSL https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
+  - curl -s -L https://nvidia.github.io/nvidia-docker/$(. /etc/os-release;echo $ID$VERSION_ID)/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
+  - apt-get update -y
+  - apt-get install -y nvidia-docker2
+  - systemctl restart docker
+  {% endif %}
+
   # Run container
   - 'docker run --net=host {%+ if gpu_instance %}--gpus=all{% endif %} {% for key in env_vars %} -e {{key}}={{env_vars[key]}} {% endfor %}{{image}} {{ command }}'