triage: --container-url-template (#1493)

This allows the container-level search to use a different series of
dated container images.

Author: olupton

.github/triage/jax_toolbox_triage/args.py

@@ -57,6 +57,14 @@ def parse_args(args=None) -> argparse.Namespace:
             Container to use. Example: jax, maxtext. Used to construct the URLs of
             nightly containers, like ghcr.io/nvidia/jax:CONTAINER-YYYY-MM-DD.""",
     )
+    container_search_args.add_argument(
+        "--container-url-template",
+        type=str,
+        help="""
+            Container URL pattern as a Python format string into which `container` and
+            `date` will be substituted, e.g. ghcr.io/nvidia/jax:{container}-{date} for
+            the JAX-Toolbox public nightlies.""",
+    )
     parser.add_argument(
         "--output-prefix",
         default=datetime.datetime.now().strftime("triage-%Y-%m-%d-%H-%M-%S"),

.github/triage/jax_toolbox_triage/main.py

@@ -99,7 +99,11 @@ def test_output_directory(
         out_dir.mkdir(mode=0o755)
         return out_dir.resolve()
 
-    container_url = functools.partial(container_url_base, container=args.container)
+    container_url = functools.partial(
+        container_url_base,
+        container=args.container,
+        template=args.container_url_template,
+    )
 
     def Container(
         url, test_output_host_directory: typing.Optional[pathlib.Path] = None

.github/triage/jax_toolbox_triage/utils.py

@@ -7,18 +7,23 @@
 import typing
 
 
-def container_url(date: datetime.date, *, container: str) -> str:
+def container_url(
+    date: datetime.date, *, container: str, template: typing.Optional[str] = None
+) -> str:
     """
     Construct the URL for --container on the given date.
 
     Arguments:
     date: YYYY-MM-DD format.
     """
+    date_str = date.isoformat()
+    if template is not None:
+        return template.format(container=container, date=date_str)
     # Around 2024-02-09 the naming scheme changed.
-    if date > datetime.date(year=2024, month=2, day=9):
-        return f"ghcr.io/nvidia/jax:{container}-{date.isoformat()}"
+    elif date > datetime.date(year=2024, month=2, day=9):
+        return f"ghcr.io/nvidia/jax:{container}-{date_str}"
     else:
-        return f"ghcr.io/nvidia/{container}:nightly-{date.isoformat()}"
+        return f"ghcr.io/nvidia/{container}:nightly-{date_str}"
 
 
 def get_logger(output_prefix: pathlib.Path) -> logging.Logger: