feat: Add debug mode to one-shot commands

docs/content/misc/debugging.md

@@ -96,6 +96,25 @@ green "play" button at the top left corner of the "Run and Debug" tab.
 
 For more information on the VSCode debugger, see [this guide](https://code.visualstudio.com/docs/debugtest/debugging).
 
+### Debugging one-shot commands
+
+`tesseract run` also supports a `--debug` flag, which lets you remotely debug a one-shot command
+rather than a long-running server. It works with any command you can run this way (`apply`,
+`jacobian`, `check-gradients`, `health`, ...):
+
+```bash
+$ tesseract run --debug helloworld apply '{"inputs": {"a": 1.0, "b": 2.0}}'
+```
+
+Because `tesseract run` executes a single command and then exits, the runtime starts a debugpy
+server inside the container and **blocks until a debugger attaches** before running anything. The
+debugger attaches before your `tesseract_api.py` is even imported, so breakpoints in module-level
+code (e.g. model loading at import time) are hit as well as those inside the endpoint itself.
+
+The host port to connect to is printed in the CLI when the command starts. Use the same VSCode
+launch config as above (filling in the printed port number) to attach; execution resumes
+automatically once the debugger is connected.
+
 (profiling)=
 
 ## Profiling

tesseract_core/runtime/cli.py

@@ -213,6 +213,40 @@ def _schema_to_docstring(schema: Any, current_indent: int = 0) -> str:
     return "\n".join(docstring)
 
 
+def _maybe_start_debugger(wait_for_client: bool, port: int = 5678) -> None:
+    """Start a debugpy server for remote debugging when debug mode is enabled.
+
+    The long-running ``serve`` command launches a non-blocking server that a
+    debugger can attach to at any time. One-shot commands instead attach early in
+    ``main`` (before the Tesseract API is imported, so module-level code can be
+    debugged too) and block until a client connects, since they would otherwise
+    finish before there is a chance to attach.
+
+    Args:
+        wait_for_client: If True, block until a debugger attaches.
+        port: Port to listen on inside the container.
+    """
+    if not get_config().debug:
+        return
+
+    # Python 3.11+ freezes stdlib bootstrap modules, which makes debugpy print a
+    # noisy "frozen modules" warning (it could only ever miss breakpoints inside
+    # those frozen modules, never in user code). Skip the validation check.
+    os.environ.setdefault("PYDEVD_DISABLE_FILE_VALIDATION", "1")
+
+    import debugpy
+
+    debugpy.listen(("0.0.0.0", port))
+    if wait_for_client:
+        print(
+            "Debug mode enabled, waiting for debugger to attach...",
+            file=sys.stderr,
+            flush=True,
+        )
+        debugpy.wait_for_client()
+        print("Debugger attached, resuming execution.", file=sys.stderr, flush=True)
+
+
 @app.command("check")
 def check() -> None:
     """Check whether the Tesseract API is valid."""
@@ -376,6 +410,8 @@ def serve(
     num_workers: Annotated[int, typer.Option(help="Number of worker processes")] = 1,
 ) -> None:
     """Start running this Tesseract's web server."""
+    # The server is long-running, so a debugger can attach at any time.
+    _maybe_start_debugger(wait_for_client=False)
     serve_(host=host, port=port, num_workers=num_workers)
 
 
@@ -468,6 +504,7 @@ def _callback_wrapper(**kwargs: Any):
         def command_func(payload: str):
             parsed_payload = _parse_payload(payload)
             return _callback_wrapper(payload=parsed_payload)
+
     else:
 
         def command_func():
@@ -544,6 +581,15 @@ def main() -> None:
 
         _configure_required_file_load()
 
+        # Attach the debugger before the Tesseract API is imported below (during
+        # command registration) so module-level code can be debugged too. The
+        # command isn't parsed yet, so we inspect argv directly (like
+        # `_configure_required_file_load` above): `serve` launches its own
+        # non-blocking debugger and must not block, and help should not block.
+        skip_debug_wait_args = {"serve", "-h", "--help"}
+        if get_config().debug and not skip_debug_wait_args.intersection(sys.argv):
+            _maybe_start_debugger(wait_for_client=True)
+
         _add_user_commands_to_cli(app, out_stream=orig_stdout)
         app(auto_envvar_prefix="TESSERACT_RUNTIME")
 

tesseract_core/runtime/serve.py

@@ -127,12 +127,6 @@ async def wrapper(*args: Any, accept: str, run_id: str | None, **kwargs: Any):
 
 def serve(host: str, port: int, num_workers: int) -> None:
     """Start the REST API."""
-    config = get_config()
-    if config.debug:
-        import debugpy
-
-        debugpy.listen(("0.0.0.0", 5678))
-
     uvicorn.run(
         "tesseract_core.runtime.app_http:app",
         host=host,

tesseract_core/sdk/cli.py

@@ -1115,6 +1115,17 @@ def run_container(
             help="Enable tracing for detailed debug output.",
         ),
     ] = False,
+    debug: Annotated[
+        bool,
+        typer.Option(
+            "--debug",
+            help=(
+                "Enable debug mode. This starts a debugpy server in the Tesseract and "
+                "blocks until a debugger attaches to the forwarded port. "
+                "WARNING: This may expose sensitive information, use with caution (and never in production)."
+            ),
+        ),
+    ] = False,
     invoke_help: Annotated[
         bool,
         typer.Option(
@@ -1229,6 +1240,9 @@ def run_container(
             user=user,
             memory=memory,
             docker_args=shlex.split(docker_args) if docker_args else None,
+            # `--debug` is meaningless when only forwarding `--help`, and would
+            # otherwise block waiting for a debugger on a help invocation.
+            debug=debug and not invoke_help,
             stream_logs=logger.info,  # Stream logs via logger
         )
 

tesseract_core/sdk/engine.py

@@ -961,6 +961,7 @@ def run_tesseract(
     output_format: Literal["json", "json+base64", "json+binref"] | None = None,
     output_file: str | None = None,
     docker_args: list[str] | None = None,
+    debug: bool = False,
     stream_logs: bool | Callable[[str], None] = False,
 ) -> tuple[str, str]:
     """Start a Tesseract and execute a given command.
@@ -985,6 +986,8 @@ def run_tesseract(
         output_file: If specified, the output will be written to this file within output_path
             instead of stdout.
         docker_args: Additional arguments to pass to the container runtime (e.g., Docker).
+        debug: Enable debug mode. This starts a debugpy server in the Tesseract and
+            blocks execution until a debugger attaches to the forwarded port.
         stream_logs: If set, stream logs in real-time. Can be True (streams to stderr)
             or a callable that accepts a string (e.g., logger.info).
 
@@ -1058,6 +1061,23 @@ def run_tesseract(
     if network is not None:
         _ensure_network_exists(network)
 
+    if debug:
+        environment["TESSERACT_DEBUG"] = "1"
+        # `network="host"` binds the container's debugpy port directly on the host,
+        # so no explicit port mapping is needed (and would actually be rejected).
+        if network == "host":
+            debugpy_port = "5678"
+        else:
+            debugpy_port = str(get_free_port())
+            if ports is None:
+                ports = {}
+            ports[f"127.0.0.1:{debugpy_port}"] = "5678"
+        logger.info(
+            f"Debug mode enabled. Attach a debugger to localhost:{debugpy_port} "
+            "to start execution (see the 'Debug mode' section of the docs for a "
+            "sample VSCode launch config)."
+        )
+
     # Run the container, optionally streaming stderr to the terminal
     result = docker_client.containers.run(
         image=image,

tests/sdk_tests/test_engine.py

@@ -286,6 +286,22 @@ def test_run_tesseract(mocked_docker):
     assert res["device_requests"] is None
 
 
+def test_run_debug(mocked_docker):
+    """Test running a tesseract in debug mode forwards a debugpy port."""
+    res_out, _ = engine.run_tesseract(
+        "foobar",
+        "apply",
+        ['{"inputs": {"a": [1, 2, 3], "b": [4, 5, 6]}}'],
+        debug=True,
+    )
+
+    res = json.loads(res_out)
+    # TESSERACT_DEBUG is set so the runtime starts debugpy and waits for a client.
+    assert res["environment"]["TESSERACT_DEBUG"] == "1"
+    # The container's debugpy port (5678) is forwarded to a free port on the host.
+    assert "5678" in res["ports"].values()
+
+
 def test_run_gpu(mocked_docker):
     """Test running a tesseract with all available GPUs."""
     res_out, _ = engine.run_tesseract(