From 16317a06e8cda4f8924eb70353903c290df1c4e8 Mon Sep 17 00:00:00 2001 From: Jon Duea Date: Thu, 12 Nov 2026 17:09:03 -0600 Subject: [PATCH 1/2] ci: add automated PR reviews (Claude, CodeRabbit, Gemini) Claude review runs on the owner's Claude subscription via CLAUDE_CODE_OAUTH_TOKEN. CodeRabbit and Gemini Code Assist configs take effect once their GitHub Apps are installed. Co-Authored-By: Claude Opus 5.5 --- .coderabbit.yaml | 14 ++++++++ .gemini/config.yaml | 10 ++++++ .github/workflows/claude-review.yml | 52 +++++++++++++++++++++++++++++ .github/workflows/tests.yml | 28 ++++++++++++++++ 4 files changed, 104 insertions(+) create mode 100644 .coderabbit.yaml create mode 100644 .gemini/config.yaml create mode 100644 .github/workflows/claude-review.yml create mode 100644 .github/workflows/tests.yml diff --git a/.coderabbit.yaml b/.coderabbit.yaml new file mode 100644 index 00000000..de793fb5 --- /dev/null +++ b/.coderabbit.yaml @@ -0,0 +1,14 @@ +# CodeRabbit config. Install the app: https://github.com/apps/coderabbitai +language: en-US +reviews: + profile: chill + request_changes_workflow: false + high_level_summary: true + poem: false + review_status: true + collapse_walkthrough: true + auto_review: + enabled: true + drafts: false +chat: + auto_reply: true diff --git a/.gemini/config.yaml b/.gemini/config.yaml new file mode 100644 index 00000000..f26750ac --- /dev/null +++ b/.gemini/config.yaml @@ -0,0 +1,10 @@ +# Gemini Code Assist config. Install the app: https://github.com/apps/gemini-code-assist +have_fun: false +code_review: + disable: false + comment_severity_threshold: MEDIUM + max_review_comments: -1 + pull_request_opened: + help: false + summary: true + code_review: true diff --git a/.github/workflows/claude-review.yml b/.github/workflows/claude-review.yml new file mode 100644 index 00000000..de1a7de2 --- /dev/null +++ b/.github/workflows/claude-review.yml @@ -0,0 +1,52 @@ +name: Claude PR Review + +# Runs on the repo owner's Claude subscription (CLAUDE_CODE_OAUTH_TOKEN from +# `claude setup-token`), not paid API credits. Triggers are kept narrow to save +# plan quota: new / ready / reopened PRs, plus the `claude-review` label to re-run. +on: + pull_request: + types: [opened, ready_for_review, reopened, labeled] + +concurrency: + group: claude-review-${{ github.event.pull_request.number }} + cancel-in-progress: true + +jobs: + review: + # Skip drafts, PRs from forks (no secret access), and unrelated label events. + if: >- + github.event.pull_request.draft == false && + github.event.pull_request.head.repo.full_name == github.repository && + (github.event.action != 'labeled' || github.event.label.name == 'claude-review') + runs-on: ubuntu-latest + timeout-minutes: 20 + permissions: + contents: read + pull-requests: write + issues: write + id-token: write + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 1 + + - uses: anthropics/claude-code-action@v1 + with: + claude_code_oauth_token: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }} + github_token: ${{ secrets.GITHUB_TOKEN }} + prompt: | + REPO: ${{ github.repository }} + PR NUMBER: ${{ github.event.pull_request.number }} + + Review this pull request. Read AGENTS.md / CONTRIBUTING.md first if present + and hold the change to the repo's own conventions. + + Focus on, in order: correctness bugs, security problems, broken invariants + or contracts, missing tests for changed behavior, and performance regressions. + Skip formatting and style nitpicks that linters already catch. + + Leave an inline comment for each concrete issue on the exact line, saying what + breaks and how to fix it. Then post one short summary comment with + `gh pr comment`. If you find nothing worth flagging, say so in one line. + claude_args: | + --allowedTools "mcp__github_inline_comment__create_inline_comment,Bash(gh pr comment:*),Bash(gh pr diff:*),Bash(gh pr view:*),Read,Glob,Grep" diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 00000000..77c3e04e --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,28 @@ +name: Tests + +on: + pull_request: + push: + branches: [main] + +concurrency: + group: tests-${{ github.ref }} + cancel-in-progress: true + +jobs: + tests: + name: tests + runs-on: ubuntu-latest + timeout-minutes: 30 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - uses: astral-sh/setup-uv@v5 + - name: Install dependencies + run: uv sync --group dev + - name: Fast test suite + run: make test-fast + env: + JAX_PLATFORMS: cpu From f628885c618287e5d76d82228ee643b09e3dd5eb Mon Sep 17 00:00:00 2001 From: Jon Duea Date: Thu, 12 Nov 2026 18:09:35 -0600 Subject: [PATCH 2/2] fix: restore Kaggle worker and repair failing fast tests on main - Restore scripts/kaggle_worker_entry.py and scripts/kaggle_runtime_env.py. 6d9fe7f removed them as obsolete, but the Kaggle/Colab launchers (kaggle_runner, kernel_package, remote_package, colab_runner) still package and run them, and benchmark_jax_rl imports kaggle_runtime_env. - rollout-phase-profile: SystemExit() takes no `file=` kwarg, so a missing integration worktree crashed with TypeError instead of a clean message. - Kaggle submission packager: add factored_decode_scan, map_pool/load and map_pool/bake to RUNTIME_FILES; packaged modules import them, so the submission would fail with ModuleNotFoundError. - Tests: src.training.curriculum moved to src.opponents.curriculum; bounds_exit_time test used a 1-unit board (BOARD_SIZE is 100); preflight sweep reseed_every_updates was deliberately raised to 100. - Docs: drop links to files removed at retirement and to uncommitted outputs/, and reword AGENTS.md so it no longer names removed doc dirs. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 2 +- docs/README.md | 2 +- docs/ROADMAP.md | 2 +- .../colab-long-run-monitor-sync-recovery.md | 2 +- scripts/kaggle_runtime_env.py | 147 +++ scripts/kaggle_worker_entry.py | 1007 +++++++++++++++++ scripts/validate_kaggle_docker_submission.py | 3 + src/cli/benchmark/rollout_phase_profile.py | 4 +- tests/test_cli_rollout_phase_profile.py | 12 + tests/test_preflight_wandb_sweep_compose.py | 2 +- tests/test_train_checkpoint.py | 2 +- tests/test_trajectory_core.py | 8 +- 12 files changed, 1182 insertions(+), 11 deletions(-) create mode 100644 scripts/kaggle_runtime_env.py create mode 100644 scripts/kaggle_worker_entry.py diff --git a/AGENTS.md b/AGENTS.md index 3b3aa867..1d01d12b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -89,7 +89,7 @@ uv run ow train ... artifacts=hybrid_promotion # strict promotion: docker + to - Canonical feature path: Kaggle/JAX obs → `encode_turn` (planet-edge `TurnBatch`); golden tests live in `tests/test_feature_encoding_golden.py`. - JAX layout: rollout `src/jax/rollout/collect.py`, PPO `src/jax/ppo_update.py`, learner sampling `src/jax/action_sampling.py`, shield `src/jax/shield/*` (Python refs in `src/game/shield.py`, `shield_config.py`), train loop `src/jax/train/`, opponents `src/opponents/jax_actions/sampling.py` (mixed-family dispatch, family-batched reorder/merge, historical pool gather with `pool_row_indices` after batched path — use `jax.lax.cond` for trace-safe branches); baseline scripted (nearest_sniper/turtle/opportunistic) and random JAX families skip trajectory shield (`_scripted_edge_action`/`_random_edge_action`); learner/latest/historical self-play stay shielded; Python tournament baselines in `src/opponents/runtime.py` use `BASELINE_RUNTIME_ENV` (shield off) — Kaggle `random_agent` must not pass through `filter_moves_with_trajectory_shield` (collapses to noop); `OPPONENT_FAMILY_*` in `src/opponents/constants.py`; import `src/opponents.pool` only for pool logic. JAX training `nearest_sniper` is not proven action-for-action vs tournament `SniperOpponent`; `latest` uses `encode_turn` without learner `feature_history` — expect learner-skewed win rates, not sniper-beating proof. Learner K-scan now uses O(K) incremental decode (`factorized_decode_step` + edge-gather carry advance) via `src/jax/factored_decode_scan.py`, wired into both rollout sampling (`action_sampling.py`) and PPO replay (`factored_sequence_scan.py`); decoder hidden is exported from the final scan carry `.state` (no trailing full `factorized_decode`). Stacked latest self-play sampler gated by static config (`_static_latest_only_self_play_sample_enabled` in collect/collect_kernel) — never traced `stage_view` on jitted collect (`float(tracer)` on noop/preflight); off for `curriculum=noop_only`, on for `curriculum=latest_only`, off for mixed/historical. Per-format timing (`*_2p`/`*_4p`) only when `metric_groups.debug`; `average_placement_4p` stays default. Profile opponent collect offline via `ow benchmark rollout-phase-profile` / `rollout-phase-breakdown` — default admission preset is `noop_only`; production_mix needs `--train-overrides curriculum=production_mix` (+ `--full-geometry` for 32×256); never `telemetry=rollout_phase_timing` on production `ow train` (host sync stalls). **`ow debug`:** Tier C driver in `src/jax/debug/driver.py`; VS Code launch uses `.venv/bin/ow` (not `python -m src.cli`) plus `ORBIT_WARS_ALLOW_CPU_JAX_ON_NVIDIA=1` for `--cpu` on NVIDIA. -- Hydra dataclass defaults in `src/config/schema.py` can differ from `conf/` YAML; verify with `uv run ow train print_resolved_config=true`. Doc-type map: `docs/README.md` (`make agent-context` → `docs.readme`); calibration index `docs/benchmarks/README.md`; operator tools `docs/tools/`. `docs/archive/` and `docs/audits/` have been removed — canonical specs live under `docs/brainstorms/` + root docs; `tests/test_docs_navigation.py` enforces no `docs/archive|plans|ideation` references and validates markdown + backtick `docs/*.md` / `scripts/*.py` links across docs, `AGENTS.md`, and `CONCEPTS.md` (do not reintroduce archive-path links). +- Hydra dataclass defaults in `src/config/schema.py` can differ from `conf/` YAML; verify with `uv run ow train print_resolved_config=true`. Doc-type map: `docs/README.md` (`make agent-context` → `docs.readme`); calibration index `docs/benchmarks/README.md`; operator tools `docs/tools/`. The former archive and audits doc folders have been removed — canonical specs live under `docs/brainstorms/` + root docs; `tests/test_docs_navigation.py` enforces no archive/plans/ideation doc references and validates markdown + backtick doc / `scripts/*.py` links across docs, `AGENTS.md`, and `CONCEPTS.md` (do not reintroduce archive-path links). - **JAX/Kaggle env parity:** Kaggle reference `kaggle_environments/.../orbit_wars/orbit_wars.py`; JAX `src/jax/env.py` (single JIT-safe path — **no** `pure_callback`, **no** `_reference_*` helpers, **no** `src/game/*` imports in the env hot path; **no** parallel Python/callback fallbacks in train/rollout/env hot path). Parity via mechanical validity invariants, not bit-exact seed replay; see `docs/solutions/conventions/jax-no-kaggle-callbacks.md`. Verify with `make test-kaggle-parity` (includes `tests/test_map_pool_reset.py` when map pool is landed). Live tournament/eval/submit paths use `jax_game_from_observation_fast` obs replay — **not** map-pool indices. Planet rotation index matches Kaggle `obs.step`. Rollout opponents call `encode_turn(game, cfg.task)` without `feature_history` (learner history is updated in env `_finish_step`); Phase B `opp_batch_cache` must not reuse learner `feature_history`; opponent batch-cache skip must use `jax.lax.cond`, not Python `if` on traced values. **Map pool:** offline bake via `ow benchmark map-pool {profile,bake,validate}` → `data/jax_map_pool/default_v1.npz`; training `batched_reset_with_pool` gather only — no inline `while_loop` rejection sampling in vmapped `reset()`; comets activate from pre-baked schedules in `step()`. Map pool is the **default** training path — `conf/task/base.yaml` sets `map_pool_path: data/jax_map_pool/default_v1.npz`, inherited by every task profile (`shield_cheap`/`default`/`map_pool`); rollout init uses gather reset whenever `map_pool_path` is set, and only configs **without** a path fall to `batched_reset` → `_reset_legacy` stub board (tests only). `task=map_pool` (`conf/task/map_pool.yaml`) is just an alias to the same default pool — not "turn on map pool"; the real knob for a *different* pool is `task.map_pool_path=` (+ `map_pool_sha256`), and throughput admission uses `task=map_pool` only as a benchmarking preset label. Failed June Colab pilots collapsed on `production_mix`/early self-play, not on map pool. - **Hybrid promotion / submit-valid:** Decision tree in `docs/AGENT_CAPABILITIES.md` (hybrid poll → `ow eval results show` for `validation_ok`, or `ow eval package --validate-docker` with JSON `"ok": true`); local replay HTML and packaging-only are never submit-valid proof. `artifacts=hybrid_promotion` sets `promotion.strategy=hybrid`, `tournament.enabled=true`, `checkpoint_eval_async=true`, `replay.enabled=false`. Scalar metric improvements queue composite `checkpoint_eval` (Docker validation → tournament → promote). Artifacts: `evaluations/checkpoint_eval_u_/{docker_validation,tournament}/`. Poll `ow eval status --run --watch`; dry-run `ow eval jobs cancel` before cancelling queue jobs. Bare `ow train` uses `artifacts=default`—not the submit-valid funnel. - **Training geometry & run layout:** `rollout_microbatch_envs` must be ≤ each rollout group's env count (`src/config/rollout_allocation.py`); canonical `training=default` inherits `2p4p_32_split` (16+16 envs, microbatch 16). Mixed 2p/4p uses Option A: separate jitted collectors per format (2p vs 4p differ in env rules/home assignment, not just seat count), sequential Python-loop collect then concat before PPO — `2p4p_*_rotate` collects one format per update. Compose tests parametrize every `conf/training/*` profile. `ow train` with `output.campaign=` writes under `outputs/campaigns//runs//` — metrics in `logs/*_jax.jsonl` (not run-root `metrics.jsonl`), checkpoints in `checkpoints/jax_ckpt_*.pkl`, artifact jobs in `queue/optional_jobs/`, evaluations in `evaluations/`. diff --git a/docs/README.md b/docs/README.md index 4fbc16da..9e03ed85 100644 --- a/docs/README.md +++ b/docs/README.md @@ -64,7 +64,7 @@ Long-lived references at `docs/` root (not folder indexes): - [agent-native-phase3-status.md](agent-native-phase3-status.md) — shipped benchmark/gate/sweep CLI primitives - [nomenclature-rfc.md](nomenclature-rfc.md) — user-facing term mappings -Repo root: [COLAB_LAUNCH_AND_INTEGRATION_PROMOTION.md](../COLAB_LAUNCH_AND_INTEGRATION_PROMOTION.md) — live Colab pilot tracker; [CONCEPTS.md](../CONCEPTS.md) — glossary. +Repo root: [CONCEPTS.md](../CONCEPTS.md) — glossary. (The Colab pilot tracker was retired with the project; see [RETROSPECTIVE.md](../RETROSPECTIVE.md).) ## Maintenance diff --git a/docs/ROADMAP.md b/docs/ROADMAP.md index 23d81b4c..d0a08ead 100644 --- a/docs/ROADMAP.md +++ b/docs/ROADMAP.md @@ -27,7 +27,7 @@ | Item | Link | | ---- | ---- | | Launch hygiene learner ablation + tier-2 assessment | Tier-2 failed (~4× sec/update); hot path exhausted; B (hygiene) wins learn-proof vs A (79162a); artifact [ablation.json](benchmarks/launch-hygiene-ablation.json) | [plan](solutions/tooling-decisions/launch-hygiene-learner-ablation-gate.md) · [runbook](operator-runbook.md) | -| Preflight learn-proof refresh (post-hygiene) | VERIFIED through `beat_random` vs `preflight-calibration.json` | [report](../outputs/preflight/factorized_post_hygiene_learn_proof.json) · [runbook](operator-runbook.md) | +| Preflight learn-proof refresh (post-hygiene) | VERIFIED through `beat_random` vs `preflight-calibration.json` | report: `outputs/preflight/factorized_post_hygiene_learn_proof.json` (generated locally, not committed) · [runbook](operator-runbook.md) | | Cursor session-start hook: `docs/CURSOR.md` + project `.cursor/hooks.json` example | [plan](CURSOR.md) · [Phase 3 §4](agent-native-phase3-status.md) | | Observability debug bundle: `mean_ships_per_launch`, PPO `_2p`/`_4p` + update-time fractions gated behind `metric_groups.debug` | [plan](solutions/developer-experience/benchmark-subprocess-training-observability.md) | | Split decoder replay batch contracts (#167) | [plan](architecture/jax-policy-encoder.md) · Closes #167 | diff --git a/docs/solutions/workflow-issues/colab-long-run-monitor-sync-recovery.md b/docs/solutions/workflow-issues/colab-long-run-monitor-sync-recovery.md index d31a0c39..1095f939 100644 --- a/docs/solutions/workflow-issues/colab-long-run-monitor-sync-recovery.md +++ b/docs/solutions/workflow-issues/colab-long-run-monitor-sync-recovery.md @@ -283,7 +283,7 @@ Eval artifacts: `outputs/colab_runner/monitor/evals//jax_ckpt_/`. Mon ## Related - Operator reference: [`docs/colab_runner.md`](../../colab_runner.md) — monitor section -- Launch tracker: [`COLAB_LAUNCH_AND_INTEGRATION_PROMOTION.md`](../../../COLAB_LAUNCH_AND_INTEGRATION_PROMOTION.md) — pilot gates, sweep invalidation, June 16 long-run state +- Launch tracker: `COLAB_LAUNCH_AND_INTEGRATION_PROMOTION.md` (removed at project retirement) — pilot gates, sweep invalidation, June 16 long-run state - Train host architecture: [`colab-train-host-preflight-long-run.md`](../architecture-patterns/colab-train-host-preflight-long-run.md) - CLI observability (no tail pipe): [`ow-long-cli-stderr-progress-no-tail-pipe.md`](../developer-experience/ow-long-cli-stderr-progress-no-tail-pipe.md) — prefer `ow train colab monitor` over piping `launch` - Sweep selection anti-patterns (generalized): [`planet-flow-sweep-gameable-objective.md`](../logic-errors/planet-flow-sweep-gameable-objective.md) diff --git a/scripts/kaggle_runtime_env.py b/scripts/kaggle_runtime_env.py new file mode 100644 index 00000000..4dd58a92 --- /dev/null +++ b/scripts/kaggle_runtime_env.py @@ -0,0 +1,147 @@ +"""Kaggle/runtime environment helpers for Orbit Wars workers. + +This module must be safe to import before JAX. The critical rule for NVIDIA +Kaggle workers is: never set ``JAX_PLATFORMS=gpu``. Newer JAX treats the +``gpu`` alias as a request to consider GPU backends broadly and may attempt +ROCm before/alongside CUDA. Kaggle NVIDIA kernels should be pinned to +``cuda,cpu``. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +KAGGLE_RUNTIME_ENV_PATCH_VERSION = "cuda-platform-pin-v10" +_WORKER_VENV_ENV = "ORBIT_WARS_WORKER_VENV" +_DRIVER_DIR_CANDIDATES: tuple[str, ...] = ( + "/usr/local/nvidia/lib64", + "/usr/local/nvidia/lib", + "/usr/local/cuda/compat", + "/usr/local/cuda-12.8/compat", + "/usr/local/cuda/lib64", + "/usr/lib/x86_64-linux-gnu", +) + + +def pin_jax_platform_from_kaggle(*, env: dict[str, str] | None = None) -> None: + """Pin JAX's platform selection for Kaggle before importing JAX. + + ``JAX_PLATFORMS=gpu`` is intentionally avoided for NVIDIA workers because it + can trigger ROCm initialization attempts. Use ``cuda,cpu`` instead. + """ + + target = os.environ if env is None else env + accelerator_id = target.get("KAGGLE_ACCELERATOR_ID", os.environ.get("KAGGLE_ACCELERATOR_ID", "")) + backend = target.get("ORBIT_WARS_JAX_BACKEND", os.environ.get("ORBIT_WARS_JAX_BACKEND", "")).strip().lower() + target.pop("JAX_PLATFORM_NAME", None) + + if backend == "tpu" or _accelerator_requests_tpu(accelerator_id): + target["JAX_PLATFORMS"] = "tpu,cpu" + elif backend == "cpu" or _truthy(target.get("ORBIT_WARS_FORCE_JAX_CPU", os.environ.get("ORBIT_WARS_FORCE_JAX_CPU", ""))): + target["JAX_PLATFORMS"] = "cpu" + elif _accelerator_requests_nvidia(accelerator_id) or _nvidia_runtime_visible(): + target["JAX_PLATFORMS"] = "cuda,cpu" + else: + # Avoid GPU plugin probing on plain CPU hosts. Preserve an explicit user + # setting if one was already provided. + target.setdefault("JAX_PLATFORMS", "cpu") + + +def isolate_worker_python_env(*, env: dict[str, str] | None = None) -> None: + """Make subprocesses prefer the managed worker venv and avoid user site dirs.""" + + target = os.environ if env is None else env + target["PYTHONNOUSERSITE"] = "1" + target.pop("PYTHONHOME", None) + venv = _worker_venv(target) + if venv is None: + return + target["VIRTUAL_ENV"] = str(venv) + target["UV_PROJECT_ENVIRONMENT"] = str(venv) + bin_dir = venv / "bin" + if bin_dir.exists(): + target["PATH"] = _prepend_path(str(bin_dir), target.get("PATH", "")) + + +def add_worker_cuda_library_path(*, env: dict[str, str] | None = None) -> None: + """Expose venv CUDA wheel libraries plus Kaggle's NVIDIA driver libraries.""" + + target = os.environ if env is None else env + accelerator_id = target.get("KAGGLE_ACCELERATOR_ID", os.environ.get("KAGGLE_ACCELERATOR_ID", "")) + if not (_accelerator_requests_nvidia(accelerator_id) or _nvidia_runtime_visible()): + return + venv = _worker_venv(target) + dirs: list[str] = [] + if venv is not None: + dirs.extend(_cuda_wheel_library_dirs(venv)) + dirs.extend(_cuda_driver_library_dirs(target.get("LD_LIBRARY_PATH", ""))) + if dirs: + target["LD_LIBRARY_PATH"] = os.pathsep.join(_dedupe(dirs)) + target.setdefault("TF_GPU_ALLOCATOR", "cuda_malloc_async") + + +def _worker_venv(env: dict[str, str]) -> Path | None: + value = env.get(_WORKER_VENV_ENV, os.environ.get(_WORKER_VENV_ENV, "")).strip() + if not value: + return None + return Path(value).resolve() + + +def _cuda_wheel_library_dirs(venv: Path) -> list[str]: + dirs: list[str] = [] + for site_packages in sorted(venv.glob("lib/python*/site-packages")): + nvidia_root = site_packages / "nvidia" + if not nvidia_root.exists(): + continue + for lib_dir in sorted(nvidia_root.glob("*/lib")): + if lib_dir.is_dir(): + dirs.append(str(lib_dir.resolve())) + return dirs + + +def _cuda_driver_library_dirs(existing_ld_library_path: str) -> list[str]: + candidates = [ + *(item for item in existing_ld_library_path.split(os.pathsep) if item), + *_DRIVER_DIR_CANDIDATES, + ] + result: list[str] = [] + for item in candidates: + path = Path(item) + if not path.is_dir(): + continue + if any(path.glob("libcuda.so*")) or any(path.glob("libnvidia-ml.so*")): + result.append(str(path.resolve())) + return _dedupe(result) + + +def _nvidia_runtime_visible() -> bool: + return Path("/usr/local/nvidia/lib64").is_dir() or any(Path("/dev").glob("nvidia[0-9]*")) + + +def _accelerator_requests_nvidia(accelerator_id: str) -> bool: + return accelerator_id.strip().lower().startswith("nvidia") + + +def _accelerator_requests_tpu(accelerator_id: str) -> bool: + return accelerator_id.strip().lower().startswith("tpu") + + +def _truthy(value: str) -> bool: + return str(value).strip().lower() in {"1", "true", "yes", "on"} + + +def _prepend_path(item: str, existing: str) -> str: + parts = [part for part in existing.split(os.pathsep) if part] + return os.pathsep.join([item, *(part for part in parts if part != item)]) + + +def _dedupe(items: list[str]) -> list[str]: + seen: set[str] = set() + result: list[str] = [] + for item in items: + if item in seen: + continue + seen.add(item) + result.append(item) + return result diff --git a/scripts/kaggle_worker_entry.py b/scripts/kaggle_worker_entry.py new file mode 100644 index 00000000..5a45cd5e --- /dev/null +++ b/scripts/kaggle_worker_entry.py @@ -0,0 +1,1007 @@ +#!/usr/bin/env python3 +"""Kaggle-side W&B population worker entrypoint.""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import sys +import time +from pathlib import Path +from typing import Any + +SCRIPT_DIR = Path(__file__).resolve().parent +for candidate in (SCRIPT_DIR, SCRIPT_DIR.parent): + candidate_text = str(candidate) + if candidate_text not in sys.path: + sys.path.insert(0, candidate_text) + +from src.orchestration.kaggle_jax import ( # noqa: E402 + format_bootstrap_failure, + log_bootstrap_failure, + sync_kaggle_worker_environment, +) +from src.orchestration.remote_worker import ( # noqa: E402 + collect_checkpoint_paths, + completed_summary, + exception_message, + exit_code, + load_packaged_env, + tail, + write_summary, +) + +# U1b: full Kaggle worker orchestration still lives here; only shared helpers are +# extracted to src/orchestration/remote_worker.py for Colab reuse. + +_JAX_BACKEND_MARKER = Path(".orbit_wars_jax_backend") +_WORKER_VENV_ENV = "ORBIT_WARS_WORKER_VENV" +_WORKER_VENV_READY_ENV = "ORBIT_WARS_WORKER_VENV_READY" +KAGGLE_WORKER_ENTRY_PATCH_VERSION = "subprocess-jax-diagnostics-v9" + +_CUDA_DRIVER_LIBRARY_DIR_CANDIDATES: tuple[str, ...] = ( + "/usr/local/nvidia/lib64", + "/usr/local/nvidia/lib", + "/usr/local/cuda/compat", + "/usr/local/cuda-12.8/compat", + "/usr/lib/x86_64-linux-gnu", +) + +_HYDRA_CONFIG_GROUP_KEYS: frozenset[str] = frozenset( + { + "model", + "training", + "opponents", + "curriculum", + "reward", + "task", + "telemetry", + "artifacts", + "output", + } +) + + +def main() -> None: + print(f"ORBIT_WARS_WORKER_ENTRY_PATCH={KAGGLE_WORKER_ENTRY_PATCH_VERSION}", flush=True) + load_packaged_env() + standalone = _is_standalone_mode() + wandb_secret: dict[str, object] | None + if standalone: + wandb_secret = {"skipped": True, "reason": "standalone worker mode"} + else: + wandb_secret = _load_wandb_api_key_from_kaggle_secret() + summary: dict[str, Any] = { + "status": "starting", + "cwd": str(Path.cwd()), + "env": _safe_worker_env(), + "worker_mode": "standalone" if standalone else "wandb", + "wandb_secret": wandb_secret, + } + write_summary(summary) + try: + _ensure_uv_environment(summary) + summary["diagnostics"] = diagnostics() + write_summary(summary) + _ensure_accelerator(summary) + if standalone: + _run_standalone_worker(summary) + else: + _run_wandb_worker(summary) + except BaseException as exc: + summary["status"] = "failed" + summary["error"] = exception_message(exc) + summary.setdefault("exit_code", exit_code(exc)) + write_summary(summary) + raise + + +def _run_standalone_worker(summary: dict[str, Any]) -> None: + config = _load_standalone_config() + summary["standalone_config"] = config + summary["standalone_overrides"] = list(_config_to_overrides(config)) + write_summary(summary) + _run_training_candidate(summary, config, wandb_run=None) + if summary.get("status") == "failed": + raise SystemExit(int(summary.get("exit_code", 1))) + summary.setdefault("exit_code", 0) + summary["status"] = "standalone_complete" + write_summary(summary) + + +def _run_wandb_worker(summary: dict[str, Any]) -> None: + sweep_id = os.environ.get("WANDB_SWEEP_ID") + if not sweep_id: + raise SystemExit( + "WANDB_SWEEP_ID is required for Kaggle population workers." + ) + if not os.environ.get("WANDB_API_KEY"): + summary["wandb_secret_retry"] = _load_wandb_api_key_from_kaggle_secret() + write_summary(summary) + + if not os.environ.get("WANDB_API_KEY"): + raise SystemExit( + "WANDB_API_KEY is not available after retrying Kaggle Secrets. " + "Add or attach a Kaggle Secret named by WANDB_API_KEY_SECRET_NAME, " + "or provide WANDB_API_KEY through the worker environment. " + f"Secret status: {summary.get('wandb_secret_retry') or summary.get('wandb_secret')}" + ) + wandb_project = os.environ.get("WANDB_PROJECT") + wandb_entity = os.environ.get("WANDB_ENTITY") + summary["sweep_id"] = sweep_id + summary["wandb_agent"] = { + "project": wandb_project or "", + "entity": wandb_entity or "", + } + write_summary(summary) + + import wandb # type: ignore + + def run_candidate() -> None: + run = wandb.init(job_type="kaggle-population-worker") + run_finished = False + summary["wandb_run"] = { + "id": str(getattr(run, "id", "")), + "name": str(getattr(run, "name", "")), + } + write_summary(summary) + try: + _run_training_candidate(summary, dict(wandb.config), wandb_run=run) + run.summary["kaggle_worker_exit_code"] = summary.get("exit_code", 0) + run.finish(exit_code=int(summary.get("exit_code", 0))) + run_finished = True + if int(summary.get("exit_code", 0)) != 0: + raise SystemExit(int(summary.get("exit_code", 1))) + except BaseException as exc: + summary["status"] = "failed" + summary["error"] = exception_message(exc) + summary.setdefault("exit_code", exit_code(exc)) + write_summary(summary) + if not run_finished: + run.finish(exit_code=exit_code(exc)) + raise + + wandb.agent( + sweep_id, + function=run_candidate, + count=1, + entity=wandb_entity, + project=wandb_project, + ) + if summary.get("status") == "failed": + raise SystemExit(int(summary.get("exit_code", 1))) + summary.setdefault("exit_code", 0) + summary["status"] = "agent_complete" + write_summary(summary) + + +def _run_training_candidate( + summary: dict[str, Any], + config: dict[str, Any], + *, + wandb_run: Any | None, +) -> None: + from src.orchestration.throughput import ( + calibration_grid, + estimate_training_overrides, + finalize_rollout_shape_overrides, + ) + + standalone = wandb_run is None + try: + observed = _hardware_profile(summary) + base_overrides = _config_to_overrides(config) + if standalone: + base_overrides.extend(_standalone_extra_overrides()) + hydra_overrides = tuple(base_overrides) + calibration_overrides = estimate_training_overrides( + observed, + _prefixed_config(config, "model."), + _prefixed_config(config, "task."), + hydra_overrides=hydra_overrides, + ) + variants = calibration_grid( + calibration_overrides, + hydra_overrides=hydra_overrides, + ) + settings = _calibration_settings() + calibration_results = _run_calibration(base_overrides, variants, settings) + selected_overrides = finalize_rollout_shape_overrides( + _select_calibration( + calibration_results, + calibration_overrides, + allow_fallback=os.environ.get( + "ORBIT_WARS_KAGGLE_ALLOW_CALIBRATION_FALLBACK" + ) + == "1", + ), + hydra_overrides, + ) + summary["hardware"] = { + "accelerator_id": observed.accelerator_id, + "gpu_name": observed.gpu_name, + "memory_gb": observed.memory_gb, + } + summary["calibration"] = { + "settings": settings, + "estimated_overrides": list(calibration_overrides), + "results": calibration_results, + "selected_overrides": list(selected_overrides), + } + write_summary(summary) + if wandb_run is not None: + import wandb # type: ignore + + wandb.log( + { + "kaggle_worker/observed_gpu": observed.gpu_name, + "kaggle_worker/observed_memory_gb": observed.memory_gb, + "kaggle_worker/calibration_results": calibration_results, + "kaggle_worker/selected_overrides": list(selected_overrides), + }, + step=0, + ) + overrides = list(base_overrides) + overrides.extend(selected_overrides) + overrides = _apply_smoke_training_overrides(overrides) + if standalone: + overrides.extend( + [ + "telemetry.wandb.enabled=false", + "telemetry.wandb.log_artifacts=false", + "artifacts.artifact_pipeline.docker_validation_async=false", + "artifacts.replay.enabled=false", + ] + ) + else: + overrides.extend( + [ + "telemetry.wandb.enabled=true", + "telemetry.wandb.log_artifacts=true", + "artifacts.artifact_pipeline.docker_validation_async=false", + "artifacts.replay.enabled=false", + ] + ) + command = _render_worker_train_command(tuple(overrides)) + summary["final_command"] = command + summary["selected_overrides"] = overrides + write_summary(summary) + print("worker command:", " ".join(command), flush=True) + env = _subprocess_env() + if wandb_run is not None: + env.setdefault("WANDB_RUN_ID", wandb_run.id) + env.setdefault("WANDB_RESUME", "allow") + completed = subprocess.run(command, check=False, text=True, env=env) + summary["exit_code"] = completed.returncode + summary["checkpoint_paths"] = collect_checkpoint_paths() + summary["status"] = "completed" if completed.returncode == 0 else "failed" + write_summary(summary) + if completed.returncode != 0: + raise SystemExit(completed.returncode) + except BaseException as exc: + summary["status"] = "failed" + summary["error"] = exception_message(exc) + summary.setdefault("exit_code", exit_code(exc)) + summary.setdefault("checkpoint_paths", collect_checkpoint_paths()) + write_summary(summary) + raise + + + +def diagnostics() -> dict[str, Any]: + """Collect runtime diagnostics without importing JAX in this parent process. + + The verified JAX CUDA configuration only works reliably when LD_LIBRARY_PATH + is present at process start. Importing JAX after mutating LD_LIBRARY_PATH in + this already-running parent process can fail even when child processes work. + Therefore all JAX/Flax checks run in a fresh worker-venv Python subprocess. + """ + + _activate_local_venv() + payload: dict[str, Any] = { + "python": sys.version, + "executable": sys.executable, + "sys_prefix": sys.prefix, + "sys_base_prefix": getattr(sys, "base_prefix", ""), + "worker_entry_patch": KAGGLE_WORKER_ENTRY_PATCH_VERSION, + "diagnostics_mode": "jax-subprocess-probe-v9", + "accelerator_id": os.environ.get("KAGGLE_ACCELERATOR_ID", ""), + "jax_platforms_env": os.environ.get("JAX_PLATFORMS", ""), + "virtual_env": os.environ.get("VIRTUAL_ENV", ""), + "worker_venv": os.environ.get(_WORKER_VENV_ENV, ""), + "ld_library_path": os.environ.get("LD_LIBRARY_PATH", ""), + "nvidia_smi": _run_optional( + ["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader"] + ), + } + probe = _run_jax_diagnostics_subprocess() + payload["jax_subprocess_probe"] = probe + if int(probe.get("returncode", 1)) == 0: + parsed = probe.get("parsed") + if isinstance(parsed, dict): + payload["jax_version"] = parsed.get("jax_version", "") + payload["jax_default_backend"] = parsed.get("jax_default_backend", "") + payload["jax_devices"] = list(parsed.get("jax_devices", [])) + payload["jax_cuda_devices"] = list(parsed.get("jax_cuda_devices", [])) + payload["jax_platforms"] = list(parsed.get("jax_platforms", [])) + payload["flax_version"] = parsed.get("flax_version", "") + payload["flax_linen_import"] = parsed.get("flax_linen_import", "") + else: + payload["jax_error"] = str( + probe.get("stderr_tail") or probe.get("stdout_tail") or probe + ) + return payload + + +def _run_jax_diagnostics_subprocess() -> dict[str, object]: + command = [str(_venv_python()), "-c", _JAX_DIAGNOSTICS_PROBE_CODE] + completed = subprocess.run( + command, + check=False, + capture_output=True, + text=True, + env=_subprocess_env(), + ) + result: dict[str, object] = { + "command": command, + "returncode": int(completed.returncode), + "stdout_tail": tail(completed.stdout, limit=6000), + "stderr_tail": tail(completed.stderr, limit=6000), + } + if completed.stdout: + for line in reversed(completed.stdout.strip().splitlines()): + try: + result["parsed"] = json.loads(line) + break + except json.JSONDecodeError: + continue + return result + + +_JAX_DIAGNOSTICS_PROBE_CODE = r""" +import json +import os +import sys + + +def platform_aliases(devices): + platforms = set() + for device in devices: + platform = str(getattr(device, "platform", "") or "") + if platform: + platforms.add(platform) + text = str(device).lower() + if "cuda" in text: + platforms.add("cuda") + platforms.add("gpu") + if platform == "gpu": + platforms.add("cuda") + if platform == "cuda": + platforms.add("gpu") + return sorted(platforms) + +payload = { + "probe_python": sys.executable, + "sys_prefix": sys.prefix, + "sys_base_prefix": getattr(sys, "base_prefix", ""), + "JAX_PLATFORMS": os.environ.get("JAX_PLATFORMS", ""), + "VIRTUAL_ENV": os.environ.get("VIRTUAL_ENV", ""), + "LD_LIBRARY_PATH": os.environ.get("LD_LIBRARY_PATH", ""), +} +try: + import flax + import flax.linen as nn # noqa: F401 + payload["flax_version"] = str(getattr(flax, "__version__", "")) + payload["flax_linen_import"] = "ok" + + import jax + devices = list(jax.devices()) + payload["jax_version"] = str(getattr(jax, "__version__", "")) + payload["jax_default_backend"] = str(jax.default_backend()) + payload["jax_devices"] = [str(device) for device in devices] + payload["jax_platforms"] = platform_aliases(devices) + try: + cuda_devices = list(jax.devices("cuda")) + except Exception as exc: + payload["jax_cuda_error"] = repr(exc) + cuda_devices = [] + payload["jax_cuda_devices"] = [str(device) for device in cuda_devices] + payload["jax_platforms"] = sorted( + set(payload["jax_platforms"]) | set(platform_aliases(cuda_devices)) + ) + if cuda_devices: + payload["jax_platforms"] = sorted(set(payload["jax_platforms"]) | {"cuda", "gpu"}) + x = jax.numpy.arange(8.0) + payload["probe_sum"] = float(jax.numpy.sum(x)) + print(json.dumps(payload, sort_keys=True)) +except BaseException as exc: + payload["error"] = repr(exc) + print(json.dumps(payload, sort_keys=True)) + raise +""" + +def _ensure_accelerator(summary: dict[str, Any]) -> None: + if os.environ.get("ORBIT_WARS_KAGGLE_ALLOW_CPU") == "1": + return + from src.orchestration.accelerators import is_tpu_accelerator + + diagnostics_payload = summary.get("diagnostics", {}) + platforms = set(diagnostics_payload.get("jax_platforms", [])) + accelerator_id = os.environ.get("KAGGLE_ACCELERATOR_ID", "") + if is_tpu_accelerator(accelerator_id): + if "tpu" not in platforms: + raise SystemExit( + f"JAX TPU backend required for Kaggle worker. Platforms: {platforms}; " + f"diagnostics={diagnostics_payload}" + ) + return + if _accelerator_requests_nvidia(accelerator_id): + if not ({"gpu", "cuda"} & platforms): + raise SystemExit( + f"JAX GPU/CUDA backend required for Kaggle worker. " + f"Platforms: {platforms}; diagnostics={diagnostics_payload}" + ) + return + + +def _hardware_profile(summary: dict[str, Any]) -> Any: + from src.orchestration.accelerators import default_memory_gb, is_tpu_accelerator + from src.orchestration.throughput import HardwareProfile + + accelerator_id = os.environ.get("KAGGLE_ACCELERATOR_ID", "unknown") + if is_tpu_accelerator(accelerator_id): + return HardwareProfile( + accelerator_id=accelerator_id, + gpu_name="tpu-v5e-8", + memory_gb=default_memory_gb(accelerator_id), + ) + + smi = str(summary.get("diagnostics", {}).get("nvidia_smi", "")) + gpu_name = smi.split(",", 1)[0].strip() if smi.strip() else "unknown" + memory_gb = default_memory_gb(accelerator_id, fallback=16.0) + if "," in smi: + memory_text = smi.split(",", 1)[1].strip().lower().replace("mib", "") + try: + memory_gb = float(memory_text) / 1024.0 + except ValueError: + pass + return HardwareProfile( + accelerator_id=accelerator_id, + gpu_name=gpu_name, + memory_gb=memory_gb, + ) + + +def _config_to_overrides(config: dict[str, Any]) -> list[str]: + overrides: list[str] = [] + for key, value in sorted(config.items()): + if not isinstance(key, str): + continue + if key.startswith("kaggle_worker."): + continue + if "." not in key and key not in _HYDRA_CONFIG_GROUP_KEYS: + continue + overrides.append(f"{key}={_hydra_value(value)}") + return overrides + + +def _prefixed_config(config: dict[str, Any], prefix: str) -> dict[str, Any]: + return { + key.removeprefix(prefix): value + for key, value in config.items() + if isinstance(key, str) and key.startswith(prefix) + } + + +def _hydra_value(value: object) -> str: + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, (list, tuple)): + return "[" + ",".join(str(item) for item in value) + "]" + return str(value) + + +def _calibration_command( + overrides: list[str], *, warmup: int, updates: int +) -> list[str]: + return [ + str(_venv_python()), + str(SCRIPT_DIR / "benchmark_jax_rl.py"), + "--overrides", + *overrides, + "--warmup", + str(warmup), + "--updates", + str(updates), + ] + + +def _render_worker_train_command(overrides: tuple[str, ...]) -> list[str]: + return [str(_venv_python()), "-m", "src.train", *overrides] + + +def _venv_python() -> Path: + candidate = _worker_venv() / "bin" / "python" + if not candidate.exists(): + raise SystemExit("Worker venv python is missing after bootstrap.") + return candidate + + +def _run_calibration( + base_overrides: list[str], + variants: list[tuple[str, ...]], + settings: dict[str, int], +) -> list[dict[str, object]]: + results: list[dict[str, object]] = [] + warmup = settings["warmup"] + updates = settings["updates"] + timeout_seconds = settings["timeout_seconds"] + selected_variants = variants[: settings["max_variants"]] + for index, variant in enumerate(selected_variants): + overrides = [ + *_without_training_shape_overrides(base_overrides), + *variant, + "telemetry.wandb.enabled=false", + "artifacts.artifact_pipeline.enabled=false", + "artifacts.replay.enabled=false", + ] + command = _calibration_command( + overrides, + warmup=warmup, + updates=updates, + ) + print("calibration command:", " ".join(command), flush=True) + result: dict[str, object] = { + "index": index, + "overrides": list(variant), + "timeout_seconds": timeout_seconds, + } + env = _subprocess_env() + try: + completed = subprocess.run( + command, + check=False, + capture_output=True, + text=True, + timeout=timeout_seconds, + env=env, + ) + except subprocess.TimeoutExpired as exc: + result["returncode"] = 124 + result["error"] = f"calibration timed out after {timeout_seconds}s" + result["stdout_tail"] = tail(exc.stdout) + result["stderr_tail"] = tail(exc.stderr) + results.append(result) + continue + result["returncode"] = completed.returncode + result["stdout_tail"] = tail(completed.stdout) + result["stderr_tail"] = tail(completed.stderr) + if completed.returncode == 0: + try: + result.update(json.loads(completed.stdout.strip().splitlines()[-1])) + except (IndexError, json.JSONDecodeError) as exc: + result["error"] = f"failed to parse benchmark output: {exc}" + else: + result["error"] = tail(completed.stderr or completed.stdout) + print( + f"calibration variant {index} failed (rc={completed.returncode}):", + result["error"], + flush=True, + ) + results.append(result) + return results + + +def _select_calibration( + results: list[dict[str, object]], + fallback: tuple[str, ...], + *, + allow_fallback: bool, +) -> tuple[str, ...]: + successful = [ + result + for result in results + if int(result.get("returncode", 1)) == 0 and "samples_per_sec" in result + ] + if not successful: + if not allow_fallback: + raise SystemExit( + "All calibration variants failed; set " + "ORBIT_WARS_KAGGLE_ALLOW_CALIBRATION_FALLBACK=1 to use estimator fallback." + ) + return fallback + best = max(successful, key=lambda item: float(item.get("samples_per_sec", 0.0))) + return tuple(str(item) for item in best.get("overrides", fallback)) + + +def _calibration_settings() -> dict[str, int]: + smoke = _is_smoke_run() + return { + "warmup": _env_int("ORBIT_WARS_KAGGLE_CALIBRATION_WARMUP", 0, minimum=0), + "updates": _env_int("ORBIT_WARS_KAGGLE_CALIBRATION_UPDATES", 1, minimum=1), + "max_variants": _env_int( + "ORBIT_WARS_KAGGLE_CALIBRATION_MAX_VARIANTS", + 1 if smoke else 3, + minimum=1, + ), + "timeout_seconds": _env_int( + "ORBIT_WARS_KAGGLE_CALIBRATION_TIMEOUT_SECONDS", + 600 if smoke else 1800, + minimum=1, + ), + } + + +def _is_smoke_run() -> bool: + return os.environ.get("ORBIT_WARS_KAGGLE_RUN_TYPE", "").strip().lower() == "smoke" + + +def _is_standalone_mode() -> bool: + return ( + os.environ.get("ORBIT_WARS_KAGGLE_WORKER_MODE", "").strip().lower() + == "standalone" + ) + + +def _load_standalone_config() -> dict[str, Any]: + from src.orchestration.wandb_sweeps import load_standalone_config + + sweep_yaml = os.environ.get("WANDB_SWEEP_YAML", "").strip() + if not sweep_yaml: + raise SystemExit( + "WANDB_SWEEP_YAML is required for standalone Kaggle workers." + ) + path = Path(sweep_yaml) + if not path.is_file(): + raise SystemExit(f"Packaged sweep YAML is missing: {path}") + return load_standalone_config(path) + + +def _standalone_extra_overrides() -> list[str]: + raw = os.environ.get("ORBIT_WARS_KAGGLE_STANDALONE_OVERRIDES", "").strip() + if not raw: + return [] + try: + parsed = json.loads(raw) + except json.JSONDecodeError: + parsed = [item.strip() for item in raw.split(",") if item.strip()] + if not isinstance(parsed, list): + raise SystemExit( + "ORBIT_WARS_KAGGLE_STANDALONE_OVERRIDES must be a JSON list of Hydra overrides." + ) + return [str(item) for item in parsed if str(item).strip()] + + +def _apply_smoke_training_overrides(overrides: list[str]) -> list[str]: + """Apply short training defaults for smoke validation runs.""" + + if not _is_smoke_run(): + return overrides + filtered = [ + item + for item in overrides + if not item.startswith("training.total_updates=") + ] + filtered.append("training.total_updates=10") + return filtered + + +def _without_training_shape_overrides(overrides: list[str]) -> list[str]: + shape_keys = { + "training.num_envs", + "training.rollout_steps", + "training.update_chunk_rows", + "training.rollout_microbatch_envs", + } + return [item for item in overrides if item.split("=", 1)[0] not in shape_keys] + + +def _run_optional(command: list[str]) -> str: + if shutil.which(command[0]) is None: + return "" + completed = subprocess.run(command, check=False, capture_output=True, text=True) + return completed.stdout.strip() or completed.stderr.strip() + + +def _load_wandb_api_key_from_kaggle_secret() -> dict[str, object]: + """Resolve the W&B key from Kaggle Secrets instead of packaging it. + + Kaggle's secrets service can be temporarily unavailable immediately after a + CLI-pushed kernel starts. Retry because manual reruns often succeed once the + service has initialized for the kernel. + """ + + if os.environ.get("WANDB_API_KEY"): + return {"loaded": True, "source": "environment"} + + secret_name = os.environ.get("WANDB_API_KEY_SECRET_NAME", "WANDB_API_KEY").strip() + if not secret_name: + return {"loaded": False, "secret_name": "", "error": "no secret name set"} + + max_wait_seconds = _env_int( + "ORBIT_WARS_KAGGLE_SECRET_WAIT_SECONDS", + 30 if _is_smoke_run() else 120, + minimum=0, + ) + sleep_seconds = 2.0 + deadline = time.monotonic() + max_wait_seconds + attempts: list[dict[str, object]] = [] + + while True: + attempt_index = len(attempts) + 1 + try: + from kaggle_secrets import UserSecretsClient # type: ignore + except ImportError as exc: + error = f"kaggle_secrets unavailable: {exc}" + attempts.append({"attempt": attempt_index, "error": error}) + return { + "loaded": False, + "secret_name": secret_name, + "attempts": attempts, + "error": error, + } + + try: + value = UserSecretsClient().get_secret(secret_name) + except Exception as exc: + error = f"could not read Kaggle secret: {exc}" + attempts.append({"attempt": attempt_index, "error": error}) + if time.monotonic() >= deadline: + return { + "loaded": False, + "secret_name": secret_name, + "attempts": attempts, + "error": error, + } + print( + f"WANDB secret unavailable on attempt {attempt_index}; " + f"retrying in {sleep_seconds:.1f}s: {error}", + flush=True, + ) + time.sleep(sleep_seconds) + sleep_seconds = min(sleep_seconds * 1.5, 15.0) + continue + + if not value: + error = "Kaggle secret is empty" + attempts.append({"attempt": attempt_index, "error": error}) + return { + "loaded": False, + "secret_name": secret_name, + "attempts": attempts, + "error": error, + } + + os.environ["WANDB_API_KEY"] = str(value) + return { + "loaded": True, + "secret_name": secret_name, + "source": "kaggle_secret", + "attempt_count": attempt_index, + } + +def _ensure_uv_environment(summary: dict[str, Any]) -> None: + if os.environ.get(_WORKER_VENV_READY_ENV) == "1": + _activate_local_venv() + return + if shutil.which("uv") is None: + completed = subprocess.run( + [sys.executable, "-m", "pip", "install", "uv"], + check=False, + capture_output=True, + text=True, + ) + summary["uv_install"] = completed_summary(completed) + write_summary(summary) + if completed.returncode != 0: + raise SystemExit("Failed to install uv inside Kaggle worker.") + if Path("pyproject.toml").exists(): + _reset_worker_venv() + sync = _bootstrap_uv_environment(os.environ.get("KAGGLE_ACCELERATOR_ID", "")) + summary["uv_sync"] = { + "returncode": sync["returncode"], + "tpu_backend": sync["tpu_backend"], + "steps": sync["steps"], + } + write_summary(summary) + if int(sync["returncode"]) != 0: + log_bootstrap_failure(sync["steps"]) + raise SystemExit(format_bootstrap_failure(sync["steps"])) + _activate_local_venv() + _reexec_with_venv_python() + + +def _reexec_with_venv_python() -> None: + """Run the worker under the verified managed venv. + + The CUDA JAX probe succeeded with a specific venv + LD_LIBRARY_PATH. Carry + that exact class of environment into the re-executed worker instead of + re-applying generic runtime helpers. + """ + + venv_python = _venv_python().resolve() + if Path(sys.executable).resolve() == venv_python: + os.environ[_WORKER_VENV_READY_ENV] = "1" + _activate_local_venv() + return + env = _subprocess_env() + env[_WORKER_VENV_READY_ENV] = "1" + script_path = Path(__file__).resolve() + os.execve(str(venv_python), [str(venv_python), str(script_path), *sys.argv[1:]], env) + + +def _bootstrap_uv_environment(accelerator_id: str) -> dict[str, object]: + """Install project deps before importing ``src`` (Kaggle has no venv yet).""" + + sync = sync_kaggle_worker_environment(accelerator_id) + return { + "returncode": sync.returncode, + "tpu_backend": sync.tpu_backend, + "steps": [dict(step) for step in sync.steps], + } + + +def _reset_worker_venv() -> None: + """Recreate the worker venv every run; Kaggle working dirs persist stale JAX installs.""" + + venv = _worker_venv() + if venv.exists(): + shutil.rmtree(venv) + backend = ( + "tpu" + if _accelerator_requests_tpu(os.environ.get("KAGGLE_ACCELERATOR_ID", "")) + else "gpu" + ) + _JAX_BACKEND_MARKER.write_text(f"{backend}\n", encoding="utf-8") + + +def _subprocess_env() -> dict[str, str]: + env = os.environ.copy() + _apply_verified_worker_runtime_env(env=env) + return env + + +def _accelerator_requests_tpu(accelerator_id: str) -> bool: + normalized = accelerator_id.strip().lower() + return bool(normalized) and normalized.startswith("tpu") + + +def _accelerator_requests_nvidia(accelerator_id: str) -> bool: + normalized = accelerator_id.strip().lower() + return bool(normalized) and normalized.startswith("nvidia") + + +def _activate_local_venv() -> None: + _apply_verified_worker_runtime_env(env=os.environ) + for site_packages in sorted(_worker_venv().glob("lib/python*/site-packages")): + site_text = str(site_packages.resolve()) + if site_text not in sys.path: + sys.path.insert(0, site_text) + + +def _apply_verified_worker_runtime_env(*, env: dict[str, str]) -> None: + """Apply the runtime env that matches the successful JAX CUDA probe.""" + + venv = _worker_venv().resolve() + bin_dir = venv / "bin" + env[_WORKER_VENV_ENV] = str(venv) + env["VIRTUAL_ENV"] = str(venv) + env["UV_PROJECT_ENVIRONMENT"] = str(venv) + env["PYTHONNOUSERSITE"] = "1" + env.pop("PYTHONHOME", None) + if bin_dir.exists(): + env["PATH"] = _prepend_path(str(bin_dir), env.get("PATH", "")) + + accelerator_id = env.get("KAGGLE_ACCELERATOR_ID", os.environ.get("KAGGLE_ACCELERATOR_ID", "")) + env.pop("JAX_PLATFORM_NAME", None) + if _accelerator_requests_tpu(accelerator_id): + env["JAX_PLATFORMS"] = "tpu,cpu" + elif _accelerator_requests_nvidia(accelerator_id): + env["JAX_PLATFORMS"] = "cuda,cpu" + merged = [ + *_cuda_wheel_library_dirs(venv), + *_cuda_driver_library_dirs(env.get("LD_LIBRARY_PATH", "")), + ] + if merged: + env["LD_LIBRARY_PATH"] = os.pathsep.join(_dedupe(merged)) + elif os.environ.get("ORBIT_WARS_FORCE_JAX_CPU", "").strip().lower() in {"1", "true", "yes", "on"}: + env["JAX_PLATFORMS"] = "cpu" + + +def _cuda_wheel_library_dirs(venv: Path) -> list[str]: + dirs: list[str] = [] + for site_packages in sorted(venv.glob("lib/python*/site-packages")): + nvidia_root = site_packages / "nvidia" + if not nvidia_root.exists(): + continue + for lib_dir in sorted(nvidia_root.glob("*/lib")): + if lib_dir.is_dir(): + dirs.append(str(lib_dir.resolve())) + return _dedupe(dirs) + + +def _cuda_driver_library_dirs(existing_ld_library_path: str = "") -> list[str]: + candidates = [*_CUDA_DRIVER_LIBRARY_DIR_CANDIDATES] + candidates.extend(item for item in existing_ld_library_path.split(os.pathsep) if item) + candidates.extend(item for item in os.environ.get("LD_LIBRARY_PATH", "").split(os.pathsep) if item) + result: list[str] = [] + for item in candidates: + path = Path(item) + if not path.exists() or not path.is_dir(): + continue + if list(path.glob("libcuda.so*")) or list(path.glob("libnvidia-ml.so*")): + result.append(str(path.resolve())) + return _dedupe(result) + + +def _jax_platform_aliases(devices: list[object]) -> list[str]: + platforms: set[str] = set() + for device in devices: + platform = str(getattr(device, "platform", "") or "") + if platform: + platforms.add(platform) + device_text = str(device).lower() + if "cuda" in device_text: + platforms.add("cuda") + platforms.add("gpu") + if platform == "gpu": + platforms.add("cuda") + if platform == "cuda": + platforms.add("gpu") + return sorted(platforms) + + +def _prepend_path(prefix: str, existing: str) -> str: + return os.pathsep.join(_dedupe([prefix, *(item for item in existing.split(os.pathsep) if item)])) + + +def _dedupe(items: list[str] | tuple[str, ...]) -> list[str]: + seen: set[str] = set() + result: list[str] = [] + for item in items: + if item in seen: + continue + seen.add(item) + result.append(item) + return result + + +def _worker_venv() -> Path: + return Path(os.environ.get(_WORKER_VENV_ENV, ".venv")) + + +def _safe_worker_env() -> dict[str, str]: + keys = ( + "KAGGLE_ACCELERATOR_ID", + "ORBIT_WARS_KAGGLE_ALLOW_CPU", + "ORBIT_WARS_KAGGLE_ALLOW_CALIBRATION_FALLBACK", + "ORBIT_WARS_KAGGLE_CALIBRATION_MAX_VARIANTS", + "ORBIT_WARS_KAGGLE_CALIBRATION_TIMEOUT_SECONDS", + "ORBIT_WARS_KAGGLE_CALIBRATION_UPDATES", + "ORBIT_WARS_KAGGLE_CALIBRATION_WARMUP", + "ORBIT_WARS_KAGGLE_RUN_TYPE", + "ORBIT_WARS_KAGGLE_STANDALONE_OVERRIDES", + "ORBIT_WARS_KAGGLE_TRUST_BASE_JAX", + "ORBIT_WARS_KAGGLE_WORKER_MODE", + "WANDB_ENTITY", + "WANDB_API_KEY_SECRET_NAME", + "WANDB_PROJECT", + "WANDB_RESUME", + "WANDB_SWEEP_ID", + "WANDB_SWEEP_YAML", + ) + return {key: os.environ[key] for key in keys if key in os.environ} + + +def _env_int(name: str, default: int, *, minimum: int) -> int: + try: + value = int(os.environ.get(name, str(default))) + except ValueError: + return default + return max(value, minimum) + + +if __name__ == "__main__": + main() diff --git a/scripts/validate_kaggle_docker_submission.py b/scripts/validate_kaggle_docker_submission.py index 29dd4027..32e5c6a6 100644 --- a/scripts/validate_kaggle_docker_submission.py +++ b/scripts/validate_kaggle_docker_submission.py @@ -57,6 +57,7 @@ "jax/encoders/planet_encoder_common.py", "jax/encoders/planet_graph_transformer.py", "jax/encoders/remat.py", + "jax/factored_decode_scan.py", "jax/factored_sequence_scan.py", "jax/launch_hygiene.py", "jax/planet_flow.py", @@ -68,8 +69,10 @@ "jax/submission_runtime.py", "jax/env.py", "jax/map_pool/__init__.py", + "jax/map_pool/bake.py", "jax/map_pool/comets.py", "jax/map_pool/home_assignment.py", + "jax/map_pool/load.py", "jax/rollout/__init__.py", "jax/rollout/types.py", "jax/ship_action.py", diff --git a/src/cli/benchmark/rollout_phase_profile.py b/src/cli/benchmark/rollout_phase_profile.py index a0486f6e..9798664e 100644 --- a/src/cli/benchmark/rollout_phase_profile.py +++ b/src/cli/benchmark/rollout_phase_profile.py @@ -6,7 +6,6 @@ import json import os import subprocess -import sys from pathlib import Path from src.cli.benchmark.common import REPO_ROOT @@ -21,8 +20,7 @@ def _integration_root(args: argparse.Namespace) -> Path: raise SystemExit( "rollout-phase-profile requires the integration worktree " f"(missing timed collect at {root}). " - "Run from orbit_wars-integration or pass --repo-root.", - file=sys.stderr, + "Run from orbit_wars-integration or pass --repo-root." ) return root diff --git a/tests/test_cli_rollout_phase_profile.py b/tests/test_cli_rollout_phase_profile.py index 0562b3b2..18eecce0 100644 --- a/tests/test_cli_rollout_phase_profile.py +++ b/tests/test_cli_rollout_phase_profile.py @@ -7,6 +7,8 @@ from pathlib import Path from unittest.mock import patch +import pytest + from src.cli import benchmark as benchmark_cli from src.cli.benchmark.common import REPO_ROOT from src.cli.benchmark.rollout_phase_profile import run_rollout_phase_profile_cli @@ -91,8 +93,18 @@ def test_benchmark_main_rollout_phase_profile_dispatches(mock_profile) -> None: "--updates", "2", "--json", + "--repo-root", + str(REPO_ROOT), ] ) == 0 ) mock_profile.assert_called_once() + + +def test_rollout_phase_profile_cli_missing_integration_root_exits_cleanly( + tmp_path: Path, +) -> None: + args = Namespace(repo_root=tmp_path / "missing-integration") + with pytest.raises(SystemExit, match="requires the integration worktree"): + run_rollout_phase_profile_cli(args) diff --git a/tests/test_preflight_wandb_sweep_compose.py b/tests/test_preflight_wandb_sweep_compose.py index ff13bbf4..f64fa996 100644 --- a/tests/test_preflight_wandb_sweep_compose.py +++ b/tests/test_preflight_wandb_sweep_compose.py @@ -22,7 +22,7 @@ def test_preflight_wandb_sweep_compose() -> None: assert params["telemetry.metric_groups.losses"]["value"] is True assert params["task"]["value"] == "rollout_selected_validate" assert params["training.total_updates"]["value"] == 100 - assert params["training.reseed_every_updates"]["value"] == 50 + assert params["training.reseed_every_updates"]["value"] == 100 assert params["curriculum"]["value"] == "noop_only" assert "train_bundle" not in params assert "opponents" not in params diff --git a/tests/test_train_checkpoint.py b/tests/test_train_checkpoint.py index eaaa8c17..f46488af 100644 --- a/tests/test_train_checkpoint.py +++ b/tests/test_train_checkpoint.py @@ -272,7 +272,7 @@ def test_checkpoint_payload_builder_includes_curriculum_state() -> None: from src.config import TrainConfig from src.jax.policy import build_jax_policy from src.jax.train import init_train_state - from src.training.curriculum import CurriculumController + from src.opponents.curriculum import CurriculumController cfg = TrainConfig() cfg.task.candidate_count = 4 diff --git a/tests/test_trajectory_core.py b/tests/test_trajectory_core.py index c0872fdb..44ac7384 100644 --- a/tests/test_trajectory_core.py +++ b/tests/test_trajectory_core.py @@ -2,6 +2,8 @@ from __future__ import annotations +import pytest + from src.shield.trajectory_core import ( SHIELD_PARITY_PAIRS, acceptable_planet_hit, @@ -40,9 +42,11 @@ def test_line_circle_intersection_hits_sun_center() -> None: def test_bounds_exit_time_detects_right_edge() -> None: - hit = bounds_exit_time(0.5, 0.5, 25.0, 0.5) + # Board coordinates span [0, BOARD_SIZE] (100.0), so start just inside the right edge. + hit = bounds_exit_time(99.5, 50.0, 125.0, 50.0) assert hit is not None - assert 0.0 < hit <= 1.0 + assert hit == pytest.approx(0.5 / 25.5) + assert bounds_exit_time(0.5, 0.5, 25.0, 0.5) is None def test_acceptable_planet_hit_modes() -> None: