fix

sky/provision/runpod/instance.py

@@ -11,6 +11,7 @@
 from sky.utils import ux_utils
 
 POLL_INTERVAL = 5
+QUERY_PORTS_TIMEOUT_SECONDS = 30
 
 logger = sky_logging.init_logger(__name__)
 
@@ -224,11 +225,19 @@ def query_ports(
 ) -> Dict[int, List[common.Endpoint]]:
     """See sky/provision/__init__.py"""
     del head_ip, provider_config  # Unused.
-    instances = _filter_instances(cluster_name_on_cloud, None, head_only=True)
-    assert len(instances) == 1
-    head_inst = list(instances.values())[0]
-    return {
-        port: [common.SocketEndpoint(**endpoint)]
-        for port, endpoint in head_inst['port2endpoint'].items()
-        if port in resources_utils.port_ranges_to_set(ports)
-    }
+    # RunPod ports sometimes take a while to be ready.
+    start_time = time.time()
+    while True:
+        instances = _filter_instances(cluster_name_on_cloud,
+                                      None,
+                                      head_only=True)
+        assert len(instances) == 1
+        head_inst = list(instances.values())[0]
+        if (all(port in head_inst['port2endpoint'] for port in ports) or
+                time.time() - start_time > QUERY_PORTS_TIMEOUT_SECONDS):
+            return {
+                port: [common.SocketEndpoint(**endpoint)]
+                for port, endpoint in head_inst['port2endpoint'].items()
+                if port in resources_utils.port_ranges_to_set(ports)
+            }
+        time.sleep(1)