Install default_env on all nodes

runhouse/constants.py

@@ -59,6 +59,9 @@
     "bash ~/miniconda.sh -b -p ~/miniconda",
     "source $HOME/miniconda/bin/activate",
 ]
+# TODO should default to user's local Python version?
+# from platform import python_version; python_version()
+CONDA_PREFERRED_PYTHON_VERSION = "3.10.9"
 
 TEST_ORG = "test-org"
 TESTING_LOG_LEVEL = "INFO"

runhouse/resources/envs/conda_env.py

@@ -5,7 +5,7 @@
 
 import yaml
 
-from runhouse.constants import ENVS_DIR
+from runhouse.constants import CONDA_PREFERRED_PYTHON_VERSION, ENVS_DIR
 from runhouse.globals import obj_store
 from runhouse.logger import get_logger
 
@@ -59,18 +59,21 @@ def config(self, condensed=True):
     def env_name(self):
         return self.conda_yaml["name"]
 
-    def _create_conda_env(self, force: bool = False, cluster: "Cluster" = None):
+    def _create_conda_env(
+        self, force: bool = False, cluster: "Cluster" = None, node: Optional[str] = None
+    ):
         yaml_path = Path(ENVS_DIR) / f"{self.env_name}.yml"
 
         env_exists = (
             f"\n{self.env_name} "
-            in run_setup_command("conda info --envs", cluster=cluster)[1]
+            in run_setup_command("conda info --envs", cluster=cluster, node=node)[1]
         )
-        run_setup_command(f"mkdir -p {ENVS_DIR}", cluster=cluster)
+        run_setup_command(f"mkdir -p {ENVS_DIR}", cluster=cluster, node=node)
         yaml_exists = (
             (Path(ENVS_DIR).expanduser() / f"{self.env_name}.yml").exists()
             if not cluster
-            else run_setup_command(f"ls {yaml_path}", cluster=cluster)[0] == 0
+            else run_setup_command(f"ls {yaml_path}", cluster=cluster, node=node)[0]
+            == 0
         )
 
         if force or not (yaml_exists and env_exists):
@@ -87,19 +90,25 @@ def _create_conda_env(self, force: bool = False, cluster: "Cluster" = None):
                 subprocess.run(f'python -c "{python_commands}"', shell=True)
             else:
                 contents = yaml.dump(self.conda_yaml)
-                run_setup_command(f"echo $'{contents}' > {yaml_path}", cluster=cluster)
+                run_setup_command(
+                    f"echo $'{contents}' > {yaml_path}", cluster=cluster, node=node
+                )
 
             # create conda env from yaml file
-            run_setup_command(f"conda env create -f {yaml_path}", cluster=cluster)
+            run_setup_command(
+                f"conda env create -f {yaml_path}", cluster=cluster, node=node
+            )
 
             env_exists = (
                 f"\n{self.env_name} "
-                in run_setup_command("conda info --envs", cluster=cluster)[1]
+                in run_setup_command("conda info --envs", cluster=cluster, node=node)[1]
             )
             if not env_exists:
                 raise RuntimeError(f"conda env {self.env_name} not created properly.")
 
-    def install(self, force: bool = False, cluster: "Cluster" = None):
+    def install(
+        self, force: bool = False, cluster: "Cluster" = None, node: Optional[str] = None
+    ):
         """Locally install packages and run setup commands.
 
         Args:
@@ -109,32 +118,42 @@ def install(self, force: bool = False, cluster: "Cluster" = None):
                 on the cluster using SSH. (default: ``None``)
         """
         if not any(["python" in dep for dep in self.conda_yaml["dependencies"]]):
-            status_codes = run_setup_command("python --version", cluster=cluster)
+            status_codes = run_setup_command(
+                "python --version", cluster=cluster, node=node
+            )
             base_python_version = (
-                status_codes[1].split()[1] if status_codes[0] == 0 else "3.10.9"
+                status_codes[1].split()[1]
+                if status_codes[0] == 0
+                else CONDA_PREFERRED_PYTHON_VERSION
             )
             self.conda_yaml["dependencies"].append(f"python=={base_python_version}")
         install_conda(cluster=cluster)
         local_env_exists = (
             f"\n{self.env_name} "
-            in run_setup_command("conda info --envs", cluster=cluster)[1]
+            in run_setup_command("conda info --envs", cluster=cluster, node=node)[1]
         )
 
