Add image api docs (#1578)

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
 

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__

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.

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.")
 

runhouse/resources/images/__init__.py

@@ -1 +1 @@
-from .image import Image, ImageSetupStepType
+from .image import Image, ImageSetupStep, ImageSetupStepType

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,25 +19,56 @@ 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
 
 
 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
 
         self.setup_steps = []
         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 "<registry>/<image>:<tag>"
+            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,