From 353069325b1c5b3fdd2c9cf361ffa4a8fd36a2a8 Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Wed, 7 Feb 2024 17:57:10 +0100
Subject: [PATCH] Add SSH keys to GCP metadata (#881)

(cherry picked from commit 1d5964d537c9bc90eb19c9a0acbfd1ee36fb3a37)
---
 src/dstack/_internal/core/backends/gcp/compute.py   | 11 +++++++----
 src/dstack/_internal/core/backends/gcp/resources.py |  4 ++++
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/src/dstack/_internal/core/backends/gcp/compute.py b/src/dstack/_internal/core/backends/gcp/compute.py
index 8942f44e1..aba7f4c86 100644
--- 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
index bea0370f0..f6919f956 100644
--- 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)