raise node deprature exceptions - local-launched clusters

Author: Alexandra Belousov

runhouse/constants.py

@@ -117,3 +117,4 @@
 SSH_SKY_SECRET_NAME = "ssh-sky-key"
 
 INSUFFICIENT_DISK_MSG = "No space left on device"
+CONNECTION_REFUSED_MSG = "Connection refused"

runhouse/exceptions.py

@@ -25,3 +25,16 @@ def __init__(
         )
         msg = f"{msg}. To resolve it, teardown the cluster and re-launch it with larger disk size."
         super().__init__(msg)
+
+
+class ClusterTerminated(Exception):
+    """Raised when we try to connect to a terminated cluster"""
+
+    def __init__(
+        self,
+        error_msg: str = None,
+    ) -> None:
+        self.error_msg = error_msg
+        msg = self.error_msg if self.error_msg else "Cluster is terminated"
+        msg = f"{msg}. To resolve it, please bring the cluster up by using the Runhouse's python APIs or CLI commands."
+        super().__init__(msg)

runhouse/resources/hardware/cluster.py

@@ -26,6 +26,7 @@
     get_running_and_not_running_clusters,
     get_unsaved_live_clusters,
     parse_filters,
+    RunhouseDaemonStatus,
 )
 
 from runhouse.resources.images import Image, ImageSetupStepType
@@ -68,6 +69,7 @@
     NUM_PORTS_TO_TRY,
     RESERVED_SYSTEM_NAMES,
 )
+from runhouse.exceptions import ClusterTerminated
 from runhouse.globals import configs, obj_store, rns_client
 from runhouse.logger import get_logger
 
@@ -1043,6 +1045,32 @@ def status(self, send_to_den: bool = False):
 
         return status
 
+    def _update_cluster_status_to_terminated_in_den(self):
+        if self.rns_address:
+            try:
+                # Update Den with the terminated status
+                status_data = {
+                    "daemon_status": RunhouseDaemonStatus.TERMINATED,
+                    "resource_type": self.__class__.__base__.__name__.lower(),
+                    "data": {},
+                }
+
+                cluster_uri = rns_client.format_rns_address(self.rns_address)
+                status_resp = requests.post(
+                    f"{rns_client.api_server_url}/resource/{cluster_uri}/cluster/status",
+                    json=status_data,
+                    headers=rns_client.request_headers(),
+                )
+
+                # Note: 404 means that the cluster is not saved in Den
+                if status_resp.status_code not in [200, 404]:
+                    logger.warning(
+                        "Failed to update Den with terminated cluster status"
+                    )
+
+            except Exception as e:
+                logger.warning(e)
+
     def ssh_tunnel(
         self, local_port, remote_port=None, num_ports_to_try: int = 0
     ) -> "SshTunnel":
@@ -1488,6 +1516,11 @@ def rsync(
         if not src_node and up and not Path(source).expanduser().exists():
             raise ValueError(f"Could not locate path to sync: {source}.")
 
+        # The cluster is probably auto-stopped / terminated via CLI while calling cluster.rsync
+        if not self.is_up():
+            self._update_cluster_status_to_terminated_in_den()
+            raise ClusterTerminated()
+
         # before syncing the object, making sure there is enough disk space for it. If we don't preform this check in
         # advance, there might be a case where we start rsyncing the object -> during the rsync no disk space will
         # be left (because the object is too big) -> the rsync will stack and no results will be return (even failures),
@@ -1613,14 +1646,19 @@ def rsync(
                 logger.info(f"Rsyncing {source} on {node} to {dest}")
                 Path(dest).expanduser().parent.mkdir(parents=True, exist_ok=True)
 
-            runner.rsync(
-                source,
-                dest,
-                up=up,
-                filter_options=filter_options,
-                stream_logs=stream_logs,
-                ignore_existing=ignore_existing,
-            )
+            try:
+                runner.rsync(
+                    source,
+                    dest,
+                    up=up,
+                    filter_options=filter_options,
+                    stream_logs=stream_logs,
+                    ignore_existing=ignore_existing,
+                )
+            except ClusterTerminated as e:
+                if not self._ping():
+                    self._update_cluster_status_to_terminated_in_den()
+                    raise e
             return
 
         # Case 3: rsync with password between cluster and local (e.g. laptop)
@@ -1636,16 +1674,21 @@ def rsync(
             logger.info(f"Rsyncing {source} on {node} to {dest}")
             Path(dest).expanduser().parent.mkdir(parents=True, exist_ok=True)
 
-        rsync_cmd = runner.rsync(
-            source,
-            dest,
-            up=up,
-            filter_options=filter_options,
-            stream_logs=stream_logs,
-            return_cmd=True,
-            ignore_existing=ignore_existing,
-        )
-        run_command_with_password_login(rsync_cmd, pwd, stream_logs)
+        try:
+            rsync_cmd = runner.rsync(
+                source,
+                dest,
+                up=up,
+                filter_options=filter_options,
+                stream_logs=stream_logs,
+                return_cmd=True,
+                ignore_existing=ignore_existing,
+            )
+            run_command_with_password_login(rsync_cmd, pwd, stream_logs)
+        except ClusterTerminated as e:
+            if not self._ping():
+                self._update_cluster_status_to_terminated_in_den()
+                raise e
 
     def _local_rsync(
         self,
@@ -2153,6 +2196,18 @@ def run_python(
             conda_env_name=conda_env_name,
         )
 
+        if node == "all":
+            status_code = set([return_code[0][0] for return_code in return_codes])
+            python_run_failed = len(status_code) > 1 or status_code.pop() != 0
+        else:
+            status_code = return_codes[0][0]
+            python_run_failed = status_code != 0
+        if python_run_failed:
+            if not self._ping():
+                self._update_cluster_status_to_terminated_in_den()
+                raise ClusterTerminated()
+            raise ValueError(f"Could not run python code on cluster, got {status_code}")
+
         return return_codes
 
     def create_conda_env(self, conda_env_name: str, conda_config: Dict):

runhouse/resources/hardware/on_demand_cluster.py

@@ -7,8 +7,6 @@
 from pathlib import Path
 from typing import Any, Dict, List, Union
 
-import requests
-
 import rich.errors
 
 try:
@@ -32,7 +30,6 @@
     ClusterStatus,
     LauncherType,
     pprint_launched_cluster_summary,
-    RunhouseDaemonStatus,
     ServerConnectionType,
     up_cluster_helper,
 )
@@ -684,29 +681,7 @@ def teardown(self, verbose: bool = True):
             LocalLauncher.teardown(cluster=self, verbose=verbose)
 
         if self.rns_address is not None:
-            try:
-                # Update Den with the terminated status
-                status_data = {
-                    "daemon_status": RunhouseDaemonStatus.TERMINATED,
-                    "resource_type": self.__class__.__base__.__name__.lower(),
-                    "data": {},
-                }
-
-                cluster_uri = rns_client.format_rns_address(self.rns_address)
-                status_resp = requests.post(
-                    f"{rns_client.api_server_url}/resource/{cluster_uri}/cluster/status",
-                    json=status_data,
-                    headers=rns_client.request_headers(),
-                )
-
-                # Note: 404 means that the cluster is not saved in Den
-                if status_resp.status_code not in [200, 404]:
-                    logger.warning(
-                        "Failed to update Den with terminated cluster status"
-                    )
-
-            except Exception as e:
-                logger.warning(e)
+            self._update_cluster_status_to_terminated_in_den()
 
     def teardown_and_delete(self, verbose: bool = True):
         """Teardown cluster and delete it from configs.

runhouse/resources/hardware/sky/subprocess_utils.py

@@ -4,8 +4,8 @@
 
 from runhouse.logger import get_logger
 
-from runhouse.exceptions import InsufficientDisk
-from runhouse.constants import INSUFFICIENT_DISK_MSG
+from runhouse.exceptions import InsufficientDisk, ClusterTerminated
+from runhouse.constants import INSUFFICIENT_DISK_MSG, CONNECTION_REFUSED_MSG
 
 logger = get_logger(__name__)
 
@@ -50,6 +50,8 @@ def handle_returncode(returncode: int,
 
         if INSUFFICIENT_DISK_MSG in stderr:
             raise InsufficientDisk(command=command, error_msg=error_msg)
+        if CONNECTION_REFUSED_MSG in stderr:
+            raise ClusterTerminated()
         if callable(error_msg):
             error_msg = error_msg()
         raise CommandError(returncode, command, error_msg, stderr)

runhouse/servers/http/http_client.py

@@ -13,6 +13,7 @@
 
 import requests
 
+from runhouse.exceptions import ClusterTerminated
 from runhouse.globals import rns_client
 from runhouse.logger import get_logger
 
@@ -199,6 +200,12 @@ def request(
             headers=headers,
         )
 
+    def _ping_cluster(self):
+        try:
+            self.check_server()
+        except ValueError:
+            raise ClusterTerminated()
+
     def request_json(
         self,
         endpoint: str,
@@ -250,12 +257,17 @@ def request_json(
             )
         resp_json = response.json()
         if isinstance(resp_json, dict) and "output_type" in resp_json:
-            return handle_response(
-                resp_json,
-                resp_json["output_type"],
-                err_str,
-                log_formatter=self.log_formatter,
-            )
+            try:
+                return handle_response(
+                    resp_json,
+                    resp_json["output_type"],
+                    err_str,
+                    log_formatter=self.log_formatter,
+                )
+            except ConnectionError:
+                self._ping_cluster()
+            except requests.exceptions.ConnectionError:
+                self._ping_cluster()
         return resp_json
 
     def check_server(self):