triage-tool: support multi-node triage (#1421)

- Add `--{passing,failing}-commits` option, which allows the
commit-level search to be manually restricted in scope
- This also allows single-container triage, as you can pass
`--{passing,failing}-container` and `--{failing,passing}-commits`
- With the Pyxis backend, re-use the same container instance within a
single triage tool process
  - This reduces the amount of time spent in container creation
- Support multi-node/multi-process triage with the Pyxis backend
- This is implemented by annotating the various commands run inside the
containers as `once` (run once, in one container instance),
`once_per_container` (run once per container instance, i.e. once per
node), and `default` (run without extra `srun` arguments -- the caller
must make this do the right thing e.g. by passing `--ntasks-per-node` to
`salloc`.
    -  `once` example: getting the JAX commit from a container
    - `once_per_container` example: building JAX
    - `default` example: running the test case

Author: olupton

.github/triage/jax_toolbox_triage/args.py

@@ -6,6 +6,16 @@
 import tempfile
 
 
+def parse_commit_argument(s):
+    ret = {}
+    for part in s.split(","):
+        sw, commit = part.split(":", 1)
+        assert sw not in ret, ret
+        ret[sw] = commit
+    assert ret.keys() == {"jax", "xla"}, ret.keys()
+    return ret
+
+
 def parse_args(args=None):
     parser = argparse.ArgumentParser(
         description="""
@@ -70,10 +80,11 @@ def parse_args(args=None):
         "--failing-container",
         help="""
             Skip the container-level search and pass this container to the commit-level
-            search. If this is passed, --passing-container must be too, but --container
-            is not required. This can be used to apply the commit-level bisection
-            search to containers not from the ghcr.io/nvidia/jax:CONTAINER-YYYY-MM-DD
-            series, although they must have a similar structure.""",
+            search. If this is passed, --passing-container or --passing-commits must be
+            too, but --container is not required. This can be used to apply the
+            commit-level bisection search to containers not from the
+            ghcr.io/nvidia/jax:CONTAINER-YYYY-MM-DD series, although they must have a
+            similar structure.""",
     )
     container_search_args.add_argument(
         "--end-date",
@@ -88,10 +99,11 @@ def parse_args(args=None):
         "--passing-container",
         help="""
             Skip the container-level search and pass this container to the commit-level
-            search. If this is passed, --failing-container must be too, but --container is
-            not required. This can be used to apply the commit-level bisection search
-            to containers not from the ghcr.io/nvidia/jax:CONTAINER-YYYY-MM-DD series,
-            although they must have a similar structure.""",
+            search. If this is passed, --failing-container or --failing-commits must be
+            too, but --container is not required. This can be used to apply the
+            commit-level bisection search to containers not from the
+            ghcr.io/nvidia/jax:CONTAINER-YYYY-MM-DD series, although they must have a
+            similar structure.""",
     )
     container_search_args.add_argument(
         "--start-date",
@@ -126,6 +138,24 @@ def parse_args(args=None):
             significantly speed up the commit-level search. By default, uses a temporary
             directory including the name of the current user.""",
     )
+    commit_search_args.add_argument(
+        "--failing-commits",
+        help="""
+            When combined with --passing-container, the commit-level triage will use
+            that container and --failing-commits will specify the end of the commit
+            range, rather than the commits being extracted from --failing-container.
+            Expects an argument of form jax:jax_commit_hash,xla:xla_commit_hash.""",
+        type=parse_commit_argument,
+    )
+    commit_search_args.add_argument(
+        "--passing-commits",
+        help="""
+            When combined with --failing-container, the commit-level triage will use
+            that container and --passing-commits will specify the start of the commit
+            range, rather than the commits being extracted from --passing-container.
+            Expects an argument of form jax:jax_commit_hash,xla:xla_commit_hash.""",
+        type=parse_commit_argument,
+    )
     parser.add_argument(
         "-v",
         "--container-mount",
@@ -145,29 +175,33 @@ def parse_args(args=None):
     )
     args = parser.parse_args(args=args)
     assert args.container_runtime in {"docker", "pyxis"}, args.container_runtime
-    num_explicit_containers = (args.passing_container is not None) + (
-        args.failing_container is not None
-    )
-    if num_explicit_containers == 1:
-        raise Exception(
-            "--passing-container and --failing-container must both be passed if either is"
+    passing_commits_known = (args.passing_container is not None) or (
+        args.passing_commits is not None
+    )
+    failing_commits_known = (args.failing_container is not None) or (
+        args.failing_commits is not None
+    )
+    sets_of_known_commits = passing_commits_known + failing_commits_known
+    if sets_of_known_commits == 2:
+        # If the container-level search is being skipped, because a valid combination
+        # of --{passing,failing}-{commits,container} is passed, then no container-level
+        # search options should be passed.
+        assert (
+            args.container is None and args.start_date is None and args.end_date is None
+        ), (
+            "No container-level search options should be passed if the passing/failing containers/commits have been passed explicitly."
         )
-    if num_explicit_containers == 2:
-        # Explicit mode, --container, --start-date and --end-date are all ignored
-        if args.container:
-            raise Exception(
-                "--container must not be passed if --passing-container and --failing-container are"
-            )
-        if args.start_date:
-            raise Exception(
-                "--start-date must not be passed if --passing-container and --failing-container are"
-            )
-        if args.end_date:
-            raise Exception(
-                "--end-date must not be passed if --passing-container and --failing-container are"
-            )
-    elif num_explicit_containers == 0 and args.container is None:
+        assert (
+            args.passing_container is not None or args.failing_container is not None
+        ), ""
+    elif sets_of_known_commits == 1:
         raise Exception(
-            "--container must be passed if --passing-container and --failing-container are not"
+            "If --passing-{commits OR container} is passed then --failing-{commits OR container} should be too"
+        )
+    else:
+        # None of --{passing,failing}-{commits,container} were passed, make sure the
+        # compulsory arguments for the container-level search were passed
+        assert args.container is not None, (
+            "--container must be passed for the container-level search"
         )
     return args

.github/triage/jax_toolbox_triage/docker.py

@@ -56,7 +56,10 @@ def __repr__(self):
         return f"Docker({self._url})"
 
     def exec(
-        self, command: typing.List[str], workdir=None
+        self,
+        command: typing.List[str],
+        policy: typing.Literal["once", "once_per_container", "default"] = "default",
+        workdir=None,
     ) -> subprocess.CompletedProcess:
         """
         Run a command inside a persistent container.

.github/triage/jax_toolbox_triage/logic.py

@@ -214,43 +214,36 @@ def commit_search(
     end_jax_commit = jax_commits[-1][0]
     end_xla_commit = xla_commits[-1][0]
     if skip_precondition_checks:
-        logger.info("Skipping check that vanilla rebuild + test reproduces failure")
+        logger.info("Skipping check that 'bad' commits reproduce failure")
     else:
-        # Verify we can build successfully and that the test fails as expected. These
-        # commits are the ones already checked out in the container, but specifying
-        # them explicitly is good for the summary JSON.
-        logger.info("Building in the range-ending container...")
+        # Verify we can build successfully and that the test fails as expected.
+        logger.info("Verifying test failure using 'bad' commits")
         range_end_result, stdout, stderr = build_and_test(
             jax_commit=end_jax_commit, xla_commit=end_xla_commit
         )
         if not range_end_result:
-            logger.info("Verified test failure after vanilla rebuild")
+            logger.info("Verified test failure using 'bad' commits")
         else:
-            logger.fatal("Vanilla rebuild did not reproduce test failure")
+            logger.fatal("Could not reproduce failure with 'bad' commits")
             logger.fatal(stdout)
             logger.fatal(stderr)
-            raise Exception(
-                "Could not reproduce failure after rebuild in 'bad' container"
-            )
+            raise Exception("Could not reproduce failure with 'bad' commits")
 
-    # Verify that we can build the commit at the start of the range and reproduce the
-    # test success there in the end-of-range container.
-    range_start_result, stdout, stderr = build_and_test(
-        jax_commit=start_jax_commit, xla_commit=start_xla_commit
-    )
-    if range_start_result:
-        logger.info(
-            "Test passed after rebuilding commits from start container in end container"
-        )
+    if skip_precondition_checks:
+        logger.info("Skipping check that 'good' commits reproduce success")
     else:
-        logger.fatal(
-            "Test failed after rebuilding commits from start container in end container"
-        )
-        logger.fatal(stdout)
-        logger.fatal(stderr)
-        raise Exception(
-            "Could not reproduce success with 'good' commits in 'bad' container"
+        # Verify that we can build successfully and that the test succeeds as expected.
+        logger.info("Verifying test success using 'good' commits")
+        range_start_result, stdout, stderr = build_and_test(
+            jax_commit=start_jax_commit, xla_commit=start_xla_commit
         )
+        if range_start_result:
+            logger.info("Verified test passes using 'good' commits")
+        else:
+            logger.fatal("Could not reproduce success with 'good' commits")
+            logger.fatal(stdout)
+            logger.fatal(stderr)
+            raise Exception("Could not reproduce success with 'good' commits")
 
     # Finally, start bisecting. This is XLA-centric; JAX is moved too but is secondary.
     while len(xla_commits) > 2:
@@ -260,6 +253,9 @@ def commit_search(
         for jax_index, (jax_hash, jax_date) in enumerate(jax_commits):
             if jax_date >= xla_date:
                 break
+        logger.info(
+            f"Chose from {len(xla_commits)} remaining XLA commits and {len(jax_commits)} remaining JAX commits"
+        )
         bisect_result, _, _ = build_and_test(jax_commit=jax_hash, xla_commit=xla_hash)
         if bisect_result:
             # Test passed, continue searching in the second half

.github/triage/jax_toolbox_triage/main.py

@@ -32,7 +32,9 @@ def get_commit(
     results = []
     for suffix in ["", "-source"]:
         dirname = f"/opt/{repo}{suffix}"
-        result = container.exec(["git", "rev-parse", "HEAD"], workdir=dirname)
+        result = container.exec(
+            ["git", "rev-parse", "HEAD"], policy="once", workdir=dirname
+        )
         results.append(result)
         if result.returncode == 0:
             commit = result.stdout.strip()
@@ -108,12 +110,7 @@ def check_container(date: datetime.date) -> TestResult:
         )
         return TestResult(result=test_pass, stdout=result.stdout, stderr=result.stderr)
 
-    if args.passing_container is not None:
-        assert args.failing_container is not None
-        # Skip the container-level search because explicit end points were given
-        passing_url = args.passing_container
-        failing_url = args.failing_container
-    else:
+    if 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(
@@ -127,20 +124,44 @@ def check_container(date: datetime.date) -> TestResult:
         )
         passing_url = container_url(range_start)
         failing_url = container_url(range_end)
+    else:
+        # Skip the container-level search because at lease one explicit end point was
+        # given
+        passing_url = args.passing_container
+        failing_url = args.failing_container
+
+    jax_dir = "/opt/jax"
+    xla_dir = "/opt/xla"
+    if args.passing_commits is not None:
+        start_jax_commit = args.passing_commits["jax"]
+        start_xla_commit = args.passing_commits["xla"]
+    else:
+        assert passing_url is not None
+        with Container(passing_url) as worker:
+            start_jax_commit, jax_dir = get_commit(worker, "jax")
+            start_xla_commit, xla_dir = get_commit(worker, "xla")
+
+    if args.failing_commits is not None:
+        end_jax_commit = args.failing_commits["jax"]
+        end_xla_commit = args.failing_commits["xla"]
+    else:
+        assert failing_url is not None
+        with Container(failing_url) as worker:
+            end_jax_commit, jax_dir = get_commit(worker, "jax")
+            end_xla_commit, xla_dir = get_commit(worker, "xla")
+
+    bisection_url = failing_url or passing_url
+    assert bisection_url is not None
+    assert jax_dir is not None
+    assert xla_dir is not None
 
     # Container-level search is now complete. Triage proceeds inside the `range_end``
     # container. First, we check that rewinding JAX and XLA inside the `range_end``
     # container to the commits used in the `range_start` container passes, whereas
     # using the `range_end` commits reproduces the failure.
 
-    with Container(passing_url) as worker:
-        start_jax_commit, _ = get_commit(worker, "jax")
-        start_xla_commit, _ = get_commit(worker, "xla")
-
     # Fire up the container that will be used for the fine search.
-    with Container(failing_url) as worker:
-        end_jax_commit, jax_dir = get_commit(worker, "jax")
-        end_xla_commit, xla_dir = get_commit(worker, "xla")
+    with Container(bisection_url) as worker:
         logger.info(
             (
                 f"Bisecting JAX [{start_jax_commit}, {end_jax_commit}] and "
@@ -150,6 +171,9 @@ def check_container(date: datetime.date) -> TestResult:
 
         # Get the full lists of JAX/XLA commits and dates
         def commits(start, end, dir):
+            worker.check_exec(
+                ["git", "fetch"], policy="once_per_container", workdir=dir
+            )
             result = worker.check_exec(
                 [
                     "git",
@@ -159,6 +183,7 @@ def commits(start, end, dir):
                     "--format=%H %cI",
                     f"{start}^..{end}",
                 ],
+                policy="once",
                 workdir=dir,
             )
             data = []
@@ -190,21 +215,39 @@ def build_and_test(
             jaxlib, and run the test command. Throws on error when checking out or
             building, and returns the status of the test command.
             """
-            worker.check_exec(["git", "stash"], workdir=xla_dir)
-            worker.check_exec(["git", "stash"], workdir=jax_dir)
-            worker.check_exec(["git", "checkout", xla_commit], workdir=xla_dir)
-            worker.check_exec(["git", "checkout", jax_commit], workdir=jax_dir)
-            logger.info(f"Checking out XLA {xla_commit} JAX {jax_commit}")
+            worker.check_exec(
+                ["git", "stash"], policy="once_per_container", workdir=xla_dir
+            )
+            worker.check_exec(
+                ["git", "stash"], policy="once_per_container", workdir=jax_dir
+            )
+            worker.check_exec(
+                ["git", "checkout", xla_commit],
+                policy="once_per_container",
+                workdir=xla_dir,
+            )
+            worker.check_exec(
+                ["git", "checkout", jax_commit],
+                policy="once_per_container",
+                workdir=jax_dir,
+            )
+            logger.info(f"Checking out XLA {xla_commit} JAX {jax_commit} in {worker}")
             # Build JAX
             before = time.monotonic()
             # Unfortunately the build system does not always seem to handle incremental
             # rebuilds correctly.
-            worker.check_exec(["bazel", "clean", "--expunge"], workdir=jax_dir)
+            worker.check_exec(
+                ["bazel", "clean", "--expunge"],
+                policy="once_per_container",
+                workdir=jax_dir,
+            )
             build_jax = [
                 "build-jax.sh",
                 f"--bazel-cache={args.bazel_cache}",
             ]
-            build_result = worker.check_exec(build_jax, workdir=jax_dir)
+            build_result = worker.check_exec(
+                build_jax, policy="once_per_container", workdir=jax_dir
+            )
             middle = time.monotonic()
             logger.info(f"Build completed in {middle - before:.1f}s")
             logger.debug(
@@ -217,7 +260,7 @@ def build_and_test(
                 "commit",
                 {
                     "build_time": middle - before,
-                    "container": failing_url,
+                    "container": bisection_url,
                     "jax": jax_commit,
                     "result": test_result.returncode == 0,
                     "test_time": test_time,