Close the pypi-cache loop: njs index merging, wants-collector, wheel-syncer, and bug fixes (#419)

**Impact:** pypi-cache module serving correctness, CI integration tests,
all clusters with pypi-cache deployed
**Risk:** medium

## What
Fixes the pypi-cache module to correctly serve packages by replacing the
error-page fallback with an njs-based index merge, adds the
wants-collector and wheel-syncer pipelines for the self-learning wheel
build system, fixes PEP 691 content negotiation bugs, and hardens
integration tests. Also bumps runner, BuildKit, and GPU plugin versions.

## Why
The pypi-cache nginx config had a fundamental index shadowing problem
(BY/BZ scenarios): when pypiserver returned HTTP 200 with wheels for the
_wrong_ CUDA variant or architecture, nginx's `proxy_intercept_errors`
never fired, so upstream PyPI was never consulted. Clients received an
index containing only incompatible wheels and failed to install. The fix
requires always merging both local and upstream indexes. Additionally,
the build pipeline that produces pre-built wheels
(pytorch/test-infra#7912) needs infrastructure to discover which
packages need building (wants-collector) and distribute built wheels to
pypiserver instances (wheel-syncer).

## How
- Replaced the `proxy_intercept_errors` + `error_page 404` fallback
pattern with an njs module (`merge_indexes.js`) that fires parallel
subrequests to both pypiserver and pypi.org, then merges the link lists
with filename-based deduplication (local wins)
- Used `parseFilesWithFallback()` to handle the format mismatch where
pypiserver v2.x returns HTML regardless of Accept header while pypi.org
returns PEP 691 JSON
- Moved upstream URL rewriting into the njs script since nginx's
`sub_filter` does not apply to subrequest response bodies
- Split `/packages/` into two regex locations: hash-based paths
(`/packages/[0-9a-f][0-9a-f]/`) route to files.pythonhosted.org, flat
paths route to pypiserver — preventing pypiserver from receiving
hash-based URLs it can't serve
- nginx now logs upstream-bound requests to EFS
(`/data/logs/upstream/fallback.YYYY-MM-DD.log`) using `map
$time_iso8601` for daily rotation, replacing the removed Python
log_rotator.py
- Increased nginx memory allocation to 64 GiB to accommodate njs
subrequest buffering (100m `subrequest_output_buffer_size` x workers x
concurrent requests)

```
  PyPI upstream cache durations

  ┌─────────────────────────────────────────────┬───────────┬───────────────┬───────────┐
  │                    What                     │ TTL (200) │ TTL (301/302) │ TTL (404) │
  ├─────────────────────────────────────────────┼───────────┼───────────────┼───────────┤
  │ Index pages (/simple/) from pypi.org        │ 10m       │ 1m            │ 1m        │
  ├─────────────────────────────────────────────┼───────────┼───────────────┼───────────┤
  │ Wheel downloads from files.pythonhosted.org │ 30 days   │ 30 days       │ 1m        │
  ├─────────────────────────────────────────────┼───────────┼───────────────┼───────────┤
  │ PyTorch wheels from download.pytorch.org    │ 30 days   │ 30 days       │ 1m        │
  ├─────────────────────────────────────────────┼───────────┼───────────────┼───────────┤
  │ PyTorch index (/whl)                        │ 10m       │ 1m            │ 1m        │
  ├─────────────────────────────────────────────┼───────────┼───────────────┼───────────┤
  │ Fallback (@pypi_fallback)                   │ 10m       │ 1m            │ 1m        │
  └─────────────────────────────────────────────┴───────────┴───────────────┴───────────┘

  Inactive eviction for all: 7 days (entries not accessed in 7d get purged regardless of TTL).

  Pypiserver (local) cache durations

  ┌───────────────────────────────────────────┬───────────┐
  │                   What                    │ TTL (200) │
  ├───────────────────────────────────────────┼───────────┤
  │ Index pages (/_internal/local/simple/)    │ 5m        │
  ├───────────────────────────────────────────┼───────────┤
  │ Local wheel downloads (.whl/.tar.gz/.zip) │ 30 days   │
  └───────────────────────────────────────────┴───────────┘
  ```

## Changes

### pypi-cache: njs index merging
- Added `merge_indexes.js` — njs module for merging local pypiserver + upstream pypi.org indexes
- Rewired `/simple/` location from `proxy_pass` + `error_page` to `js_content merge_indexes.mergeSimple`
- Added `/_internal/local/simple/` and `/_internal/upstream/simple/` internal locations with separate cache key prefixes and TTLs
- Split `/packages/` into hash-based (upstream proxy) and flat (pypiserver) locations with correct ordering
- Added daily log rotation via nginx `map`/`open_log_file_cache` for upstream fallback requests
- Removed `pypi-cache-scripts` ConfigMap and `log_rotator.py` sidecar — no longer needed
- Fixed `sub_filter_types` duplicate MIME type warning for `text/html`

### pypi-cache: wants-collector pipeline
- Added `wants_collector.py` — daemon that scans EFS access logs, checks PyPI JSON API for pre-built wheel availability across a target matrix (python versions x architectures x manylinux), uploads wants list and shared prebuilt cache to S3
- Added `wants-collector-deployment.yaml.tpl` with IRSA-annotated ServiceAccount, liveness probe, read-only root filesystem
- Added Terraform IRSA role with S3 PutObject/GetObject/ListBucket permissions
- Added `needbuild.txt` support for force-build overrides
- Added `cleanup_old_logs()` for EFS log hygiene
- Added comprehensive unit tests (`test_wants_collector.py` — 1043 lines covering all functions, edge cases, SIGTERM handling, loop behavior)

### pypi-cache: wheel-syncer pipeline
- Added `wheel_syncer.py` — daemon that syncs pre-built wheels from S3 to EFS wheelhouse per CUDA slug with atomic writes and path-traversal protection
- Added `wheel-syncer-deployment.yaml.tpl` with IRSA-annotated ServiceAccount, liveness probe, read-only root filesystem
- Added Terraform IRSA role with read-only S3 GetObject/ListBucket permissions
- Added comprehensive unit tests (`test_wheel_syncer.py` — 411 lines)

### pypi-cache: shared S3 bucket
- Added standalone Terraform root (`terraform/wheel-cache-bucket/`) for the shared `pytorch-pypi-wheel-cache` S3 bucket
- Public read on `wants/*`, `prebuilt-cache.txt`, `needbuild.txt`; 7-day lifecycle expiration on `wants/`
- Added Trivy ignore rules for intentional public-read S3 bucket (AWS-0086 through AWS-0094)

### pypi-cache: deploy.sh enhancements
- Deploy now reads `wants_collector_role_arn` and `wheel_syncer_role_arn` from Terraform outputs
- Annotates ServiceAccounts with IRSA role ARNs
- Creates ConfigMaps for wants-collector and wheel-syncer scripts
- Optional nginx cache clearing before restart (interactive prompt with `PYPI_CACHE_CLEAR` env override)
- Deploys and waits for wants-collector and wheel-syncer rollouts

### pypi-cache: smoke and e2e tests
- Added `test_pypi_cache.py` — 9 categories of live-cluster smoke tests (namespace, ServiceAccounts with IRSA, ConfigMaps, storage, deployments, services, NetworkPolicy, pod specs, NodePools)
- Added `test_pypi_cache_e2e.py` — e2e tests via kubectl exec (health, index responses, access logs, wants-collector cycle health + S3 validation, wheel-syncer cycle health + EFS validation, merged index PEP 691 content negotiation)
- Added `conftest.py` with session-scoped fixtures for pypi-cache pods, slugs, and pipeline pods

### pypi-cache: unit tests
- Added `test_merge_indexes.py` — 1281-line test suite for the njs merge algorithm reimplemented in Python, including a 9-scenario matrix (3 local states x 3 upstream states) for both HTML and JSON formats, and a dedicated mixed-format test class reproducing the exact PEP 691 bug
- Updated `test_generate_manifests.py` — adjusted for nginx memory increase (2Gi -> 64Gi), removed log_rotator references, updated pypiserver volume mount assertions

### pypi-cache: configuration
- Added `target_architectures` and `target_manylinux` to `clusters.yaml` defaults
- Increased nginx default memory from 2 GiB to 64 GiB for njs subrequest buffering

### Integration tests
- Added retry loops with configurable max retries for pypi-cache health checks
- Changed health checks from `/simple/` to `/health` endpoint
- Added environment variable verification step
- Added numpy wheel cache validation tests (local wheel download, platform tag validation, functional validation) — gracefully skips when wheel cache is not yet populated
- Expanded CPU and CUDA action test jobs with full test battery (uv compile+sync, uv project interface, URL rewriting verification, UV_INDEX_STRATEGY check)
- CUDA test job now runs on GPU runner (`x86iavx512-29-115-t4`) instead of CPU-only runner

### Version bumps
- GitHub Actions Runner: 2.313.0 -> 2.333.1 (Dockerfile + clusters.yaml + generate_runners.py)
- BuildKit: v0.28.0 -> v0.29.0 (generate_buildkit.py + integration test workflow)
- NVIDIA device plugin: v0.14.5 -> v0.19.0

### Smoke test reliability
- Added agent mode (`AGENT_ENVIRONMENT=true`) to `just test`, `just lint`, and `just smoke` — suppresses output on success, dumps full output on failure
- Increased `MIN_NODE_AGE_SECONDS` from 120s to 600s for more stable node filtering
- Added `get_recently_stable_node_names()` helper to tolerate Pending DaemonSet pods on nodes younger than 20 minutes
- `assert_daemonset_healthy` now accepts a configurable `min_node_age` parameter
- `assert_deployment_ready` now detects and waits for active rollouts instead of immediately failing
- cache-enforcer smoke test now skips when no runner nodes exist (desired == 0)
- Alloy logging smoke test tolerates Pending pods on recently-stable nodes
- git-cache smoke test uses `min_node_age=900` for workload-type filtered DaemonSet checks
- Node compactor e2e rollout timeout increased from 120s to 200s (exceeds 120s terminationGracePeriodSeconds)

### Helm value fixes
- Alloy logging: moved `mounts` config under `alloy` key (was incorrectly under `controller`)
- Alloy events: moved `resources` from `controller` to `alloy` key
- Alloy monitoring: added explicit resource requests/limits (250m/512Mi to 2/1Gi)

### Cleanup
- Removed 11 module-level CLAUDE.md files (base, node-compactor, arc-runners, arc, buildkit, cache-enforcer, eks, karpenter, logging, monitoring, nodepools, pypi-cache) — this information lives in skills and the upstream CLAUDE.md
- Deleted orphaned git-cache-server StatefulSet migration added to git-cache deploy.sh

## Notes
- The wheel build pipeline itself is in a separate PR (pytorch/test-infra#7912) — this PR provides the infrastructure (wants-collector discovers demand, wheel-syncer distributes results)
- The `pytorch-pypi-wheel-cache` S3 bucket is intentionally public-read for cross-cluster access without cross-account IAM
- The nginx memory increase to 64 GiB is driven by njs subrequest buffering needs — pypi.org's root `/simple/` listing is ~40 MB, and with 8 worker processes and concurrent requests the worst-case buffer usage approaches 1.6 GB; 64 GiB provides ample headroom on dedicated r5d.12xlarge nodes

## Testing
- All unit tests pass (`just test`) — including ~2700 lines of new test code for merge_indexes, wants_collector, and wheel_syncer
- All linters pass (`just lint`)
- Integration tests expanded with retry logic, health endpoint checks, environment variable verification, numpy wheel cache validation, and full pip/uv test battery across CPU and CUDA configurations
- Smoke tests cover all new infrastructure (ServiceAccounts, IRSA annotations, ConfigMaps, Deployments, Services, pipeline health, merged index content negotiation)

### Integration tests expanded to carefully test every aspect of pypi-cache
```
 $  just integration-test arc-staging
Updating kubeconfig for pytorch-arc-staging (us-west-1)...
Updated context pytorch-arc-staging in /Users/jschmidt/.kube/config
10:34:03 [INFO] Integration test for cluster: arc-staging
(pytorch-arc-staging)
10:34:03 [INFO]   Runner prefix: 'c-mt-'
10:34:03 [INFO]   B200 enabled: False
10:34:03 [INFO]   Cache enforcer: True
10:34:03 [INFO]   PyPI cache slugs: cpu cu128 cu130
10:34:03 [INFO]   Smoke tests: skip
10:34:03 [INFO]   Compactor tests: skip
10:34:03 [INFO]   Branch: osdc-integration-test-arc-staging
10:34:03 [INFO] Phase 0: Cleaning up stale PRs...
10:34:04 [INFO] Phase 1: Checking for active runner pods (arc-staging
only)...
10:34:07 [INFO]   No runner pods active. Skipping pool clear.
10:34:07 [INFO] Canary repo already cloned at
/Users/jschmidt/meta/ciforge/osdc/upstream/osdc/.scratch/pytorch-canary,
fetching...
10:34:08 [INFO] Phase 2: Preparing PR...
10:34:14 [INFO] PR #373 created:
pytorch/pytorch-canary#373
10:34:14 [INFO] Phase 3: Running parallel validation...
10:34:14 [INFO] Phase 4: Waiting for PR workflow runs (timeout: 50 min,
buffer: 10 min)...
10:34:14 [INFO] Filtering to runs created after
2026-04-02T17:34:08.442162+00:00
10:34:15 [INFO]   No runs found yet, waiting...
10:34:46 [INFO] Run: OSDC Integration Test —
https://github.com/pytorch/pytorch-canary/actions/runs/23913510770
10:34:46 [INFO]   1/1 runs still in progress...
10:35:18 [INFO]   1/1 runs still in progress...
10:35:49 [INFO]   1/1 runs still in progress...
10:36:20 [INFO]   1/1 runs still in progress...
10:36:51 [INFO]   1/1 runs still in progress...
10:37:22 [INFO]   1/1 runs still in progress...
10:37:53 [INFO]   1/1 runs still in progress...
10:38:25 [INFO]   1/1 runs still in progress...
10:38:56 [INFO]   1/1 runs still in progress...
10:39:27 [INFO]   All 1 run(s) completed.


============================================================
  OSDC Integration Test Results
============================================================
  Cluster: arc-staging (pytorch-arc-staging)
  Date:    2026-04-02 17:39 UTC

  PR Workflow Jobs:
    ✓ test-cpu-x86-amx               success
    ✓ test-harbor                    success
    ✓ test-cpu-arm64                 success
    ✓ test-gpu-t4                    success
    ✓ test-cache-enforcer            success
    ✓ test-pypi-cache-defaults       success
    ✓ test-git-cache                 success
    ✓ test-pypi-cache-action-cpu     success
    ✓ test-cpu-x86-avx512            success
    ✓ test-gpu-t4-multi              success
    ✓ build-amd64 / build            success
    ✓ test-pypi-cache-action-cuda    success
    ✓ build-arm64 / build            success

  Smoke            ⊘ SKIPPED
  Compactor        ⊘ SKIPPED

  Overall: PASSED
============================================================

10:39:29 [INFO] Phase 5: Closing PR #373...
10:39:31 [INFO] Total integration test time: 5m28s
```

### Failing smoke tests are related to yangs changes/experiment with monitoring not yet reflected on this PR. Everything related to this PR is green.
```

====================================================================================================================================================
test session starts
=====================================================================================================================================================
platform darwin -- Python 3.12.12, pytest-9.0.2, pluggy-1.6.0
rootdir: /Users/jschmidt/meta/ciforge/osdc/upstream/osdc
configfile: pyproject.toml
plugins: anyio-4.12.1, xdist-3.8.0, cov-7.0.0
16 workers [193 items]

.....................................................................................................s..................s...s......................s............s.........................F.FssF.
[100%]

==========================================================================================================================================================
FAILURES
==========================================================================================================================================================

_________________________________________________________________________________________________________________
TestMonitoringPerTargetVerification.test_kube_prom_stack_target_fresh[kube-state-metrics]
__________________________________________________________________________________________________________________
[gw3] darwin -- Python 3.12.12
/Users/jschmidt/meta/ciforge/osdc/upstream/osdc/.venv/bin/python3
modules/monitoring/tests/smoke/test_monitoring.py:364: in
test_kube_prom_stack_target_fresh
    assert_metric_fresh_in_mimir(
tests/smoke/helpers.py:590: in assert_metric_fresh_in_mimir
raise AssertionError(f"[{label}] No metric series found after
{REMOTE_RETRIES} attempts")
E AssertionError: [kube-prom-stack target: kube-state-metrics] No metric
series found after 5 attempts

__________________________________________________________________________________________________________________________________
TestLoggingPerSourceVerification.test_pod_logs_arriving
___________________________________________________________________________________________________________________________________
[gw9] darwin -- Python 3.12.12
/Users/jschmidt/meta/ciforge/osdc/upstream/osdc/.venv/bin/python3
modules/logging/tests/smoke/test_logging.py:328: in
test_pod_logs_arriving
    assert_logs_fresh_in_loki(
tests/smoke/helpers.py:654: in assert_logs_fresh_in_loki
    assert last_age < max_staleness, (
           ^^^^^^^^^^^^^^^^^^^^^^^^
E AssertionError: [pod logs (kube-system)] Logs are stale after 5
attempts: newest entry is 1093s old (threshold: 600s)

_______________________________________________________________________________________________________________________________________
TestMonitoringHelm.test_helm_release_deployed
________________________________________________________________________________________________________________________________________
[gw14] darwin -- Python 3.12.12
/Users/jschmidt/meta/ciforge/osdc/upstream/osdc/.venv/bin/python3
modules/monitoring/tests/smoke/test_monitoring.py:60: in
test_helm_release_deployed
assert release is not None, "Helm release 'kube-prometheus-stack' not
found"
E   AssertionError: Helm release 'kube-prometheus-stack' not found
E   assert None is not None

==================================================================================================================================================
short test summary info
===================================================================================================================================================
SKIPPED [1] modules/pypi-cache/tests/smoke/test_pypi_cache.py:394: No
instance_type configured — pypi-cache uses shared base nodes
SKIPPED [1] modules/pypi-cache/tests/smoke/test_pypi_cache.py:385: No
instance_type configured — pypi-cache uses shared base nodes
SKIPPED [1] modules/pypi-cache/tests/smoke/test_pypi_cache.py:357: No
instance_type configured — no dedicated node taint to tolerate
SKIPPED [1]
modules/cache-enforcer/tests/smoke/test_cache_enforcer.py:107: No
cache-enforcer pods desired (no runner nodes in cluster)
SKIPPED [1] modules/monitoring/tests/smoke/test_monitoring.py:172: No
dcgm-exporter pods found (no GPU nodes)
SKIPPED [1]
modules/cache-enforcer/tests/smoke/test_cache_enforcer.py:346: No
cache-enforcer pods found (no runner nodes in cluster)
SKIPPED [1]
modules/cache-enforcer/tests/smoke/test_cache_enforcer.py:306: No
cache-enforcer pods found (no runner nodes in cluster)

====================================================================================================================================
3 failed, 183 passed, 7 skipped in 158.04s (0:02:38)
====================================================================================================================================
```

### Workload tests succeed up to the pypi install and only fail as they are not properly handling repository changes (not being tested yet)
https://github.com/pytorch/pytorch-canary/actions/runs/23914496937

---------

Signed-off-by: Jean Schmidt <contato@jschmidt.me>

Author: jeanschmidt

osdc/.trivyignore

@@ -29,6 +29,19 @@ AWS-0132
 # already restrict inbound traffic.
 AWS-0164
 
+# PyPI wheel cache S3 bucket is intentionally public-read. The bucket stores
+# only package availability lists (wants/*.txt) and a prebuilt cache — no
+# sensitive data. Public access lets any cluster read the shared cache without
+# cross-account IAM.
+AWS-0086
+AWS-0087
+AWS-0088
+AWS-0089
+AWS-0090
+AWS-0091
+AWS-0093
+AWS-0094
+
 # =============================================================================
 # Kubernetes — privileged access (intentional for system DaemonSets)
 # =============================================================================

osdc/base/CLAUDE.md

@@ -16,7 +16,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     && rm -rf /var/lib/apt/lists/*
 
 # Install GitHub Actions Runner
-ARG RUNNER_VERSION=2.313.0
+ARG RUNNER_VERSION=2.333.1
 RUN useradd -m -s /bin/bash runner \
     && mkdir -p /home/runner/actions-runner
 

osdc/base/docker/runner-base/Dockerfile

@@ -76,6 +76,15 @@ if kubectl get deployment git-cache-central -n kube-system &>/dev/null; then
   sleep 5
 fi
 
+# --- Migration: delete orphaned git-cache-server StatefulSet ---
+# git-cache-server was an earlier iteration replaced by git-cache-central.
+# Clean it up if still present.
+if kubectl get statefulset git-cache-server -n kube-system &>/dev/null; then
+  echo "Deleting orphaned git-cache-server StatefulSet..."
+  kubectl delete statefulset git-cache-server -n kube-system --wait=false
+  sleep 5
+fi
+
 # --- Apply static resources (kustomize) ---
 echo "Applying static git-cache resources..."
 kubectl_apply_if_changed -k "$SCRIPT_DIR/"

osdc/base/kubernetes/git-cache/deploy.sh

@@ -68,6 +68,7 @@ def test_daemonset_exists_and_ready(self, all_daemonsets, all_nodes):
             GIT_CACHE_NAMESPACE,
             "git-cache-warmer",
             node_selector={"workload-type": ["github-runner", "buildkit"]},
+            min_node_age=900,
         )
 
 

osdc/base/kubernetes/git-cache/tests/smoke/test_git_cache.py

@@ -42,7 +42,7 @@ spec:
           effect: NoSchedule
       priorityClassName: system-node-critical
       containers:
-        - image: nvcr.io/nvidia/k8s-device-plugin:v0.14.5
+        - image: nvcr.io/nvidia/k8s-device-plugin:v0.19.0
           name: nvidia-device-plugin-ctr
           args:
             - "--fail-on-init-error=false"

osdc/base/kubernetes/nvidia-device-plugin.yaml

@@ -502,12 +502,16 @@ def get_compactor_pod_names(client: Client) -> set[str]:
     return names
 
 
-def wait_for_compactor_rollout(client: Client, old_pod_names: set[str], timeout_s: int = 120) -> None:
+def wait_for_compactor_rollout(client: Client, old_pod_names: set[str], timeout_s: int = 200) -> None:
     """Wait for the Deployment rollout to replace old pods with a new Running pod.
 
     After ``patch_compactor_env`` modifies the pod template, Kubernetes
     automatically creates a new ReplicaSet and pod.  This waits for a
     Running pod whose name is NOT in *old_pod_names*.
+
+    The default timeout (200s) must exceed ``terminationGracePeriodSeconds``
+    (120s) because the Recreate strategy requires the old pod to fully
+    terminate before the new pod can start.
     """
 
     def _new_pod_running() -> bool:

osdc/base/node-compactor/CLAUDE.md

@@ -50,7 +50,7 @@ defaults:
     pdb_min_available: 1
   arc:
     chart_version: "0.14.0"  # MUST match for controller + runner charts (minor version mismatch = ARS deletion)
-    runner_image_tag: "2.333.0"  # ghcr.io/actions/actions-runner tag — pin to avoid :latest
+    runner_image_tag: "2.333.1"  # ghcr.io/actions/actions-runner tag — pin to avoid :latest
     replica_count: 2
     log_level: info
     controller_cpu_request: "1"
@@ -90,6 +90,10 @@ defaults:
       - "3.13t"
       - "3.14"
       - "3.14t"
+    target_architectures:
+      - "x86_64"
+      - "aarch64"
+    target_manylinux: "2_17"
 
 clusters:
   arc-staging:

osdc/base/node-compactor/tests/e2e/helpers.py

@@ -22,7 +22,7 @@ jobs:
     # The arch input selects which BuildKit endpoint to connect to.
     runs-on: ${{ inputs.runner_label }}
     container:
-      image: moby/buildkit:v0.28.0
+      image: moby/buildkit:v0.29.0
       # Run as root to use buildctl
       options: --privileged
     steps: