Selective watch, fixes BT-12924

Author: marius-baseten

truss-chains/truss_chains/deployment/deployment_client.py

@@ -470,6 +470,7 @@ class _Watcher:
     _console: "rich_console.Console"
     _error_console: "rich_console.Console"
     _show_stack_trace: bool
+    _included_chainlets: set[str]
 
     def __init__(
         self,
@@ -480,6 +481,7 @@ def __init__(
         console: "rich_console.Console",
         error_console: "rich_console.Console",
         show_stack_trace: bool,
+        included_chainlets: Optional[list[str]],
     ) -> None:
         self._source = source
         self._entrypoint = entrypoint
@@ -498,6 +500,16 @@ def __init__(
                 for desc in _get_ordered_dependencies([entrypoint_cls])
             )
 
+        if included_chainlets:
+            if not_matched := (set(included_chainlets) - chainlet_names):
+                raise definitions.ChainsDeploymentError(
+                    "Requested to watch specific chainlets, but did not find "
+                    f"{not_matched} among available chainlets {chainlet_names}."
+                )
+            self._included_chainlets = set(included_chainlets)
+        else:
+            self._included_chainlets = chainlet_names
+
         chain_id = b10_core.get_chain_id_by_name(
             self._remote_provider.api, self._deployed_chain_name
         )
@@ -594,6 +606,13 @@ def _patch(self, executor: concurrent.futures.Executor) -> None:
                     )
                     future_to_display_name = {}
                     for chainlet_descr in chainlet_descriptors:
