Update fern workflows

Author: aybchan

.github/triage/jax_toolbox_triage/args.py

@@ -1,8 +1,11 @@
 import argparse
+import ast
 import datetime
 import getpass
 import os
 import pathlib
+import re
+import sys
 import tempfile
 import typing
 import warnings
@@ -16,6 +19,90 @@
 optional_software = ["flax", "maxtext", "transformer-engine"]
 
 
+_INFO_LOG_ARGUMENTS_RE = re.compile(
+    r"^\[INFO\]\s+\d{4}-\d{2}-\d{2}\s+\d{2}:\d{2}:\d{2}\s+Arguments:\s*$"
+)
+_INFO_LOG_ARGUMENT_RE = re.compile(
+    r"^\[INFO\]\s+\d{4}-\d{2}-\d{2}\s+\d{2}:\d{2}:\d{2}\s+"
+    r"(?P<name>[A-Za-z_][A-Za-z0-9_]*):\s*(?P<value>.*)$"
+)
+
+
+def logged_arg_value(s: str) -> typing.Any:
+    try:
+        return ast.literal_eval(s)
+    except (SyntaxError, ValueError):
+        return s
+
+
+def load_args_from_info_log(info_log: pathlib.Path) -> typing.Dict[str, str]:
+    ret: typing.Dict[str, str] = {}
+    in_arguments = False
+    with info_log.open("r", encoding="utf-8") as f:
+        for line in f:
+            if not in_arguments:
+                if _INFO_LOG_ARGUMENTS_RE.match(line):
+                    in_arguments = True
+                continue
+            match = _INFO_LOG_ARGUMENT_RE.match(line)
+            if not match:
+                if ret:
+                    break
+                continue
+            ret[match.group("name")] = match.group("value")
+    if not ret:
+        raise Exception(f"Could not read arguments from {info_log}")
+    return ret
+
+
+def parser_actions_by_dest(parser: argparse.ArgumentParser) -> typing.Dict[str, argparse.Action]:
+    return {
+        action.dest: action
+        for action in parser._actions
+        if action.dest != argparse.SUPPRESS
+    }
+
+
+def provided_arg_dests(
+    parser: argparse.ArgumentParser, args: typing.Sequence[str]
+) -> typing.Set[str]:
+    option_actions = {
+        option: action for action in parser._actions for option in action.option_strings
+    }
+    ret: typing.Set[str] = set()
+    i = 0
+    while i < len(args):
+        arg = args[i]
+        if arg == "--":
+            if i + 1 < len(args):
+                ret.add("test_command")
+            break
+        option = arg.split("=", 1)[0]
+        action = option_actions.get(option)
+        if action is None:
+            ret.add("test_command")
+            break
+        ret.add(action.dest)
+        if "=" in arg or action.nargs == 0:
+            i += 1
+        elif action.nargs is None:
+            i += 2
+        elif isinstance(action.nargs, int):
+            i += 1 + action.nargs
+        else:
+            i += 1
+    return ret
+
+
+def coerce_logged_arg(
+    action: argparse.Action, value: str
+) -> typing.Any:
+    value = logged_arg_value(value)
+    if isinstance(value, str) and action.type is not None:
+        return action.type(value)
+    return value
+
+
 def parse_cherry_picks(s: str) -> typing.Dict[str, typing.List[str]]:
     ret: typing.Dict[str, typing.List[str]] = {}
     for part in s.split(","):