-        # Hash the config_for_rns to check if we need to create/install the conda env
-        env_config = self.config()
-        # Remove the name because auto-generated names will be different, but the installed components are the same
-        env_config.pop("name")
-        install_hash = hash(str(env_config))
-        # Check the existing hash
-        if local_env_exists and install_hash in obj_store.installed_envs and not force:
-            logger.debug("Env already installed, skipping")
-            return
-        obj_store.installed_envs[install_hash] = self.name
-
-        self._create_conda_env(force=force, cluster=cluster)
-
-        self._install_reqs(cluster=cluster)
-        self._run_setup_cmds(cluster=cluster)
+        # If we're doing the install remotely via SSH (e.g. for default_env), there is no cache
+        if not cluster:
+            # Hash the config_for_rns to check if we need to create/install the conda env
+            env_config = self.config()
+            # Remove the name because auto-generated names will be different, but the installed components are the same
+            env_config.pop("name")
+            install_hash = hash(str(env_config))
+            # Check the existing hash
+            if (
+                local_env_exists
+                and install_hash in obj_store.installed_envs
+                and not force
+            ):
+                logger.debug("Env already installed, skipping")
+                return
+            obj_store.installed_envs[install_hash] = self.name
+
+        self._create_conda_env(force=force, cluster=cluster, node=node)
+
+        self._install_reqs(cluster=cluster, node=node)
+        self._run_setup_cmds(cluster=cluster, node=node)
 
         return
 

runhouse/resources/envs/env.py

@@ -140,7 +140,9 @@ def _secrets_to(self, system: Union[str, Cluster]):
             new_secrets.append(secret.to(system=system, env=self))
         return new_secrets
 
-    def _install_reqs(self, cluster: Cluster = None, reqs: List = None):
+    def _install_reqs(
+        self, cluster: Cluster = None, reqs: List = None, node: str = "all"
+    ):
         reqs = reqs or self.reqs
         if reqs:
             for package in reqs:
@@ -154,9 +156,11 @@ def _install_reqs(self, cluster: Cluster = None, reqs: List = None):
                     raise ValueError(f"package {package} not recognized")
 
                 logger.debug(f"Installing package: {str(pkg)}")
-                pkg._install(env=self, cluster=cluster)
+                pkg._install(env=self, cluster=cluster, node=node)
 
-    def _run_setup_cmds(self, cluster: Cluster = None, setup_cmds: List = None):
+    def _run_setup_cmds(
+        self, cluster: Cluster = None, setup_cmds: List = None, node: str = "all"
+    ):
         setup_cmds = setup_cmds or self.setup_cmds
 
         if not setup_cmds:
@@ -165,31 +169,37 @@ def _run_setup_cmds(self, cluster: Cluster = None, setup_cmds: List = None):
         for cmd in setup_cmds:
             cmd = self._full_command(cmd)
             run_setup_command(
-                cmd, cluster=cluster, env_vars=_process_env_vars(self.env_vars)
+                cmd,
+                cluster=cluster,
+                env_vars=_process_env_vars(self.env_vars),
+                node=node,
             )
 
-    def install(self, force: bool = False, cluster: Cluster = None):
+    def install(self, force: bool = False, cluster: Cluster = None, node: str = "all"):
         """Locally install packages and run setup commands.
 
         Args:
             force (bool, optional): Whether to setup the installation again if the env already exists
                 on the cluster. (Default: ``False``)
             cluster (Clsuter, optional): Cluster to install the env on. If not provided, env is installed
                 on the current cluster. (Default: ``None``)
+            node (str, optional): Node to install the env on. (Default: ``"all"``)
         """
-        # Hash the config_for_rns to check if we need to install
-        env_config = self.config()
-        # Remove the name because auto-generated names will be different, but the installed components are the same
-        env_config.pop("name")
-        install_hash = hash(str(env_config))
-        # Check the existing hash
-        if install_hash in obj_store.installed_envs and not force:
-            logger.debug("Env already installed, skipping")
-            return
-        obj_store.installed_envs[install_hash] = self.name
-
-        self._install_reqs(cluster=cluster)
-        self._run_setup_cmds(cluster=cluster)
+        # If we're doing the install remotely via SSH (e.g. for default_env), there is no cache
+        if not cluster:
+            # Hash the config_for_rns to check if we need to install
+            env_config = self.config()
+            # Remove the name because auto-generated names will be different, but the installed components are the same
+            env_config.pop("name")
+            install_hash = hash(str(env_config))
+            # Check the existing hash
+            if install_hash in obj_store.installed_envs and not force:
+                logger.debug("Env already installed, skipping")
+                return
+            obj_store.installed_envs[install_hash] = self.name
+
+        self._install_reqs(cluster=cluster, node=node)
+        self._run_setup_cmds(cluster=cluster, node=node)
 
     def _full_command(self, command: str):
         if self._run_cmd:
