feat: add TrainerRank API

.codex/environments/environment.toml

@@ -0,0 +1,21 @@
+version = 1
+name = "default"
+
+[setup]
+script = '''
+set -euo pipefail
+
+repo_root="$(git rev-parse --show-toplevel)"
+git_common_dir="$(git rev-parse --git-common-dir)"
+
+case "$git_common_dir" in
+  /*) common_dir="$git_common_dir" ;;
+  *) common_dir="$repo_root/$git_common_dir" ;;
+esac
+
+canonical_repo="$(dirname "$common_dir")"
+
+if [ "$canonical_repo" != "$repo_root" ] && [ -f "$canonical_repo/.env" ] && [ ! -e "$repo_root/.env" ]; then
+  install -m 600 "$canonical_repo/.env" "$repo_root/.env"
+fi
+'''

.githooks/post-checkout

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+repo_root="$(git rev-parse --show-toplevel 2>/dev/null || true)"
+git_common_dir="$(git rev-parse --git-common-dir 2>/dev/null || true)"
+
+if [ -z "$repo_root" ] || [ -z "$git_common_dir" ]; then
+  exit 0
+fi
+
+case "$git_common_dir" in
+  /*) common_dir="$git_common_dir" ;;
+  *) common_dir="$repo_root/$git_common_dir" ;;
+esac
+
+canonical_repo="$(dirname "$common_dir")"
+
+if [ "$canonical_repo" != "$repo_root" ] && [ -f "$canonical_repo/.env" ] && [ ! -e "$repo_root/.env" ]; then
+  install -m 600 "$canonical_repo/.env" "$repo_root/.env"
+fi

.githooks/post-index-change

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+repo_root="$(git rev-parse --show-toplevel 2>/dev/null || true)"
+git_common_dir="$(git rev-parse --git-common-dir 2>/dev/null || true)"
+
+if [ -z "$repo_root" ] || [ -z "$git_common_dir" ]; then
+  exit 0
+fi
+
+case "$git_common_dir" in
+  /*) common_dir="$git_common_dir" ;;
+  *) common_dir="$repo_root/$git_common_dir" ;;
+esac
+
+canonical_repo="$(dirname "$common_dir")"
+
+if [ "$canonical_repo" != "$repo_root" ] && [ -f "$canonical_repo/.env" ] && [ ! -e "$repo_root/.env" ]; then
+  install -m 600 "$canonical_repo/.env" "$repo_root/.env"
+fi

.github/workflows/prek.yml

@@ -82,7 +82,7 @@ jobs:
 
   build-cache:
     needs: cache-status
-    if: needs.cache-status.outputs.cache-hit != 'true'
+    if: github.event_name == 'push' && needs.cache-status.outputs.cache-hit != 'true'
     runs-on: art-cache-builder
     container:
       image: pytorch/pytorch:2.9.0-cuda12.8-cudnn9-devel
@@ -122,7 +122,7 @@ jobs:
             --python-mm "${CI_PYTHON_MM}"
 
   quality-checks:
-    needs: [cache-status, build-cache]
+    needs: cache-status
     if: ${{ !failure() && !cancelled() }}
     runs-on: art-large-runner
     container:
@@ -157,14 +157,86 @@ jobs:
             "${release_api}" || true)"
 
           if [ -z "${release_json}" ]; then
-            echo "::error::Missing cache release '${CI_UV_CACHE_RELEASE_TAG}'."
-            exit 1
+            echo "::warning::Missing cache release '${CI_UV_CACHE_RELEASE_TAG}'; continuing with an empty uv cache."
+            mkdir -p "${UV_CACHE_DIR}"
+            exit 0
           fi
 
           part_selection_file="/tmp/uv-cache-part-selection.txt"
-          if ! RELEASE_JSON="${release_json}" PART_PREFIX="${part_prefix}" python3 -c "import json, os, re, sys; payload=json.loads(os.environ['RELEASE_JSON']); part_prefix=os.environ['PART_PREFIX']; pattern=re.compile(r'^' + re.escape(part_prefix) + r'(\\d{3})$'); parts=[]; [parts.append((int(m.group(1)), int(a.get('id')), a.get('name'))) for a in payload.get('assets', []) for m in [pattern.match(a.get('name', ''))] if m and a.get('id') is not None]; parts.sort(key=lambda x: x[0]); indices=[p[0] for p in parts]; expected=list(range(len(parts))); print('\\n'.join(f'{asset_id} {name}' for _, asset_id, name in parts)) if parts and indices == expected else (_ for _ in ()).throw(SystemExit(2 if not parts else 3))" > "${part_selection_file}"; then
-            echo "::error::No complete uv cache part set found for prefix '${part_prefix}'."
-            exit 1
+          if ! RELEASE_JSON="${release_json}" EXACT_PREFIX="${part_prefix}" ASSET_PREFIX="${CI_UV_CACHE_ASSET_PREFIX}" python3 - <<'PY' > "${part_selection_file}"
+          import json
+          import os
+          import re
+          import sys
+
+          payload = json.loads(os.environ["RELEASE_JSON"])
+          exact_prefix = os.environ["EXACT_PREFIX"]
+          asset_prefix = os.environ["ASSET_PREFIX"]
+          exact_pattern = re.compile(r"^" + re.escape(exact_prefix) + r"(\d{3})$")
+          any_pattern = re.compile(
+              r"^"
+              + re.escape(asset_prefix)
+              + r"-([0-9a-f]+)\.tar\.zst\.part-(\d{3})$"
+          )
+
+          def complete(parts):
+              indices = [part[0] for part in parts]
+              return bool(parts) and indices == list(range(len(parts)))
+
+          def emit(parts):
+              for _, asset_id, name in parts:
+                  print(f"{asset_id} {name}")
+
+          exact_parts = []
+          groups = {}
+          for asset in payload.get("assets", []):
+              name = asset.get("name", "")
+              asset_id = asset.get("id")
+              if asset_id is None:
+                  continue
+              exact_match = exact_pattern.match(name)
+              if exact_match is not None:
+                  exact_parts.append((int(exact_match.group(1)), int(asset_id), name))
+              any_match = any_pattern.match(name)
+              if any_match is not None:
+                  fingerprint = any_match.group(1)
+                  group = groups.setdefault(
+                      fingerprint,
+                      {"updated_at": "", "parts": []},
+                  )
+                  group["updated_at"] = max(
+                      group["updated_at"],
+                      asset.get("updated_at") or asset.get("created_at") or "",
+                  )
+                  group["parts"].append((int(any_match.group(2)), int(asset_id), name))
+
+          exact_parts.sort(key=lambda part: part[0])
+          if complete(exact_parts):
+              emit(exact_parts)
+              sys.exit(0)
+
+          fallback_groups = []
+          for fingerprint, group in groups.items():
+              parts = group["parts"]
+              parts.sort(key=lambda part: part[0])
+              if complete(parts):
+                  fallback_groups.append(((group["updated_at"], fingerprint), parts))
+          if fallback_groups:
+              fallback_groups.sort(key=lambda item: item[0], reverse=True)
+              emit(fallback_groups[0][1])
+              print(
+                  "::warning::Exact uv cache is incomplete; using newest complete "
+                  "fallback cache.",
+                  file=sys.stderr,
+              )
+              sys.exit(0)
+
+          sys.exit(2)
+          PY
+          then
+            echo "::warning::No complete uv cache part set found for prefix '${part_prefix}' or any fallback; continuing with an empty uv cache."
+            mkdir -p "${UV_CACHE_DIR}"
+            exit 0
           fi
 
           part_count="$(wc -l < "${part_selection_file}" | tr -d ' ')"

.worktreeinclude

@@ -0,0 +1 @@
+.env

dev/megatron_review_perf.py

@@ -5,6 +5,7 @@
 import json
 from pathlib import Path
 import time
+from types import SimpleNamespace
 
 import numpy as np
 import torch
@@ -16,8 +17,8 @@
     build_block_mask_from_context,
     prepare_block_mask_context,
 )
-from art.megatron.context_parallel.builder import build_shared_prefix_attention_spec
-from art.megatron.context_parallel.executor import _build_stage_execution_spec
+from art.megatron.context_parallel.builder import build_prefix_tree_attention_spec
+from art.megatron.context_parallel.executor import _resolve_stage_execution_spec
 from art.megatron.context_parallel.runtime import (
     _RUNTIME_PLAN_CACHE,
     get_or_build_runtime_plan,
@@ -33,7 +34,7 @@
     normalize_sparse_block_size,
     sparse_compiled_flex_attention,
 )
-from art.megatron.shared_prefix_packing import SharedPrefixPack, pack_shared_prefixes
+from art.megatron.prefix_tree_packing import PrefixTreePack, prefix_tree_pack
 
 
 def main(
@@ -75,7 +76,7 @@ def main(
         branches_per_prefix=branches_per_prefix,
         completion_len=completion_len,
     )
-    spec = build_shared_prefix_attention_spec(
+    spec = build_prefix_tree_attention_spec(
         group_ids=pack.group_ids,
         parent_ids=pack.parent_ids,
     )
@@ -177,7 +178,7 @@ def main(
             branches_per_prefix=branches_per_prefix,
             completion_len=completion_len + variant * 11,
         )
-        variant_spec = build_shared_prefix_attention_spec(
+        variant_spec = build_prefix_tree_attention_spec(
             group_ids=variant_pack.group_ids,
             parent_ids=variant_pack.parent_ids,
         )
@@ -239,7 +240,7 @@ def _pack_workload(
     mid_prefix_len: int,
     branches_per_prefix: int,
     completion_len: int,
-) -> SharedPrefixPack:
+) -> PrefixTreePack:
     sequences = (
         _austin_sequences()
         if workload == "austin_198k"
@@ -254,7 +255,7 @@ def _pack_workload(
             completion_len=completion_len,
         )
     )
-    return pack_shared_prefixes(sequences, max_depth=max_depth)
+    return prefix_tree_pack(sequences, max_depth=max_depth)
 
 
 def _austin_sequences() -> tuple[torch.Tensor, ...]:
@@ -332,7 +333,7 @@ def _tokens(offset: int, length: int) -> torch.Tensor:
 
 
 def _build_cp_plan(
-    pack: SharedPrefixPack,
+    pack: PrefixTreePack,
     spec: object,
     topology: ParallelTopology,
     config: ContextParallelConfig,
@@ -346,7 +347,7 @@ def _build_cp_plan(
 
 
 def _build_stage_masks(
-    pack: SharedPrefixPack,
+    pack: PrefixTreePack,
     plan: object,
     config: ContextParallelConfig,
 ) -> tuple[tuple[BlockMask, tuple[object, ...]], ...]:
@@ -381,7 +382,7 @@ def _build_stage_masks(
 
 
 def _flex_records(
-    pack: SharedPrefixPack,
+    pack: PrefixTreePack,
     plan: object,
     config: ContextParallelConfig,
     *,
@@ -518,7 +519,7 @@ class _StageFlexCase:
 
 
 def _build_stage_flex_cases(
-    pack: SharedPrefixPack,
+    pack: PrefixTreePack,
     plan: object,
     config: ContextParallelConfig,
     *,
@@ -651,8 +652,9 @@ def _stage_execution_spec(
     stage: StagePlan,
     config: ContextParallelConfig,
 ) -> StageExecutionSpec:
-    return _build_stage_execution_spec(
+    return _resolve_stage_execution_spec(
         stage_plan=stage,
+        state=SimpleNamespace(config=config, execution_cache=None),
         block_size=_sparse_block_size(config),
     )
 
@@ -826,7 +828,7 @@ def _block_entries(
     return entries
 
 
-def _logical_tokens(pack: SharedPrefixPack) -> int:
+def _logical_tokens(pack: PrefixTreePack) -> int:
     return sum(int(positions.numel()) for positions in pack.positions_by_sequence)