fix ping-and-up test for den launched clusters

runhouse/resources/hardware/cluster.py

@@ -213,7 +213,7 @@ def check_connect_server():
 
         if not self._http_client and not check_connect_server():
             if self.__class__.__name__ == "OnDemandCluster":
-                self._update_from_sky_status(dryrun=False)
+                self._update_from_status(dryrun=False)
             if not self._ping(retry=False):
                 raise ConnectionError(
                     f"Could not reach {self.name} {self.head_ip}. Is cluster up?"

runhouse/resources/hardware/launcher_utils.py

@@ -409,7 +409,7 @@ def up(cls, cluster, verbose: bool = True):
                 **cluster.sky_kwargs.get("launch", {}),
             )
 
-            cluster._update_from_sky_status()
+            cluster._update_from_status()
             if cluster.domain:
                 logger.info(
                     f"Cluster has been launched with the custom domain '{cluster.domain}'. "

runhouse/resources/hardware/on_demand_cluster.py

@@ -24,7 +24,6 @@
     DEFAULT_SERVER_PORT,
     LOCAL_HOSTS,
 )
-
 from runhouse.globals import configs, obj_store, rns_client
 from runhouse.logger import get_logger
 from runhouse.resources.hardware.utils import (
@@ -35,6 +34,7 @@
     ServerConnectionType,
     up_cluster_helper,
 )
+from runhouse.rns.utils.api import read_resp_data
 from .cluster import Cluster
 from .launcher_utils import DenLauncher, LocalLauncher
 
@@ -141,7 +141,7 @@ def __init__(
         # Checks if state info is in local sky db, populates if so.
         if not dryrun and not self.ips and self.launcher == LauncherType.LOCAL:
             # Cluster status is set to INIT in the Sky DB right after starting, so we need to refresh once
-            self._update_from_sky_status(dryrun=True)
+            self._update_from_status(dryrun=True)
 
     @property
     def ips(self):
@@ -158,7 +158,7 @@ def client(self):
         except ValueError as e:
             if not self.ips:
                 # Try loading in from local Sky DB
-                self._update_from_sky_status(dryrun=True)
+                self._update_from_status(dryrun=True)
                 if not self.ips:
                     raise ValueError(
                         f"Could not determine ips for ondemand cluster <{self.name}>. "
@@ -402,7 +402,7 @@ def _sky_status(self, refresh: bool = True, retry: bool = True):
 
     def _start_ray_workers(self, ray_port, env_vars):
         if not self.internal_ips:
-            self._update_from_sky_status()
+            self._update_from_status()
 
         super()._start_ray_workers(ray_port, env_vars)
 
@@ -523,18 +523,6 @@ def _populate_connection_from_status_dict(self, cluster_dict: Dict[str, Any]):
                 self._kube_namespace = self.compute_properties.get("kube_namespace")
                 self._kube_context = self.compute_properties.get("kube_context")
 
-    def _update_from_sky_status(self, dryrun: bool = False):
-        if self.launcher != LauncherType.LOCAL:
-            return
-
-        # Try to get the cluster status from SkyDB
-        if self.is_shared:
-            # If the cluster is shared can ignore, since the sky data will only be saved on the machine where
-            # the cluster was initially upped
-            return
-
-        cluster_dict = self._sky_status(refresh=not dryrun)
-        self._populate_connection_from_status_dict(cluster_dict)
 
     def _setup_default_creds(self):
         if self.launcher == LauncherType.DEN:
@@ -544,6 +532,45 @@ def _setup_default_creds(self):
         else:
             raise ValueError(f"Invalid launcher '{self.launcher}'")
 
+    def _update_from_status(self, dryrun: bool = False):
+
+        if self.launcher == LauncherType.LOCAL:
+            # Try to get the cluster status from SkyDB
+            if self.is_shared:
+                # If the cluster is shared can ignore, since the sky data will only be saved on the machine where
+                # the cluster was initially upped
+                return
+
+            cluster_dict = self._sky_status(refresh=not dryrun)
+            self._populate_connection_from_status_dict(cluster_dict)
+        if self.launcher == LauncherType.DEN:
+            # checking if the function is called during the init process, before the cluster is saved to den.
+            if not self.rns_address:
+                return
+            cluster_uri = rns_client.format_rns_address(self.rns_address)
+            resp = requests.get(
+                f"{rns_client.api_server_url}/resource/{cluster_uri}/cluster/status",
+                headers=rns_client.request_headers(),
+            )
+            if resp.status_code != 200:
+                logger.warning(
+                    f"Received [{resp.status_code}] from Den GET '{cluster_uri}': Failed to load cluster status"
+                )
+                return
+            cluster_den_status: list = read_resp_data(resp)
+
+            if not cluster_den_status:
+                logger.warning("Failed to update cluster info")
+                return
+
+            resource_info = cluster_den_status[0].get("resource_info")
+            if not resource_info:
+                logger.warning("Failed to load resource info from cluster status")
+                return
+
+            cluster_dict = resource_info.get("cluster_config")
+            DenLauncher._update_from_den_response(cluster=self, config=cluster_dict)
+
     def get_instance_type(self):
         """Returns instance type of the cluster."""
         if self.instance_type and "--" in self.instance_type:  # K8s specific syntax
@@ -815,6 +842,7 @@ def _ping(self, timeout=5, retry=False):
             return True
 
         if retry:
-            self._update_from_sky_status(dryrun=False)
+            dryrun = False if self.launcher == LauncherType.LOCAL else None
+            self._update_from_status(dryrun=dryrun)
             return super()._ping(timeout=timeout, retry=False)
         return False