Fix gateway in GCP

dstackai · Aug 21, 2023 · 00d554c · 00d554c
1 parent f940280
commit 00d554c
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 10 deletions.
diff --git a/cli/dstack/_internal/backend/base/pricing.py b/cli/dstack/_internal/backend/base/pricing.py
@@ -34,7 +34,6 @@ def get_prices(
         self.fetch(instances, spot)
         offers = []
         for instance in instances:
-            print(self.instance_key(instance))
             for (region, is_spot), price in self.registry[self.instance_key(instance)].items():
                 if not self.region_match(instance.available_regions, region):
                     continue

diff --git a/cli/dstack/_internal/backend/gcp/compute.py b/cli/dstack/_internal/backend/gcp/compute.py
@@ -190,13 +190,14 @@ def cancel_spot_request(self, runner: Runner):
         )
 
     def create_gateway(self, instance_name: str, ssh_key_pub: str) -> GatewayHead:
+        region = self.gcp_config.regions[0]
         instance = gateway.create_gateway_instance(
             instances_client=self.instances_client,
             firewalls_client=self.firewalls_client,
             project_id=self.gcp_config.project_id,
             network=_get_network_resource(self.gcp_config.vpc),
-            subnet=_get_subnet_resource(self.gcp_config.region, self.gcp_config.subnet),
-            zone=self.gcp_config.zone,
+            subnet=_get_subnet_resource(region, self.gcp_config.subnet),
+            zone=_get_zones(self.regions_client, self.gcp_config.project_id, [region])[0],
             instance_name=instance_name,
             service_account=self.credentials.service_account_email,
             labels=dict(
@@ -945,12 +946,19 @@ def _restart_instance(
 def _terminate_instance(
     client: compute_v1.InstancesClient, gcp_config: GCPConfig, instance_name: str
 ):
-    _delete_instance(
-        client=client,
-        instance_name=instance_name,
-        project_id=gcp_config.project_id,
-        zone=gcp_config.zone,
-    )
+    request = compute_v1.AggregatedListInstancesRequest()
+    request.project = gcp_config.project_id
+    request.filter = f"name eq {instance_name}"
+    request.max_results = 1
+    agg_list = client.aggregated_list(request=request)
+    for zone, response in agg_list:
+        for _ in response.instances:
+            _delete_instance(
+                client=client,
+                instance_name=instance_name,
+                project_id=gcp_config.project_id,
+                zone=zone[6:],  # strip zones/
+            )
 
 
 def _delete_instance(

diff --git a/cli/dstack/_internal/backend/gcp/pricing.py b/cli/dstack/_internal/backend/gcp/pricing.py
@@ -12,7 +12,6 @@ def _fetch_catalog(self):
             if int(row["gpu_count"]) > 0:
                 gpu_memory = round(float(row["gpu_memory"]) * 1024)
                 instance_key += f'-{row["gpu_count"]}x{row["gpu_name"]}-{gpu_memory}'
-            print(instance_key)
             self.registry[instance_key][(row["location"][:-2], row["spot"] == "True")] = float(
                 row["price"]
             )