diff --git a/docs/api/python.rst b/docs/api/python.rst index 2df693070..0a7e4ca29 100644 --- a/docs/api/python.rst +++ b/docs/api/python.rst @@ -33,6 +33,11 @@ and rich debugging and accessibility interfaces built-in. python/cluster +.. toctree:: + :maxdepth: 1 + + python/image + .. toctree:: :maxdepth: 1 diff --git a/docs/api/python/image.rst b/docs/api/python/image.rst new file mode 100644 index 000000000..4b0922105 --- /dev/null +++ b/docs/api/python/image.rst @@ -0,0 +1,32 @@ +Image +===== + +Image Class +~~~~~~~~~~~ + +.. autoclass:: runhouse.Image + :members: + :exclude-members: + + .. automethod:: __init__ + +ImageSteupStepType +~~~~~~~~~~~~~~~~~~ + +.. autoclass:: runhouse.resources.images.ImageSetupStepType + + .. autoattribute:: PACKAGES + .. autoattribute:: CMD_RUN + .. autoattribute:: SETUP_CONDA_ENV + .. autoattribute:: RSYNC + .. autoattribute:: SYNC_SECRETS + .. autoattribute:: SET_ENV_VARS + +ImageSetupStep +~~~~~~~~~~~~~~ + +.. autoclass:: runhouse.resources.images.ImageSetupStep + :members: + :exclude-members: + + .. automethod:: __init__ diff --git a/runhouse/__init__.py b/runhouse/__init__.py index ff0143b35..7d9d5d23b 100644 --- a/runhouse/__init__.py +++ b/runhouse/__init__.py @@ -11,7 +11,7 @@ ondemand_cluster, OnDemandCluster, ) -from runhouse.resources.images.image import Image +from runhouse.resources.images import Image # WARNING: Any built-in module that is imported here must be capitalized followed by all lowercase, or we will # will not find the module class when attempting to reconstruct it from a config. diff --git a/runhouse/resources/hardware/cluster.py b/runhouse/resources/hardware/cluster.py index 944040e53..983fa28aa 100644 --- a/runhouse/resources/hardware/cluster.py +++ b/runhouse/resources/hardware/cluster.py @@ -1603,11 +1603,22 @@ def _copy_certs_to_cluster(self): def run_bash( self, commands: Union[str, List[str]], - node: Optional[str] = None, + node: Union[int, str, None] = None, process: Optional[str] = None, stream_logs: bool = True, require_outputs: bool = True, ): + """Run bash commands on the cluster through the Runhouse server. + + Args: + commands (str or List[str]): Commands to run on the cluster. + node (int, str or None): Node to run the command on. Node can an int referring to the node index, + string referring to the ips, or "all" to run on all nodes. If not specified, run the command + on the head node. (Default: ``None``) + process (str or None): Process to run the command on. (Default: ``None``) + stream_logs (bool): Whether to stream logs. (Default: ``True``) + require_outputs (bool): Whether to return outputs in addition to status code. (Default: ``True``) + """ if isinstance(commands, str): commands = [commands] @@ -1722,12 +1733,23 @@ async def print_logs(): def run_bash_over_ssh( self, commands: Union[str, List[str]], - node: Optional[str] = None, + node: Union[int, str, None] = None, stream_logs: bool = True, require_outputs: bool = True, _ssh_mode: str = "interactive", # Note, this only applies for non-password SSH conda_env_name: Optional[str] = None, ): + """Run bash commands on the cluster over SSH. + + Args: + commands (str or List[str]): Commands to run on the cluster. + node (int, str or None): Node to run the command on. Node can an int referring to the node index, + string referring to the ips, or "all" to run on all nodes. If not specified, run the command + on the head node. (Default: ``None``) + stream_logs (bool): Whether to stream logs. (Default: ``True``) + require_outputs (bool): Whether to return outputs in addition to status code. (Default: ``True``) + conda_env_name (str or None): Name of conda env to run the command in, if applicable. (Defaut: ``None``) + """ if self.on_this_cluster(): raise ValueError("Run bash over SSH is not supported on the local cluster.") diff --git a/runhouse/resources/images/__init__.py b/runhouse/resources/images/__init__.py index 7e4d9d6c5..db2c3fb42 100644 --- a/runhouse/resources/images/__init__.py +++ b/runhouse/resources/images/__init__.py @@ -1 +1 @@ -from .image import Image, ImageSetupStepType +from .image import Image, ImageSetupStep, ImageSetupStepType diff --git a/runhouse/resources/images/image.py b/runhouse/resources/images/image.py index 545c6ccc1..1e93cafe0 100644 --- a/runhouse/resources/images/image.py +++ b/runhouse/resources/images/image.py @@ -3,6 +3,8 @@ # Internal class to represent the image construction process class ImageSetupStepType(Enum): + """Enum for valid Image setup step types""" + PACKAGES = "packages" CMD_RUN = "cmd_run" SETUP_CONDA_ENV = "setup_conda_env" @@ -17,6 +19,15 @@ def __init__( step_type: ImageSetupStepType, **kwargs: Dict[str, Any], ): + """ + A component of the Runhouse Image, consisting of the step type (e.g. packages, set_env_vars), + along with arguments to provide to the function corresponding to the step type. + + Args: + step_type (ImageSetupStepType): Type of setup step used to provide the Image. + kwargs (Dict[str, Any]): Please refer to the corresponding functions in `Image` to determine + the correct keyword arguments to provide. + """ self.step_type = step_type self.kwargs = kwargs @@ -24,10 +35,24 @@ def __init__( class Image: def __init__(self, name: str, image_id: str = None): """ + Runhouse Image object, specifying cluster setup properties and steps. + Args: name (str): Name to assign the Runhouse image. image_id (str): Machine image to use, if any. (Default: ``None``) + + Example: + >>> my_image = ( + >>> rh.Image(name="base_image") + >>> .setup_conda_env( + >>> conda_env_name="base_env", + >>> conda_yaml={"dependencies": ["python=3.11"], "name": "base_env"}, + >>> ) + >>> .install_packages(["numpy", "pandas"]) + >>> .set_env_vars({"OMP_NUM_THREADS": 1}) + >>> ) """ + self.name = name self.image_id = image_id @@ -35,7 +60,15 @@ def __init__(self, name: str, image_id: str = None): self.conda_env_name = None self.docker_secret = None - def from_docker(self, image_id, docker_secret: Union["Secret", str] = None): + def from_docker(self, image_id: str, docker_secret: Union["Secret", str] = None): + """Set up and use an existing Docker image. + + Args: + image_id (str): Docker image in the following format "/:" + docker_secret (Secret or str, optional): Runhouse secret corresponding to information + necessary to pull the image from a private registry, such as Docker Hub or cloud provider. + See ``DockerRegistrySecret`` for more information. + """ if self.image_id: raise ValueError( "Setting both a machine image and docker image is not yet supported." @@ -89,7 +122,19 @@ def from_config(cls, config: Dict[str, Any]): # Steps to build the image ######################################################## - def install_packages(self, reqs: List[str], conda_env_name: Optional[str] = None): + def install_packages( + self, reqs: List[Union["Package", str]], conda_env_name: Optional[str] = None + ): + """Install the given packages. + + Args: + reqs (List[Package or str]): List of packages to install on cluster and env. + node (str, optional): Cluster node to install the package on. If specified, will use ssh to install the + package. (Default: ``None``) + conda_env_name (str, optional): Name of conda env to install the package in, if relevant. If left empty, + defaults to base environment. (Default: ``None``) + """ + self.setup_steps.append( ImageSetupStep( step_type=ImageSetupStepType.PACKAGES, @@ -100,6 +145,12 @@ def install_packages(self, reqs: List[str], conda_env_name: Optional[str] = None return self def run_bash(self, command: str, conda_env_name: Optional[str] = None): + """Run bash commands. + + Args: + command (str): Commands to run on the cluster. + conda_env_name (str or None): Name of conda env to run the command in, if applicable. (Defaut: ``None``) + """ self.setup_steps.append( ImageSetupStep( step_type=ImageSetupStepType.CMD_RUN, @@ -110,6 +161,12 @@ def run_bash(self, command: str, conda_env_name: Optional[str] = None): return self def setup_conda_env(self, conda_env_name: str, conda_yaml: Union[str, Dict]): + """Setup Conda env + + Args: + conda_env_name (str): Name of conda env to create. + conda_yaml (str or Dict): Path or Dict referring to the conda yaml config to use to create the conda env. + """ self.setup_steps.append( ImageSetupStep( step_type=ImageSetupStepType.SETUP_CONDA_ENV, @@ -123,6 +180,19 @@ def setup_conda_env(self, conda_env_name: str, conda_yaml: Union[str, Dict]): def rsync( self, source: str, dest: str, contents: bool = False, filter_options: str = None ): + """Sync the contents of the source directory into the destination. + + Args: + source (str): The source path. + dest (str): The target path. + contents (Optional[bool], optional): Whether the contents of the source directory or the directory + itself should be copied to destination. + If ``True`` the contents of the source directory are copied to the destination, and the source + directory itself is not created at the destination. + If ``False`` the source directory along with its contents are copied ot the destination, creating + an additional directory layer at the destination. (Default: ``False``). + filter_options (Optional[str], optional): The filter options for rsync. + """ self.setup_steps.append( ImageSetupStep( step_type=ImageSetupStepType.RSYNC, @@ -134,6 +204,11 @@ def rsync( ) def sync_secrets(self, providers: List[Union[str, "Secret"]]): + """Send secrets for the given providers. + + Args: + providers(List[str or Secret]): List of providers to send secrets for. + """ self.setup_steps.append( ImageSetupStep( step_type=ImageSetupStepType.SYNC_SECRETS, @@ -143,6 +218,12 @@ def sync_secrets(self, providers: List[Union[str, "Secret"]]): return self def set_env_vars(self, env_vars: Union[str, Dict]): + """Set environment variables. + + Args: + env_vars (str or Dict): Dict of environment variables and values to set, or string pointing + to local .env file consisting of env vars to set. + """ self.setup_steps.append( ImageSetupStep( step_type=ImageSetupStepType.SET_ENV_VARS,