Decorator cli

runhouse/cli_utils.py

@@ -5,6 +5,7 @@
 
 from datetime import datetime, timezone
 from enum import Enum
+from functools import wraps
 from typing import Any, Dict, List
 
 import requests
@@ -589,3 +590,85 @@ def get_wrapped_server_start_cmd(flags: List[str], screen: bool, nohup: bool):
         )
 
     return wrapped_cmd
+
+
+####################################################################################################
+# Deploy utils
+####################################################################################################
+
+class PartialModule:
+    def __init__(self,
+                 cls,
+                 system,
+                    module_name,
+                module_process,
+                 distribute_args):
+        self.cls = cls
+        self.system = system
+        self.module_name = module_name
+        self.module_process = module_process
+        self.distribute_args = distribute_args
+
+    def __call__(self, *args, **kwargs):
+        # Raise an error that the user can't just use the @rh.distribute decorator, they need to apply
+        # the @rh.deploy decorator first.
+        raise ValueError("You need to use the @rh.deploy decorator before using @rh.distribute.")
+
+# @runhouse.deploy decorator that the user can use to wrap a function they want to deploy to a cluster,
+# and then deploy it with `runhouse deploy my_app.py` (we collect all the decorated functions imported in the file
+# to deploy them).
+def deploy(**kwargs, get_if_exists=False):
+
+    @wraps
+    def decorator(func_or_cls):
+        from runhouse.globals import _deploying, _to_deploy
+
+        if isinstance(func_or_cls, PartialModule):
+            distribute_args = func_or_cls.distribute_args
+            func_or_cls = func_or_cls.cls
+        else:
+            distribute_args = None
+
+        module_name = kwargs.pop("name", None)
+        module_process = kwargs.pop("process", None)
+        if "system" in kwargs:
+            system = kwargs.pop("system")
+            if kwargs:
+                raise ValueError("Please specify system or compute kwargs, but not both.")
+        else:
+            system: rh.Cluster = rh.cluster(**kwargs)
+
+        if isinstance(func_or_cls, type):
+            new_module = rh.module(func_or_cls, name=module_name, **kwargs)
+        else:
+            new_module = rh.function(func_or_cls, name=module_name, **kwargs)
+
+        if _deploying:
+            # Construct and register a partial module that will be deployed when the deploy command is run
+            if get_if_exists:
+                new_module.get_or_to(system=system, name=module_name, process=module_process)
+            else:
+                # Should we kill the existing process here too?
+                new_module.to(system=system, name=module_name, process=module_process)
+
+            if distribute_args:
+                new_module.distribute(*distribute_args[0], **distribute_args[1])
+            return new_module
+        else:
+            deployed_module = system.get(new_module.name, default=None, remote=True)
+            if not deployed_module:
+                raise ValueError(f"Module {new_module.name} not found on cluster {system.name}. Please deploy it first"
+                                 f" with `runhouse deploy <my_app.py>`.")
+            return deployed_module
+
+    return decorator
+
+def distribute(*args, **kwargs):
+    @wraps
+    def decorator(func_or_cls):
+        # This is a partial so the order of decorator chaining can be reversed for best aesthetics
+        # the deploy method will actually call .distribute on the function or class after it's been deployed
+        return PartialModule(func_or_cls, distribute_args=(args, kwargs))
+
+    return decorator
+

runhouse/globals.py

@@ -28,3 +28,5 @@ def clean_up_ssh_connections():
 # Note: this initalizes a dummy global object. The obj_store must
 # be properly initialized by a servlet via initialize.
 obj_store = ObjStore()
+
+_deploying = False

runhouse/main.py

@@ -133,6 +133,38 @@ def logout():
     raise typer.Exit()
 
 
+###############################
+# Deploy CLI commands
+@app.command("deploy")
+def deploy(
+        python_file: str = typer.Argument(..., help="Path to the Python file to deploy."),
+        get_if_exists: bool = typer.Option(
+            False,
+            "--get-if-exists",
+            help="Whether to get existing deployed functions and modules if they already exist.",
+        ),
+):
+    """Deploy a Python file to Runhouse. This will deploy all functions and modules decorated with
+    @rh.deploy in the file."""
+    # First, switch to deployment mode so deployment decorators know to behave accordingly
+    rh.globals._deploying = True
+
+    # Next, import the file to trigger all the deployment decorators
+    try:
+        with open(python_file, "r") as f:
+            code = f.read()
+            exec(code)
+    except FileNotFoundError:
+        console.print(f"File {python_file} not found.")
+        raise typer.Exit(1)
+
+    # Success
+    console.print(f"Successfully deployed functions and modules from {python_file}.")
+
+
+
+
+
 ###############################
 # Cluster CLI commands
 ###############################