triage: allow container+commits arg (#1482)

This allows combining `--{passing,failing}-container` with
`--{passing,failing}-commits`.
The behaviour is that entries in the latter override values read from
the former.

Author: olupton

.github/triage/jax_toolbox_triage/args.py

@@ -5,14 +5,21 @@
 import pathlib
 import tempfile
 
+# Software we know may exist in the containers that we might be able to triage
+# We know how to recompile JAX/XLA, so it's OK that they include C++ code
+# TransformerEngine is intentionally excluded because build-te.sh is not plumbed yet.
+# Flax and MaxText are pure Python, so it's OK we don't have a way of compiling them,
+# but they are not always installed in containers we want to triage.
+compulsory_software = {"xla", "jax"}
+optional_software = {"flax", "maxtext"}
+
 
 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
 
 
@@ -189,14 +196,24 @@ def parse_args(args=None):
         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."
+            "No container-level search options should be passed if the passing/failing"
+            " containers/commits have been passed explicitly."
         )
         assert (
             args.passing_container is not None or args.failing_container is not None
-        ), ""
+        ), "At least one of --passing-container and --failing-container must be passed."
+        for prefix in ["passing", "failing"]:
+            assert (
+                getattr(args, f"{prefix}_container") is not None
+                or getattr(args, f"{prefix}_commits").keys() >= compulsory_software
+            ), (
+                f"--{prefix}-commits must specify all of {compulsory_software} if "
+                f"--{prefix}-container is not specified"
+            )
     elif sets_of_known_commits == 1:
         raise Exception(
-            "If --passing-{commits OR container} is passed then --failing-{commits OR container} should be too"
+            "If --passing-{commits AND/OR container} is passed then "
+            "--failing-{commits AND/OR container} should be too"
         )
     else:
         # None of --{passing,failing}-{commits,container} were passed, make sure the

.github/triage/jax_toolbox_triage/main.py

@@ -7,7 +7,7 @@
 import time
 import typing
 
-from .args import parse_args
+from .args import compulsory_software, optional_software, parse_args
 from .docker import DockerContainer
 from .logic import commit_search, container_search, TestResult
 from .pyxis import PyxisContainer
@@ -54,12 +54,6 @@ def get_commit(
 def get_commits_and_dirs(
     worker: typing.Union[DockerContainer, PyxisContainer], logger: logging.Logger
 ) -> typing.Tuple[typing.Dict[str, str], typing.Dict[str, str]]:
-    # Software we know may exist in the containers that we might be able to triage
-    # We know how to recompile JAX/XLA, so it's OK that they include C++ code
-    # TransformerEngine is intentionally excluded because there isn't a build-te.sh yet
-    # Flax and MaxText are pure Python, so it's OK we don't have a way of compiling them
-    compulsory_software = ["xla", "jax", "flax"]
-    optional_software = ["maxtext"]
     package_commits, dirs = {}, {}
     for package in compulsory_software:
         package_commits[package], dirs[package] = get_commit(worker, logger, package)
@@ -97,15 +91,22 @@ def get_commits(
         container, or return the explicitly passed set of hashes if they are
         given.
 
-        It is an error for both/neither of the arguments to be None.
+        If both arguments are non-None, the commits read from inside the given
+        container will be overwritten by explicit passed entries.
+
+        If is an error for both of the arguments to be None.
         """
         if explicit_commits is None:
             assert container_url is not None
             with Container(container_url) as worker:
                 return get_commits_and_dirs(worker, logger)
         else:
-            assert container_url is None
-            return explicit_commits, None
+            if container_url is None:
+                return explicit_commits, None
+            with Container(container_url) as worker:
+                commits, package_dirs = get_commits_and_dirs(worker, logger)
+            commits.update(explicit_commits)
+            return commits, package_dirs
 
     def add_summary_record(section, record, scalar=False):
         """

.github/triage/tests/test_arg_parsing.py

@@ -21,6 +21,22 @@
         "--passing-commits",
         "xla:fedcba9876543210,jax:0123456789",
     ],
+    [
+        "--passing-container",
+        "passing-url",
+        "--passing-commits",
+        "jax:123",  # xla not needed, the value from passing-container can be read
+        "--failing-container",
+        "failing-url",
+    ],
+    [
+        "--failing-container",
+        "failing-url",
+        "--failing-commits",
+        "xla:456",  # jax not needed, the value from failing-container can be read
+        "--passing-container",
+        "passing-url",
+    ],
 ]
 valid_start_end_date_args = [
     ["--container", "jax"],
@@ -67,14 +83,29 @@ def test_bad_container_arg_combinations_across_groups(date_args):
             "--failing-commits",
             "xla:fedcba9876543210,jax:0123456789",
         ],
-        # Cannot combine --passing-container with --passing-commits etc.
-        ["--passing-container", "passing-url", "--passing-commits", "jax:123,xla:456"],
-        ["--failing-container", "failing-url", "--failing-commits", "jax:123,xla:456"],
         # --{passing,failing}-commits must be formatted correctly
-        ["--passing-container", "passing-url", "--failing-commits", "jax:123"],
+        [
+            "--passing-container",
+            "passing-url",
+            "--failing-commits",
+            # no xla
+            "jax:123",
+        ],
         ["--passing-container", "passing-url", "--failing-commits", "jax:123,jax:456"],
-        ["--passing-container", "passing-url", "--failing-commits", "xla:123"],
-        ["--passing-container", "passing-url", "--failing-commits", "bob:123"],
+        [
+            "--passing-container",
+            "passing-url",
+            "--failing-commits",
+            # no jax
+            "xla:123",
+        ],
+        [
+            "--passing-container",
+            "passing-url",
+            "--failing-commits",
+            # neither jax nor xla
+            "bob:123",
+        ],
     ],
 )
 def test_bad_container_arg_combinations_within_groups(container_args):