+                        if chainlet_descr.display_name not in self._included_chainlets:
+                            self._console.print(
+                                f"⏩ Skipping patching `{chainlet_descr.display_name}`.",
+                                style="grey50",
+                            )
+                            continue
+
                         future = executor.submit(
                             self._code_gen_and_patch_thread,
                             chainlet_descr,
@@ -658,17 +677,17 @@ def _check_patch_results(
 
             if patch_result.status == b10_remote.PatchStatus.SUCCESS:
                 self._console.print(
-                    f"Patched Chainlet `{display_name}`.{logs_output}", style="green"
+                    f"✅ Patched Chainlet `{display_name}`.{logs_output}", style="green"
                 )
             elif patch_result.status == b10_remote.PatchStatus.SKIPPED:
                 self._console.print(
-                    f"Nothing to do for Chainlet `{display_name}`.{logs_output}",
+                    f"💤 Nothing to do for Chainlet `{display_name}`.{logs_output}",
                     style="grey50",
                 )
             else:
                 has_errors = True
                 self._error_console.print(
-                    f"Failed to patch Chainlet `{display_name}`. "
+                    f"❌ Failed to patch Chainlet `{display_name}`. "
                     f"{patch_result.message}{logs_output}"
                 )
 
@@ -702,6 +721,7 @@ def watch(
     console: "rich_console.Console",
     error_console: "rich_console.Console",
     show_stack_trace: bool,
+    included_chainlets: Optional[list[str]],
 ) -> None:
     console.print(
         (
@@ -711,6 +731,13 @@ def watch(
         style="blue",
     )
     patcher = _Watcher(
-        source, entrypoint, name, remote, console, error_console, show_stack_trace
+        source,
+        entrypoint,
+        name,
+        remote,
+        console,
+        error_console,
+        show_stack_trace,
+        included_chainlets,
     )
     patcher.watch()

truss/cli/cli.py

@@ -86,7 +86,6 @@
     ]
 }
 
-
 console = Console()
 
 error_console = Console(stderr=True, style="bold red")
@@ -555,11 +554,15 @@ def _create_chains_table(service) -> Tuple[rich.table.Table, List[str]]:
     help="Name of the remote in .trussrc to push to.",
 )
 @click.option(
-    "--user_env",
-    required=False,
+    "--experimental-watch-chainlet-names",
     type=str,
-    help="[DEPRECATED], use ``environment`` instead.",
-    hidden=True,
+    required=False,
+    help=(
+        "Runs `watch`, but only applies patches to specified chainlets. The option is "
+        "a comma-separated list of chainlet (display) names. This option can give "
+        "faster dev loops, but also lead to inconsistent deployments. Use with caution "
+        "and refer to docs."
+    ),
 )
 @log_level_option
 @error_handling
@@ -572,9 +575,9 @@ def push_chain(
     wait: bool,
     watch: bool,
     dryrun: bool,
-    user_env: Optional[str],
     remote: Optional[str],
     environment: Optional[str],
+    experimental_watch_chainlet_names: Optional[str],
 ) -> None:
     """
     Deploys a chain remotely.
@@ -589,6 +592,9 @@ def push_chain(
     from truss_chains import framework
     from truss_chains.deployment import deployment_client
 
+    if experimental_watch_chainlet_names:
+        watch = True
+
     if watch:
         if publish or promote:
             raise ValueError(
@@ -600,13 +606,10 @@ def push_chain(
             )
             wait = True
 
-    if user_env:
-        raise ValueError("`user_env` is deprecated, use `environment` instead.")
-
     if promote and environment:
         promote_warning = (
             "`promote` flag and `environment` flag were both specified. "
-            "Ignoring the value of `promote`"
+            "Ignoring the value of `promote`."
         )
         console.print(promote_warning, style="yellow")
 
@@ -674,7 +677,14 @@ def push_chain(
                 )
             console.print(deploy_success_text, style="bold green")
             console.print(f"You can run the chain with:\n{curl_snippet}")
+
             if watch:  # Note that this command will print a startup message.
+                if experimental_watch_chainlet_names:
+                    included_chainlets = [
+                        x.strip() for x in experimental_watch_chainlet_names.split(",")
+                    ]
+                else:
+                    included_chainlets = None
                 deployment_client.watch(
                     source,
                     entrypoint,
@@ -683,6 +693,7 @@ def push_chain(
                     console,
                     error_console,
                     show_stack_trace=not is_humanfriendly_log_level,
+                    included_chainlets=included_chainlets,
                 )
         else:
             console.print(f"Deployment failed ({num_failed} failures).", style="red")
@@ -710,20 +721,24 @@ def push_chain(
     help="Name of the remote in .trussrc to push to.",
 )
 @click.option(
-    "--user_env",
-    required=False,
+    "--experimental-chainlet-names",
     type=str,
-    help="[DEPRECATED], use `environment` instead.",
-    hidden=True,
+    required=False,
+    help=(
+        "Runs `watch`, but only applies patches to specified chainlets. The option is "
+        "a comma-separated list of chainlet (display) names. This option can give "
+        "faster dev loops, but also lead to inconsistent deployments. Use with caution "
+        "and refer to docs."
+    ),
 )
 @log_level_option
 @error_handling
 def watch_chains(
     source: Path,
     entrypoint: Optional[str],
     name: Optional[str],
-    user_env: Optional[str],
     remote: Optional[str],
+    experimental_chainlet_names: Optional[str],
 ) -> None:
     """
     Watches the chains source code and applies live patches to a development deployment.
@@ -738,12 +753,14 @@ def watch_chains(
     # These imports are delayed, to handle pydantic v1 envs gracefully.
     from truss_chains.deployment import deployment_client
 
-    if user_env:
-        raise ValueError("`user_env` is deprecated, use `environment` instead.")
-
     if not remote:
         remote = inquire_remote_name(RemoteFactory.get_available_config_names())
 
+    if experimental_chainlet_names:
+        included_chainlets = [x.strip() for x in experimental_chainlet_names.split(",")]
+    else:
+        included_chainlets = None
+
     deployment_client.watch(
         source,
         entrypoint,
@@ -752,6 +769,7 @@ def watch_chains(
         console,
         error_console,
         show_stack_trace=not is_humanfriendly_log_level,
+        included_chainlets=included_chainlets,
     )