Add SSH keys to GCP metadata (#881)

(cherry picked from commit 1d5964d)
dstackai · Feb 7, 2024 · 3530693 · 3530693
1 parent 4bbb2f7
commit 3530693
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/src/dstack/_internal/core/backends/gcp/compute.py b/src/dstack/_internal/core/backends/gcp/compute.py
@@ -93,6 +93,10 @@ def run_job(
             project_id=self.config.project_id,
         )
         disk_size = round(instance_offer.instance.resources.disk.size_mib / 1024)
+        authorized_keys = [
+            run.run_spec.ssh_key_pub.strip(),
+            project_ssh_public_key.strip(),
+        ]
         for zone in _get_instance_zones(instance_offer):
             request = compute_v1.InsertInstanceRequest()
             request.zone = zone
@@ -112,12 +116,10 @@ def run_job(
                 user_data=get_user_data(
                     backend=BackendType.GCP,
                     image_name=job.job_spec.image_name,
-                    authorized_keys=[
-                        run.run_spec.ssh_key_pub.strip(),
-                        project_ssh_public_key.strip(),
-                    ],
+                    authorized_keys=authorized_keys,
                     registry_auth_required=job.job_spec.registry_auth is not None,
                 ),
+                authorized_keys=authorized_keys,
                 labels={
                     "owner": "dstack",
                     "dstack_project": project_id,
@@ -179,6 +181,7 @@ def create_gateway(
             accelerators=[],
             spot=False,
             user_data=get_gateway_user_data(ssh_key_pub),
+            authorized_keys=[ssh_key_pub],
             labels={
                 "owner": "dstack",
                 "dstack_project": project_id,

diff --git a/src/dstack/_internal/core/backends/gcp/resources.py b/src/dstack/_internal/core/backends/gcp/resources.py
@@ -30,6 +30,7 @@ def create_instance_struct(
     accelerators: List[compute_v1.AcceleratorConfig],
     spot: bool,
     user_data: str,
+    authorized_keys: List[str],
     labels: Dict[str, str],
     tags: List[str],
     instance_name: str,
@@ -75,6 +76,9 @@ def create_instance_struct(
 
     metadata_items = [
         compute_v1.Items(key="user-data", value=user_data),
+        compute_v1.Items(
+            key="ssh-keys", value="\n".join(f"ubuntu:{key}" for key in authorized_keys)
+        ),
     ]
     instance.metadata = compute_v1.Metadata(items=metadata_items)