@@ -128,9 +215,16 @@ def parse_args(args=None) -> argparse.Namespace:
             container. Careful use of this option, along with --start-date, --end-date and
             --threshold-days, allows the container-level search to be skipped.""",
     )
+    parser.add_argument(
+        "--no-skip-precondition-checks",
+        dest="skip_precondition_checks",
+        action="store_false",
+        default=False,
+        help="Re-enable precondition checks when restarting from a run that skipped them.",
+    )
     parser.add_argument(
         "test_command",
-        nargs="+",
+        nargs="*",
         help="""
             Command to execute inside the container. This should be as targeted as
             possible.""",
@@ -273,6 +367,13 @@ def parse_args(args=None) -> argparse.Namespace:
         action="store_true",
         help="Exclude transformer-engine from the list of optional software to triage.",
     )
+    version_search_args.add_argument(
+        "--no-exclude-transformer-engine",
+        dest="exclude_transformer_engine",
+        action="store_false",
+        default=False,
+        help="Include transformer-engine when restarting from a run that excluded it.",
+    )
     version_search_args.add_argument(
         "--transformer-engine-ccache-env",
         action="append",
@@ -316,7 +417,9 @@ def parse_args(args=None) -> argparse.Namespace:
         default="main",
         help="The name of the main branch (e.g. main) to derive cherry-picks from",
     )
-    args = parser.parse_args(args=args)
+    raw_args = sys.argv[1:] if args is None else list(args)
+    provided_args = provided_arg_dests(parser, raw_args)
+    args = parser.parse_args(args=raw_args)
     if args.restart:
         if args.output_prefix is None:
             raise Exception("--restart requires --output-prefix")
@@ -326,11 +429,21 @@ def parse_args(args=None) -> argparse.Namespace:
             raise Exception(
                 f"--output-prefix must contain summary.json: {summary_file}"
             )
+        info_log = args.output_prefix / "info.log"
+        if not info_log.exists():
+            raise Exception(f"--output-prefix must contain info.log: {info_log}")
+        actions_by_dest = parser_actions_by_dest(parser)
+        for arg, value in load_args_from_info_log(info_log).items():
+            if arg in provided_args or arg not in actions_by_dest:
+                continue
+            setattr(args, arg, coerce_logged_arg(actions_by_dest[arg], value))
     else:
         if args.output_prefix is None:
             args.output_prefix = pathlib.Path(
                 datetime.datetime.now().strftime("triage-%Y-%m-%d-%H-%M-%S")
             )
+    if not args.test_command:
+        raise Exception("test_command must be passed or restored from info.log")
 
     assert args.container_runtime in {
         "docker",
@@ -426,3 +539,4 @@ def parse_args(args=None) -> argparse.Namespace:
         args.optional_software.remove("transformer-engine")
 
     return args
+

.github/workflows/fern-docs-ci.yml

@@ -1,27 +1,24 @@
 # Validates Fern docs configuration and checks for broken internal links on PRs.
 # Runs without a FERN_TOKEN — suitable for untrusted forks.
-name: Fern Docs CI
+name: Fern Docs Validate
 
 on:
-  push:
-    branches:
-      - "pull-request/[0-9]+"
   pull_request:
+    types: [synchronize]
     paths:
       - 'docs/**'
       - 'fern/**'
       - 'package.json'
+      - '.github/workflows/*fern*.yml'
 
 jobs:
   validate:
     name: Validate docs
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v7
 
-      - uses: actions/setup-node@v4
-        with:
-          node-version: '20'
+      - uses: actions/setup-node@v6
 
       - name: Install dependencies
         run: npm install

.github/workflows/fern-docs-preview.yml

@@ -6,11 +6,11 @@ on:
   push:
     branches:
       - "pull-request/[0-9]+"
-  pull_request:
     paths:
       - 'docs/**'
       - 'fern/**'
       - 'package.json'
+      - '.github/workflows/*fern*.yml'
 
 jobs:
   preview:
@@ -19,11 +19,9 @@ jobs:
     permissions:
       pull-requests: write
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v7
 
-      - uses: actions/setup-node@v4
-        with:
-          node-version: '20'
+      - uses: actions/setup-node@v6
 
       - name: Install dependencies
         run: npm install
@@ -47,9 +45,9 @@ jobs:
         uses: actions/github-script@v7
         with:
           script: |
-            github.rest.issues.createComment({
+            await github.rest.issues.createComment({
               issue_number: context.issue.number,
               owner: context.repo.owner,
               repo: context.repo.repo,
-              body: '📖 **Docs preview:** ${{ steps.preview.outputs.url }}'
+              body: '📖 [**Docs preview**](${{ steps.preview.outputs.url }})'
             })

.github/workflows/publish-fern-docs.yml

@@ -17,11 +17,9 @@ jobs:
     name: Publish docs
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v7
 
-      - uses: actions/setup-node@v4
-        with:
-          node-version: '20'
+      - uses: actions/setup-node@v6
 
       - name: Install dependencies
         run: npm install