Triage tool deals with non-linear history

.github/triage/jax_toolbox_triage/args.py

@@ -26,6 +26,23 @@ def parse_version_argument(s: str) -> typing.Dict[str, str]:
     return ret
 
 
+def parse_override_remotes(s: str) -> typing.Dict[str, str]:
+    """Function to parse the override remote
+
+    Inputs:
+        s: (str) e.g. https://<token>@host/repo.git
+
+    Returns:
+        ret: (typing.Dict[str,str]) Dictionary with software as key and git-url as value.
+    """
+    ret: typing.Dict[str, str] = {}
+    for part in s.split(","):
+        sw, url = part.split(":", 1)
+        assert sw not in ret, ret
+        ret[sw] = url
+    return ret
+
+
 def parse_args(args=None) -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="""
@@ -208,6 +225,14 @@ def parse_args(args=None) -> argparse.Namespace:
             in question has different versions at the endpoints of the bisection range.
         """,
     )
+    version_search_args.add_argument(
+        "--override-remotes",
+        type=parse_override_remotes,
+        default={},
+        help="""Remote URLs to be used for fetching, including auth token. E.g.:
+            jax:https://<token>@host/repo.git,xla:https://<token>@host/repo.git
+            """,
+    )
     parser.add_argument(
         "-v",
         "--container-mount",
@@ -225,10 +250,18 @@ def parse_args(args=None) -> argparse.Namespace:
         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", "local"}, (
-        args.container_runtime
+    parser.add_argument(
+        "--main-branch",
+        type=str,
+        default="main",
+        help="The name of the main branch (e.g. main) to derive cherry-picks from",
     )
+    args = parser.parse_args(args=args)
+    assert args.container_runtime in {
+        "docker",
+        "pyxis",
+        "local",
+    }, args.container_runtime
     # --{passing,failing}-commits are deprecated aliases for --{passing,failing}-versions.
     for prefix in ["passing", "failing"]:
         commits = getattr(args, f"{prefix}_commits")

.github/triage/jax_toolbox_triage/bisect.py

@@ -0,0 +1,105 @@
+import datetime
+import subprocess
+
+
+def get_commit_history(
+    worker,
+    package,
+    start,
+    end,
+    dir,
+    main_branch,
+    logger=None,
+    args=None,
+):
+    """
+    Get the commit history for a given package between two commits.
+
+    Args:
+        worker (Container): The container worker to execute commands.
+        package (str): The name of the package.
+        start (str): The starting commit hash.
+        end (str): The ending commit hash.
+        dir (str): The directory where the git repository is located.
+        main_branch (str): The main branch name. Defaults is the default branch of the repo.
+        logger (Logger, optional): Logger for debug information. Defaults to None.
+        args: Additional arguments that may contain cherry-pick commits.
+
+    Returns:
+        data: list, list of all the commits
+        cherry_pick_range: str, range of cherry pick commits if any
+    """
+    # In particular the end commit might not already be known if the older,
+    # passing, container is being used for triage.
+    commits_known = worker.exec(
+        [
+            "sh",
+            "-c",
+            f"git cat-file commit {start} && git cat-file commit {end}",
+        ],
+        policy="once_per_container",
+        workdir=dir,
+    )
+    if commits_known.returncode != 0:
+        worker.check_exec(
+            ["git", "fetch", args.override_remotes.get(package, "origin"), start, end],
+            workdir=dir,
+        )
+
+    # detect non-linear history
+    is_ancestor_result = worker.exec(
+        ["git", "merge-base", "--is-ancestor", start, end],
+        workdir=dir,
+    )
+    is_linear = is_ancestor_result.returncode == 0
+    cherry_pick_range = {}
+
+    if not is_linear:
+        logger.info(f"Using non-linear history logic with branch {main_branch}")
+
+        # 1. find the linear range on the main branch
+        passing_and_failing_cmd = worker.check_exec(
+            [
+                "sh",
+                "-c",
+                f"git merge-base {start} {end} && git merge-base {end} {main_branch}",
+            ],
+            workdir=dir,
+        ).stdout.strip()
+        passing_main_commit, failing_main_commit = passing_and_failing_cmd.splitlines()
+
+        # 2. find commits to cherry-pick from the failing branch
+        # TODO: as an alternative approach we may need to consider `{passing_main_commit}..{start}`
+        cherry_pick_range[package] = f"{failing_main_commit}..{end}"
+
+        # 3. now we can use the main branch  commits for bisection
+        start = passing_main_commit
+        end = failing_main_commit
+
+    logger.info(f"cherry_pick_range: {cherry_pick_range}, start: {start}, end: {end}")
+
+    # now create the right git command to retrieve the history between start..end
+    result = worker.check_exec(
+        [
+            "git",
+            "log",
+            "--first-parent",
+            "--reverse",
+            "--format=%H %cI",
+            f"{start}^..{end}",
+        ],
+        policy="once",
+        stderr=subprocess.PIPE,
+        workdir=dir,
+    )
+
+    data = []
+    for line in result.stdout.splitlines():
+        commit, date = line.split()
+        # for python < 3.11 we nee dto fix:
+        if date.endswith("Z"):
+            date = date[:-1] + "+00:00"
+        date = datetime.datetime.fromisoformat(date).astimezone(datetime.timezone.utc)
+        data.append((commit, date))
+
+    return data, cherry_pick_range

.github/triage/jax_toolbox_triage/container_factory.py

@@ -0,0 +1,28 @@
+import logging
+from .container import Container
+from .docker import DockerContainer
+from .pyxis import PyxisContainer
+from .local import LocalContainer
+
+
+def make_container(
+    runtime: str, url: str, mounts: list, logger: logging.Logger, **kwargs
+) -> Container:
+    """
+    This function creates a container object, based on the specified runtime
+
+    Args:
+        runtime (str): The container runtime to use (e.g., 'docker', 'pyxis', 'local').
+        url (str): The URL of the container.
+        mounts (list): List of mounts to be used in the container.
+        logger (logging.Logger): Logger instance for logging messages.
+        **kwargs: Additional keyword arguments for specific container types.
+
+    Returns:
+        Container: A container class associated with the specified runtime.
+    """
+    if runtime == "local":
+        return LocalContainer(logger=logger)
+
+    container_impl = DockerContainer if runtime == "docker" else PyxisContainer
+    return container_impl(url, logger=logger, mounts=mounts, **kwargs)