Add local runtime support for in-situ unit test bisection (#1501)

Introduces a new `local` mode to the triage utility that allows
developers to perform a commit-level bisection for unit tests without
requiring container orchestration.

This is designed for the case where a user is already inside a failing
container and wants to find a culprit commit directly.

**Changes:**
1. Adds a new option, **--container-runtime=local**.
2. When this mode is active, all `git` operations and build/test
commands are executed directly on the local machine using `subprocess`.
3. The `run_and_log` utility was updated to accept a `cwd` parameter.
This was necessary to allow the `LocalContainer` to execute commands
like `git checkout` and `bazel` builds within the correct local
repository directories.
4. The `local` runtime skips the container search, so it requires the
user to provide the bisection range explicitly using:
    - `--passing-commits="jax:<hash>,xla:<hash>"`
    - `--failing-commits="jax:<hash>,xla:<hash>"`
5. The argument parser has been updated to enforce these new rules and
prevent mixing `local` mode with container-specific arguments (i.e,
`--container`, `--start-date`).
6. Unit tests to the new argument parsing logic

Now, this command can be run inside a JAX development container:
```
jax-toolbox-triage \
    --container-runtime=local \
    --passing-commits="jax:<some_good_hash>,xla:<some_good_hash>" \
    --failing-commits="jax:<some_bad_hash>,xla:<some_bad_hash>" \
    /opt/jax-toolbox/.github/container/test-jax.sh "some_target"
```

Author: camille-004

.github/triage/jax_toolbox_triage/args.py

@@ -187,11 +187,27 @@ def parse_args(args=None) -> argparse.Namespace:
     parser.add_argument(
         "--container-runtime",
         default="docker",
-        help="Container runtime used, this can be either docker or pyxis.",
+        help="Container runtime used, can be docker, pyxis, or local.",
         type=lambda s: s.lower(),
     )
     args = parser.parse_args(args=args)
-    assert args.container_runtime in {"docker", "pyxis"}, args.container_runtime
+    assert args.container_runtime in {"docker", "pyxis", "local"}, (
+        args.container_runtime
+    )
+
+    if args.container_runtime == "local":
+        assert args.passing_commits is not None and args.failing_commits is not None, (
+            "For local runtime, --passing-commits and --failing-commits must be provided."
+        )
+        assert (
+            args.container is None
+            and args.start_date is None
+            and args.end_date is None
+            and args.passing_container is None
+            and args.failing_container is None
+        ), "Container-level search options are not applicable for local runtime."
+        return args
+
     passing_commits_known = (args.passing_container is not None) or (
         args.passing_commits is not None
     )

.github/triage/jax_toolbox_triage/local.py

@@ -0,0 +1,34 @@
+import logging
+import subprocess
+import typing
+
+from .container import Container
+from .utils import run_and_log
+
+
+class LocalContainer(Container):
+    def __init__(self, *, logger: logging.Logger):
+        super().__init__(logger=logger)
+
+    def __enter__(self) -> Container:
+        self._logger.debug("Running local mode inside current container")
+        return self
+
+    def __exit__(self, *exc_info) -> None:
+        pass
+
+    def __repr__(self) -> str:
+        return "Local"
+
+    def exec(
+        self,
+        command: typing.List[str],
+        policy: typing.Literal["once", "once_per_container", "default"] = "default",
+        stderr: typing.Literal["interleaved", "separate"] = "interleaved",
+        workdir=None,
+    ) -> subprocess.CompletedProcess:
+        return run_and_log(command, logger=self._logger, stderr=stderr, cwd=workdir)
+
+    def exists(self) -> bool:
+        """The local environment always exists."""
+        return True

.github/triage/jax_toolbox_triage/main.py

@@ -13,6 +13,7 @@
 from .args import compulsory_software, optional_software, parse_args
 from .container import Container
 from .docker import DockerContainer
+from .local import LocalContainer
 from .logic import commit_search, container_search, TestResult
 from .pyxis import PyxisContainer
 from .utils import (
@@ -108,6 +109,9 @@ def test_output_directory(
     def Container(
         url, test_output_host_directory: typing.Optional[pathlib.Path] = None
     ):
+        if args.container_runtime == "local":
+            return LocalContainer(logger=logger)
+
         Imp = DockerContainer if args.container_runtime == "docker" else PyxisContainer
         mounts = bazel_cache_mounts + args.container_mount
         if test_output_host_directory is not None:
@@ -198,7 +202,10 @@ def check_container(date: datetime.date) -> TestResult:
             host_output_directory=out_dir, result=test_pass, stdouterr=result.stdout
         )
 
-    if args.passing_container is None and args.failing_container is None:
+    if args.container_runtime == "local":
+        passing_url = "local"
+        failing_url = "local"
+    elif args.passing_container is None and args.failing_container is None:
         # Search through the published containers, narrowing down to a pair of dates with
         # the property that the test passed on `range_start` and fails on `range_end`.
         range_start, range_end = container_search(
@@ -235,7 +242,10 @@ def check_container(date: datetime.date) -> TestResult:
 
     # Choose a container to do the commit-level bisection in; use directory
     # names that match it.
-    if failing_url is not None:
+    if args.container_runtime == "local":
+        bisection_url = "local"
+        package_dirs = failing_package_dirs
+    elif failing_url is not None:
         bisection_url = failing_url
         package_dirs = failing_package_dirs
     else:

.github/triage/jax_toolbox_triage/utils.py

@@ -82,14 +82,18 @@ def prepare_bazel_cache_mounts(
 
 
 def run_and_log(
-    command, logger: logging.Logger, stderr: typing.Literal["interleaved", "separate"]
+    command,
+    logger: logging.Logger,
+    stderr: typing.Literal["interleaved", "separate"],
+    cwd: typing.Optional[str] = None,
 ) -> subprocess.CompletedProcess:
-    logger.debug(shlex.join(command))
+    logger.debug(f"Executing in {cwd or '.'}: {shlex.join(command)}")
     result = subprocess.Popen(
         command,
         encoding="utf-8",
         stderr=subprocess.STDOUT if stderr == "interleaved" else subprocess.PIPE,
         stdout=subprocess.PIPE,
+        cwd=cwd,
     )
     assert result.stdout is not None
     stdouterr = ""

.github/triage/tests/test_arg_parsing.py

@@ -44,11 +44,19 @@
     ["--container", "jax", "--end-date", "2024-10-02"],
     ["--container", "jax", "--start-date", "2024-10-01", "--end-date", "2024-10-02"],
 ]
+valid_local_args = [
+    "--container-runtime",
+    "local",
+    "--passing-commits",
+    "jax:0123,xla:4567",
+    "--failing-commits",
+    "jax:89ab,xla:cdef",
+]
 
 
 @pytest.mark.parametrize(
     "good_args",
-    [valid_start_end_container]
+    [valid_start_end_container, valid_local_args]
     + valid_start_end_date_args
     + valid_container_and_commits,
 )
@@ -57,6 +65,14 @@ def test_good_container_args(good_args):
     assert args.test_command == test_command
 
 
+def test_good_local_args():
+    args = parse_args(valid_local_args + test_command)
+    assert args.test_command == test_command
+    assert args.container_runtime == "local"
+    assert "jax" in args.passing_commits
+    assert "xla" in args.failing_commits
+
+
 @pytest.mark.parametrize("date_args", valid_start_end_date_args)
 def test_bad_container_arg_combinations_across_groups(date_args):
     # Can't combine --{start,end}-container with --container/--{start,end}-date
@@ -129,3 +145,20 @@ def test_unparsable_container_args(container_args):
 def test_invalid_container_runtime():
     with pytest.raises(Exception):
         parse_args(["--container-runtime=magic-beans"] + test_command)
+
+
+@pytest.mark.parametrize(
+    "bad_local_args",
+    [
+        ["--container-runtime", "local"],
+        ["--container-runtime", "local", "--passing-commits", "jax:1,xla:2"],
+        ["--container-runtime", "local", "--failing-commits", "jax:1,xla:2"],
+        valid_local_args + ["--container", "jax"],
+        valid_local_args + ["--start-date", "2024-01-01"],
+        valid_local_args + ["--passing-container", "url"],
+        valid_local_args + ["--failing-container", "url"],
+    ],
+)
+def test_bad_local_arg_combinations(bad_local_args):
+    with pytest.raises(Exception):
+        parse_args(bad_local_args + test_command)

docs/triage-tool.md

@@ -3,21 +3,20 @@
 `jax-toolbox-triage` is a tool to automate the process of attributing regressions to an
 individual commit of JAX or XLA.
 It takes as input a command that returns an error (non-zero) code when run in "recent"
-containers, but which returns a success (zero) code when run in some "older" container.
-The command must be executable within the containers, *i.e.* it cannot refer to files
+environments, but which returns a success (zero) code when run in some "older" environment.
+The command must be executable within the triage environment, *i.e.* it cannot refer to files
 that only exist on the host system, unless those are explicitly mounted in using the
-`-v` (`--container-mount`) option.
+`-v` (`--container-mount`) option when using a container-based runtime.
 
-The tool follows a three-step process:
+The tool follows a process that can include up to three steps:
   1. A container-level search backwards from the "recent" container where the test is
      known to fail, which identifies an "older" container where the test passes. This
      search proceeds with an exponentially increasing step size and is based on the
      `YYYY-MM-DD` tags under `ghcr.io/nvidia/jax`.
   2. A container-level binary search to refine this to the **latest** available
      container where test passes and the **earliest** available container where it
      fails.
-  3. A commit-level binary search, repeatedly building + testing inside the same
-     container, to identify a single commit of a software package known to the tool
+  3. A commit-level binary search, repeatedly building + testing, to identify a single commit of a software package known to the tool
      (JAX, XLA, Flax, optionally MaxText) that causes the test to start failing, and a
      set of reference commits of the other packages that can be used to reproduce the
      regression.
@@ -64,20 +63,20 @@ or more machines with appropriate GPUs, *e.g.* inside an `salloc` session.
 Appropriate arguments (number of nodes, number of tasks per node, *etc.*) should be
 passed to `salloc` or set via `SLURM_` environment variables so that a bare `srun` will
 correctly launch the test case.
+If `--container-runtime=local` is used, the tool assumes it is already inside a JAX container and will execute all build and test commands directly.
 
 ## Usage
 
 To use the tool, there are two compulsory inputs:
    * A test command to triage.
-   * A specification of which containers to triage in. There are two choices here:
-     * `--container`: which of the `ghcr.io/nvidia/jax:CONTAINER-YYYY-MM-DD` container
-       families to execute the test command in. Example: `jax` for a JAX unit test
-       failure, `maxtext` for a MaxText model execution failure. The `--start-date` and
+   * A specification of the triage scope. There are three choices here:
+      * **Container Search**: Usage `--container` to specify which of the `ghcr.io/nvidia/jax:CONTAINER-YYYY-MM-DD` container families to search through. Example: `jax` for a JAX unit test failure, `maxtext` for a MaxText model execution failure. The `--start-date` and
        `--end-date` options can be combined with `--container` to tune the search; see
        below for more details.
-     * `--passing-container` and `--failing-container`: a pair of URLs to containers to
+     * **Commit Search between Containers**: `--passing-container` and `--failing-container`: a pair of URLs to containers to
        use in the commit-level search; if these are passed then no container-level
        search is performed.
+      * **Local Commit Search**: Use `--container-runtime=local` when you are already inside a JAX container. This mode skips all container orchestration and performs a commit-level search directly in the local container. It requires you to specify the commit range with `--passing-commits` and `--failing-commits`.
 
 The test command will be executed directly in the container, not inside a shell, so be
 sure not to add excessive quotation marks (*i.e.* run
@@ -88,7 +87,7 @@ as fast and targeted as possible.
 If you want to run multiple commands, you might want to use something like
 `jax-toolbox-triage --container=jax sh -c "command1 && command2"`.
 
-Alternatively, you can use `-v` (`--container-mount`) to mount a host directory
+Alternatively, when using a container runtime, you can use `-v` (`--container-mount`) to mount a host directory
 containing test scripts into the container and execute a script from there, *e.g.*
 `-v $PWD:/work /work/test.sh`.