@@ -205,7 +215,7 @@ def _run_command(self, command: str, **kwargs):
     def to(
         self,
         system: Union[str, Cluster],
-        node_idx: int = None,
+        node_idx: Optional[int] = None,
         path: str = None,
         force_install: bool = False,
     ):
@@ -227,19 +237,21 @@ def to(
             >>> s3_env = env.to("s3", path="s3_bucket/my_env")
         """
         system = _get_cluster_from(system)
+        if (
+            isinstance(system, Cluster)
+            and node_idx is not None
+            and node_idx >= len(system.ips)
+        ):
+            raise ValueError(
+                f"Cluster {system.name} has only {len(system.ips)} nodes. Requested node index {node_idx} is out of bounds."
+            )
+
         new_env = copy.deepcopy(self)
         new_env.reqs, new_env.working_dir = self._reqs_to(system, path)
 
         if isinstance(system, Cluster):
             if node_idx is not None:
-                if node_idx >= len(system.ips):
-                    raise ValueError(
-                        f"Cluster {system.name} has only {len(system.ips)} nodes. Requested node index {node_idx} is out of bounds."
-                    )
-
-                if new_env.compute is None:
-                    new_env.compute = {}
-
+                new_env.compute = new_env.compute or {}
                 new_env.compute["node_idx"] = node_idx
 
             key = (

runhouse/resources/envs/utils.py

@@ -2,7 +2,7 @@
 import subprocess
 
 from pathlib import Path
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import yaml
 
@@ -131,6 +131,7 @@ def run_setup_command(
     cluster: "Cluster" = None,
     env_vars: Dict = None,
     stream_logs: bool = True,
+    node: Optional[str] = None,
 ):
     """
     Helper function to run a command during possibly the cluster default env setup. If a cluster is provided,
@@ -152,14 +153,14 @@ def run_setup_command(
         return run_with_logs(cmd, stream_logs=stream_logs, require_outputs=True)[:2]
 
     return cluster._run_commands_with_runner(
-        [cmd], stream_logs=stream_logs, env_vars=env_vars
+        [cmd], stream_logs=stream_logs, env_vars=env_vars, node=node
     )[0]
 
 
-def install_conda(cluster: "Cluster" = None):
-    if run_setup_command("conda --version", cluster=cluster)[0] != 0:
+def install_conda(cluster: "Cluster" = None, node: Optional[str] = None):
+    if run_setup_command("conda --version", cluster=cluster, node=node)[0] != 0:
         logging.info("Conda is not installed. Installing...")
         for cmd in CONDA_INSTALL_CMDS:
-            run_setup_command(cmd, cluster=cluster, stream_logs=True)
-        if run_setup_command("conda --version", cluster=cluster)[0] != 0:
+            run_setup_command(cmd, cluster=cluster, node=node, stream_logs=True)
+        if run_setup_command("conda --version", cluster=cluster, node=node)[0] != 0:
             raise RuntimeError("Could not install Conda.")

runhouse/resources/hardware/cluster.py

@@ -444,6 +444,11 @@ def _command_runner(self, node: Optional[str] = None) -> "CommandRunner":
             SkySSHRunner,
         )
 
+        if node == "all":
+            raise ValueError(
+                "CommandRunner can only be instantiated for individual nodes"
+            )
+
         node = node or self.address
 
         if (
@@ -518,7 +523,8 @@ def _sync_default_env_to_cluster(self):
             logger.info(f"Using log level {log_level} on cluster's default env")
 
         logger.info(f"Syncing default env {self._default_env.name} to cluster")
-        self._default_env.install(cluster=self)
+        for node in self.ips:
+            self._default_env.install(cluster=self, node=node)
 
     def _sync_runhouse_to_cluster(
         self,

runhouse/resources/hardware/utils.py

@@ -125,18 +125,18 @@ def _unnamed_default_env_name(cluster_name):
     return f"{cluster_name}_default_env"
 
 
-def detect_cuda_version_or_cpu(cluster: "Cluster" = None):
+def detect_cuda_version_or_cpu(cluster: "Cluster" = None, node: Optional[str] = None):
     """Return the CUDA version on the cluster. If we are on a CPU-only cluster return 'cpu'.
 
     Note: A cpu-only machine may have the CUDA toolkit installed, which means nvcc will still return
     a valid version. Also check if the NVIDIA driver is installed to confirm we are on a GPU."""
 
-    status_codes = run_setup_command("nvcc --version", cluster=cluster)
+    status_codes = run_setup_command("nvcc --version", cluster=cluster, node=node)
     if not status_codes[0] == 0:
         return "cpu"
     cuda_version = status_codes[1].split("release ")[1].split(",")[0]
 
-    if run_setup_command("nvidia-smi", cluster=cluster)[0] == 0:
+    if run_setup_command("nvidia-smi", cluster=cluster, node=node)[0] == 0:
         return cuda_version
     return "cpu"
 

runhouse/resources/packages/package.py

@@ -223,7 +223,12 @@ def _reqs_install_cmd(
         install_cmd = self._prepend_env_command(install_cmd, env=env)
         return install_cmd
 
-    def _install(self, env: Union[str, "Env"] = None, cluster: "Cluster" = None):
+    def _install(
+        self,
+        env: Union[str, "Env"] = None,
+        cluster: "Cluster" = None,
+        node: Optional[str] = None,
+    ):
         """Install package.
 
         Args:
@@ -240,6 +245,7 @@ def _install(self, env: Union[str, "Env"] = None, cluster: "Cluster" = None):
                     dest=str(self.install_target.path_to_sync_to_on_cluster),
                     up=True,
                     contents=True,
+                    node=node,
                 )
 
                 self.install_target.local_path = (
@@ -265,6 +271,7 @@ def _install(self, env: Union[str, "Env"] = None, cluster: "Cluster" = None):
                 retcode = run_setup_command(
                     f"python -c \"import importlib.util; exit(0) if importlib.util.find_spec('{self.install_target}') else exit(1)\"",
                     cluster=cluster,
+                    node=node,
                 )[0]
                 if retcode != 0:
                     self.install_target = (
@@ -273,7 +280,7 @@ def _install(self, env: Union[str, "Env"] = None, cluster: "Cluster" = None):
 
             install_cmd = self._pip_install_cmd(env=env, cluster=cluster)
             logger.info(f"Running via install_method pip: {install_cmd}")
-            retcode = run_setup_command(install_cmd, cluster=cluster)[0]
+            retcode = run_setup_command(install_cmd, cluster=cluster, node=node)[0]
             if retcode != 0:
                 raise RuntimeError(
                     f"Pip install {install_cmd} failed, check that the package exists and is available for your platform."
@@ -282,7 +289,7 @@ def _install(self, env: Union[str, "Env"] = None, cluster: "Cluster" = None):
         elif self.install_method == "conda":
             install_cmd = self._conda_install_cmd(env=env, cluster=cluster)
             logger.info(f"Running via install_method conda: {install_cmd}")
-            retcode = run_setup_command(install_cmd, cluster=cluster)[0]
+            retcode = run_setup_command(install_cmd, cluster=cluster, node=node)[0]
             if retcode != 0:
                 raise RuntimeError(
                     f"Conda install {install_cmd} failed, check that the package exists and is "
@@ -293,7 +300,7 @@ def _install(self, env: Union[str, "Env"] = None, cluster: "Cluster" = None):
             install_cmd = self._reqs_install_cmd(env=env, cluster=cluster)
             if install_cmd:
                 logger.info(f"Running via install_method reqs: {install_cmd}")
-                retcode = run_setup_command(install_cmd, cluster=cluster)[0]
+                retcode = run_setup_command(install_cmd, cluster=cluster, node=node)[0]
                 if retcode != 0:
                     raise RuntimeError(
                         f"Reqs install {install_cmd} failed, check that the package exists and is available for your platform."
@@ -315,6 +322,7 @@ def _install(self, env: Union[str, "Env"] = None, cluster: "Cluster" = None):
                 ) if not cluster else run_setup_command(
                     f"export PATH=$PATH;{self.install_target.full_local_path_str()}",
                     cluster=cluster,
+                    node=node,
                 )
             elif not cluster:
                 if Path(self.install_target).resolve().expanduser().exists():