test: exercise TrainerRank GPU gate

.github/CODEOWNERS

@@ -0,0 +1 @@
+/src/art/trainer_rank/** @bradhilton

.github/workflows/trainer-rank-gpu.yml

@@ -0,0 +1,137 @@
+name: TrainerRank GPU validation
+
+on:
+  pull_request:
+    types: [opened, reopened, synchronize, ready_for_review]
+
+permissions:
+  contents: read
+
+concurrency:
+  group: trainer-rank-gpu-${{ github.event.pull_request.number }}
+  cancel-in-progress: true
+
+jobs:
+  classify:
+    runs-on: ubuntu-latest
+    outputs:
+      required: ${{ steps.changes.outputs.required }}
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - id: changes
+        name: Classify changed files
+        env:
+          BASE_SHA: ${{ github.event.pull_request.base.sha }}
+          HEAD_SHA: ${{ github.event.pull_request.head.sha }}
+        run: |
+          set -euo pipefail
+          pattern='^src/art/megatron/(prefix_tree(_packing|_state)?\.py|context_parallel/|flex_attn/|gdn/|lora\.py|megatron_patches\.py|training/(finalize_grads|microbatches)\.py)'
+          changed="$(git diff --name-only "${BASE_SHA}...${HEAD_SHA}")"
+          critical="$(printf '%s\n' "${changed}" | grep -E "${pattern}" || true)"
+          if [ -n "${critical}" ]; then
+            echo "required=true" >> "${GITHUB_OUTPUT}"
+            {
+              echo "### TrainerRank GPU validation required"
+              echo
+              echo '```text'
+              printf '%s\n' "${critical}"
+              echo '```'
+            } >> "${GITHUB_STEP_SUMMARY}"
+          else
+            echo "required=false" >> "${GITHUB_OUTPUT}"
+            echo "### TrainerRank GPU validation not required" >> "${GITHUB_STEP_SUMMARY}"
+          fi
+
+  validate:
+    name: Run on 2x H200
+    needs: classify
+    if: needs.classify.outputs.required == 'true'
+    runs-on: ubuntu-latest
+    timeout-minutes: 45
+    environment: trainer-rank-gpu-validation
+    env:
+      SKY_INFRA: k8s/cks-wb3
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event.pull_request.head.sha }}
+
+      - name: Install launch dependencies
+        run: |
+          curl -LsSf https://astral.sh/uv/install.sh | sh
+          echo "${HOME}/.local/bin" >> "${GITHUB_PATH}"
+          sudo apt-get update
+          sudo apt-get install -y --no-install-recommends socat
+          kubectl_version="$(curl -LsSf https://dl.k8s.io/release/stable.txt)"
+          kubectl_path="${RUNNER_TEMP}/kubectl"
+          curl -LsSf -o "${kubectl_path}" \
+            "https://dl.k8s.io/release/${kubectl_version}/bin/linux/amd64/kubectl"
+          sudo install -m 0755 "${kubectl_path}" /usr/local/bin/kubectl
+
+      - name: Configure Kubernetes
+        env:
+          KUBECONFIG_VALUE: ${{ secrets.CKS_WB3_KUBECONFIG }}
+        run: |
+          set -euo pipefail
+          test -n "${KUBECONFIG_VALUE}"
+          mkdir -p "${HOME}/.kube"
+          KUBECONFIG_PATH="${HOME}/.kube/config" python3 - <<'PY'
+          import base64
+          import binascii
+          import os
+          from pathlib import Path
+
+          value = os.environ["KUBECONFIG_VALUE"].encode()
+          try:
+              decoded = base64.b64decode(b"".join(value.split()), validate=True)
+              data = decoded if decoded.lstrip().startswith(b"apiVersion:") else value
+          except binascii.Error:
+              data = value
+          Path(os.environ["KUBECONFIG_PATH"]).write_bytes(data)
+          PY
+          chmod 600 "${HOME}/.kube/config"
+          kubectl --context cks-wb3 get nodes >/dev/null
+
+      - name: Run TrainerRank GPU validation
+        run: |
+          set -euo pipefail
+          sky_venv="${RUNNER_TEMP}/skypilot"
+          uv venv --python 3.10 "${sky_venv}"
+          uv pip install --python "${sky_venv}/bin/python" \
+            'skypilot[kubernetes,remote]==0.12.0' 'kubernetes==34.1.0'
+          sky="${sky_venv}/bin/sky"
+          cluster="trainer-rank-gpu-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}"
+          cleanup() { "${sky}" down -y "${cluster}" || true; }
+          trap cleanup EXIT
+          "${sky}" check kubernetes
+          "${sky}" launch -y -c "${cluster}" \
+            --infra "${SKY_INFRA}" \
+            scripts/ci/trainer-rank-gpu.sky.yaml
+
+  gate:
+    name: trainer-rank-gpu-validation
+    needs: [classify, validate]
+    if: always()
+    runs-on: ubuntu-latest
+    steps:
+      - env:
+          CLASSIFY_RESULT: ${{ needs.classify.result }}
+          REQUIRED: ${{ needs.classify.outputs.required }}
+          RESULT: ${{ needs.validate.result }}
+        run: |
+          if [ "${CLASSIFY_RESULT}" != "success" ]; then
+            echo "TrainerRank GPU classification failed." >&2
+            exit 1
+          fi
+          if [ "${REQUIRED}" = "true" ] && [ "${RESULT}" != "success" ]; then
+            echo "TrainerRank GPU validation required but result was ${RESULT}." >&2
+            exit 1
+          fi
+          if [ "${REQUIRED}" = "true" ]; then
+            echo "TrainerRank GPU validation passed."
+          else
+            echo "TrainerRank GPU validation was not required."
+          fi

scripts/ci/trainer-rank-gpu-tests.sh

@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+export CUDA_VISIBLE_DEVICES=0,1
+export PYTHONUNBUFFERED=1
+
+uv run --no-sync pytest --tb=short \
+  tests/integration/megatron/cp_attn/test_attention_packed_vs_flattened.py \
+  'tests/integration/megatron/gdn_shared_prefix/test_gdn_cp_packed_correctness.py::test_gdn_cp_packed_sibling_order_matches_cp1_oracle[2]' \
+  'tests/integration/megatron/gdn_shared_prefix/test_gdn_cp_packed_correctness.py::test_gdn_cp_tree_chain_matches_cp1_oracle[2]' \
+  'tests/integration/megatron/gdn_shared_prefix/test_gdn_cp_packed_correctness.py::test_gdn_cp_tree_trainability_updates_parameters[2]' \
+  tests/integration/megatron/gdn_shared_prefix/test_real_gdn_tp_lora.py::test_real_qwen35_gdn_tp2_gradients_match_flattened \
+  tests/integration/megatron/lora/test_dynamic_lora_slots.py::test_dynamic_lora_slots_capture_recompute_context_and_step_independently \
+  'tests/integration/megatron/lora/test_dynamic_lora_slots.py::test_trainer_rank_tp_head_backward_matches_unsharded_oracle[2]'
+
+ART_MEGATRON_CONTEXT_PARALLEL_SIZE=2 \
+  uv run --no-sync torchrun --standalone --nproc-per-node=2 \
+    dev/trainer_rank_check.py \
+    --model Qwen/Qwen3-0.6B \
+    --layers 1 \
+    --depths 0,4 \
+    --chunks 17,8192 \
+    --slots 0

scripts/ci/trainer-rank-gpu.sky.yaml

@@ -0,0 +1,27 @@
+name: trainer-rank-gpu-validation
+
+workdir: .
+
+resources:
+  accelerators: H200:2
+  image_id: docker:docker.io/bradhiltonnw/art-gpu:latest
+
+setup: |
+  uv sync --frozen --extra megatron --group dev
+
+run: |
+  timeout --signal=TERM --kill-after=30s 20m \
+    bash scripts/ci/trainer-rank-gpu-tests.sh
+
+config:
+  kubernetes:
+    pod_config:
+      spec:
+        schedulerName: binpack-scheduler
+        activeDeadlineSeconds: 3600
+        containers:
+          - name: ray-node
+            imagePullPolicy: Always
+            env:
+              - name: UV_LINK_MODE
+                value: copy

src/art/megatron/prefix_tree.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+# Temporary no-op used to validate the manual TrainerRank GPU workflow.
 from dataclasses import dataclass
 
 import torch