Add --reload flag to the serve run subcommand (ray-project#38389)

I have implemented a custom watch using the `watchfiles` package, formerly named `watchgod`, for my own local development of Ray Serve applications. It would be nice if this functionality existed in the `serve` Command Line Interface (CLI). I have searched the GitHub issues and Ray forums. This feature has been requested off and on several times.

This PR adds the `--reload` flag to the `serve run` subcommand. It will listen for file changes in the specified working directory and redeploy the application. This is similar uvicorn and other web frameworks. It is useful for application development.

The `watchfiles` package is added to the `requirements.txt` file and the `setup.py` file as indicated.

Co-authored-by: Edward Oakes <ed.nmi.oakes@gmail.com>
Co-authored-by: Christopher Field <chris.field@theiascientific.com>
Signed-off-by: e428265 <arvind.chandramouli@lmco.com>

Author: 3 people

.gitignore

@@ -183,6 +183,7 @@ java/**/pom.xml
 
 # python virtual env
 venv
+.venv
 
 # pyenv version file
 .python-version

doc/requirements-doc.txt

@@ -71,4 +71,7 @@ myst-nb==0.13.1
 jupytext==1.13.6
 
 # Pin urllib to avoid downstream ssl incompatibility issues
-urllib3 < 1.27
+urllib3 < 1.27
+
+# For `serve run --reload` CLI.
+watchfiles==0.19.0

python/ray/_private/utils.py

@@ -1285,9 +1285,11 @@ def new_func(*args, **kwargs):
     return deprecated_wrapper
 
 
-def import_attr(full_path: str):
+def import_attr(full_path: str, *, reload_module: bool = False):
     """Given a full import path to a module attr, return the imported attr.
 
+    If `reload_module` is set, the module will be reloaded using `importlib.reload`.
+
     For example, the following are equivalent:
         MyClass = import_attr("module.submodule:MyClass")
         MyClass = import_attr("module.submodule.MyClass")
@@ -1312,6 +1314,8 @@ def import_attr(full_path: str):
         attr_name = full_path[last_period_idx + 1 :]
 
     module = importlib.import_module(module_name)
+    if reload_module:
+        importlib.reload(module)
     return getattr(module, attr_name)
 
 

python/ray/serve/scripts.py

@@ -10,6 +10,7 @@
 import yaml
 import traceback
 import re
+import watchfiles
 from pydantic import ValidationError
 
 import ray
@@ -323,6 +324,17 @@ def deploy(config_file_name: str, address: str):
         "as the ingress deployment."
     ),
 )
+@click.option(
+    "--reload",
+    "-r",
+    is_flag=True,
+    help=(
+        "Listens for changes to files in the working directory, --working-dir "
+        "or the working_dir in the --runtime-env, and automatically redeploys "
+        "the application. This will block until Ctrl-C'd, then clean up the "
+        "app."
+    ),
+)
 def run(
     config_or_import_path: str,
     arguments: Tuple[str],
@@ -335,6 +347,7 @@ def run(
     port: int,
     blocking: bool,
     gradio: bool,
+    reload: bool,
 ):
     sys.path.insert(0, app_dir)
     args_dict = convert_args_to_dict(arguments)
@@ -453,11 +466,36 @@ def run(
 
             visualizer = GraphVisualizer()
             visualizer.visualize_with_gradio(handle)
-        else:
-            if blocking:
-                while True:
-                    # Block, letting Ray print logs to the terminal.
-                    time.sleep(10)
+        elif reload:
+            if not blocking:
+                raise click.ClickException(
+                    "The --non-blocking option conflicts with the --reload option."
+                )
+            if working_dir:
+                watch_dir = working_dir
+            else:
+                watch_dir = app_dir
+
+            for changes in watchfiles.watch(
+                watch_dir,
+                rust_timeout=10000,
+                yield_on_timeout=True,
+            ):
+                if changes:
+                    cli_logger.info(
+                        f"Detected file change in path {watch_dir}. Redeploying app."
+                    )
+                    # The module needs to be reloaded with `importlib` in order to pick
+                    # up any changes.
+                    app = _private_api.call_app_builder_with_args_if_necessary(
+                        import_attr(import_path, reload_module=True), args_dict
+                    )
+                    serve.run(app, host=host, port=port)
+
+        if blocking:
+            while True:
+                # Block, letting Ray print logs to the terminal.
+                time.sleep(10)
 
     except KeyboardInterrupt:
         cli_logger.info("Got KeyboardInterrupt, shutting down...")

python/ray/serve/tests/test_cli_2.py

@@ -825,5 +825,58 @@ def test_deployment_contains_utils(ray_start_stop):
     )
 
 
+@pytest.mark.skipif(sys.platform == "win32", reason="File path incorrect on Windows.")
+def test_run_reload_basic(ray_start_stop, tmp_path):
+    """Test `serve run` with reload."""
+
+    code_template = """
+from ray import serve
+
+@serve.deployment
+class MessageDeployment:
+    def __init__(self, msg):
+        self.msg = msg
+
+    def __call__(self):
+        return self.msg
+
+
+msg_app = MessageDeployment.bind("Hello {message}!")
+    """
+
+    def write_file(message: str):
+        with open(os.path.join(tmp_path, "reload_serve.py"), "w") as f:
+            code = code_template.format(message=message)
+            print(f"Writing updated code:\n{code}")
+            f.write(code)
+            f.flush()
+
+    write_file("World")
+
+    p = subprocess.Popen(
+        [
+            "serve",
+            "run",
+            "--app-dir",
+            tmp_path,
+            "--reload",
+            "reload_serve:msg_app",
+        ]
+    )
+    wait_for_condition(lambda: ping_endpoint("") == "Hello World!", timeout=10)
+
+    # Sleep to ensure the `serve run` command is in the file watching loop when we
+    # write the change, else it won't be picked up.
+    time.sleep(5)
+
+    # Write the file: an update should be auto-triggered.
+    write_file("Updated")
+    wait_for_condition(lambda: ping_endpoint("") == "Hello Updated!", timeout=10)
+
+    p.send_signal(signal.SIGINT)
+    p.wait()
+    assert ping_endpoint("") == CONNECTION_ERROR_MSG
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", "-s", __file__]))

python/requirements.txt

@@ -17,6 +17,7 @@ pyyaml
 aiosignal
 frozenlist
 requests
+watchfiles
 
 # Python version-specific requirements
 dataclasses; python_version < '3.7'
@@ -61,3 +62,4 @@ fsspec
 pandas>=1.3
 pydantic<2
 py-spy>=0.2.0
+watchfiles

python/requirements/test-requirements.txt

@@ -118,4 +118,7 @@ sympy==1.10.1; python_version <= '3.7'
 
 # For test_basic.py::test_omp_threads_set
 threadpoolctl==3.1.0
-numexpr==2.8.4
+numexpr==2.8.4
+
+# For `serve run --reload` CLI.
+watchfiles==0.19.0

python/setup.py

@@ -271,7 +271,14 @@ def get_packages(self):
             if sys.platform == "darwin"
             else "grpcio",
         ],
-        "serve": ["uvicorn", "requests", "starlette", "fastapi", "aiorwlock"],
+        "serve": [
+            "uvicorn",
+            "requests",
+            "starlette",
+            "fastapi",
+            "aiorwlock",
+            "watchfiles",
+        ],
         "tune": ["pandas", "tensorboardX>=1.9", "requests", pyarrow_dep, "fsspec"],
         "observability": [
             "opentelemetry-api",