ci: make CI resilient to pip/uv network timeouts

[ko3n1g]

Claude summary

Background / motivation

Several CI jobs fail on transient network timeouts during dependency / image fetches, with no retry cushion:

• Host setup-python pip upgrade — on a fresh toolcache, actions/setup-python@v5 fetches the latest pip (26.1.2) from PyPI; the download hit a read timeout and aborted the whole job before any test ran:

  TimeoutError: The read operation timed out
  Error: The process '/usr/bin/bash' failed with exit code 2
  

• manylinux wheel build — the megatron-fsdp build's docker run implicitly pulled quay.io/pypa/manylinux_2_28_x86_64 and the registry request timed out:

  docker: Error response from daemon: Get "https://quay.io/v2/": net/http: request canceled while waiting for connection (Client.Timeout exceeded while awaiting headers)
  ##[error]Process completed with exit code 125.
  

pip's default per-read socket timeout is ~15s with no retries; the docker pull had no retry at all. The same flakiness hits in-container uv sync / pip.

What changed

• cicd-main.yml — add PIP_DEFAULT_TIMEOUT=120 + PIP_RETRIES=5 to the three job-level env: blocks, inherited by the host setup-python pip upgrade so it retries instead of hard-failing.
• action.yml — export PIP_DEFAULT_TIMEOUT / PIP_RETRIES / UV_HTTP_TIMEOUT=120 in both generated run-scripts so in-container uv sync / pip get the same cushion.
• _build_test_publish_wheel.yml — pre-pull the manylinux image with the repo's standard 3× retry loop before docker run, so a transient quay.io timeout no longer fails the job (later runs in the job reuse the cached image).
• launch_nemo_run_workload.py — classify read/connection timeouts (The read operation timed out, Read timed out, TimeoutError, Connection broken, Temporary failure in name resolution) as flaky so the launcher auto-restarts; and collect harness/task logs (**/*.log), not just per-rank attempt_0/*/std*.log, so those signatures are actually visible to is_flaky_failure.

Details

• Host pip env vars sit beside the existing PIP_DISABLE_PIP_VERSION_CHECK / PIP_NO_PYTHON_VERSION_WARNING / PIP_ROOT_USER_ACTION block — three identical job-level env: maps (latest, h100, gb200).
• The docker pre-pull mirrors the existing apt-get / uv install retry loops in action.yml (for i in 1 2 3; do … && break; sleep 10; done).
• _collect_failure_logs() dedups by resolved path and is OSError-tolerant; it widens the glob so the golden-value assertion (written to the nemo-run task log, outside the per-rank tree) is no longer missed.

Tested

• ruff check, isort --check, black==24.4.2 --check — all pass on the changed Python file.
• YAML parse of all three workflow files — valid.
• py_compile — OK.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[ko3n1g]

/ok to test

[ko3n1g]

/ok to test

[greptile-apps]

Greptile Summary

This PR improves CI resilience against transient network failures by setting pip/uv timeout and retry environment variables at the workflow and container levels, widening the flaky-failure detection in the NeMo-run launcher, and adding a retry loop around docker pull in the wheel-build workflow.

• cicd-main.yml / action.yml: PIP_DEFAULT_TIMEOUT=120, PIP_RETRIES=5, and UV_HTTP_TIMEOUT=120 are added to three job-level env: blocks and to both generated run-scripts so host-side setup-python pip upgrades and in-container uv sync operations retry on slow PyPI responses.
• launch_nemo_run_workload.py: Five new flaky-failure strings are added (including a broad \"TimeoutError\"), and log collection is refactored into _collect_failure_logs() which widens the glob to **/*.log so harness-side assertion failures are visible to the retry logic.
• _build_test_publish_wheel.yml: A three-attempt retry loop is added around docker pull before the wheel-build container run, consistent with existing retry patterns for apt-get and uv installs.

Confidence Score: 3/5

Safe to merge for the CI infrastructure changes; the Python launcher change introduces a flaky-failure pattern that could silently retry genuine training failures.

The YAML changes are straightforward and low-risk. Adding the bare string TimeoutError to the flaky-failure list means any training job that raises a Python TimeoutError will be retried rather than reported as a failure, making regressions harder to catch. The **/*.log glob is also unbounded and could read large amounts of data into memory.

tests/test_utils/python_scripts/launch_nemo_run_workload.py — the new flaky-failure patterns and the recursive log-collection glob warrant a closer look before merging.

Important Files Changed

Filename	Overview
tests/test_utils/python_scripts/launch_nemo_run_workload.py	New _collect_failure_logs helper widens log collection to all **/*.log files; new flaky-failure strings include an overly broad "TimeoutError" that could mask real training failures.
.github/actions/action.yml	Adds PIP_DEFAULT_TIMEOUT, PIP_RETRIES, and UV_HTTP_TIMEOUT exports to both generated run-scripts (unit-test and e2e) to prevent in-container pip/uv network failures.
.github/workflows/cicd-main.yml	Adds PIP_DEFAULT_TIMEOUT: 120 and PIP_RETRIES: 5 to three job-level env: blocks (latest, h100, gb200) — straightforward, consistent addition.
.github/workflows/_build_test_publish_wheel.yml	Adds a 3-attempt retry loop around docker pull before the wheel-build container run; consistent with existing retry patterns in the repository.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[CI Job Starts] --> B[actions/setup-python\nPIP_DEFAULT_TIMEOUT=120\nPIP_RETRIES=5]
    B --> C[Install uv - retry loop x3]
    C --> D[Generate run-script\nexport PIP_DEFAULT_TIMEOUT=120\nexport PIP_RETRIES=5\nexport UV_HTTP_TIMEOUT=120]
    D --> E[uv sync - with timeout/retry]
    E --> F[launch_nemo_run_workload.py]
    F --> G{Job Status}
    G -->|Success| H[Exit 0]
    G -->|Failure| I[_collect_failure_logs\nattempt_0/**/*.log\n+ **/*.log]
    I --> J[is_flaky_failure check]
    J -->|Matched| K[Auto-retry]
    J -->|Not matched| L[Exit with error]
    K --> F

Loading

_{Reviews (1): Last reviewed commit: "ci: retry manylinux image pull in wheel ..." | Re-trigger Greptile}

[greptile-apps]

"TimeoutError" match is too broad

TimeoutError is a Python builtin that appears in stack traces for any kind of timeout — NCCL collective timeouts, test-harness wait timeouts, asyncio cancellations, custom training deadlines — not just network I/O. Any genuine performance regression that causes a training job to exceed its allotted time will match this string, get classified as flaky, and be silently retried up to the configured number of attempts, masking the regression entirely. The other two new network patterns ("The read operation timed out" and "Read timed out") are already specific enough to catch the pip/uv scenario described in the PR. Consider removing the bare "TimeoutError" or replacing it with a more qualified string such as "TimeoutError: The read operation" to keep it scoped to network I/O.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[greptile-apps]

Unbounded **/*.log glob may read very large log sets into memory

workdir.glob("**/*.log") is recursive and matches every .log file under the working directory, including build artefacts, prior run logs, or any container-mounted directories that happen to be under cwd. In a large CI environment with multiple prior attempts or deep output trees, this can pull hundreds of megabytes of log data into a single in-memory list, potentially causing the launcher process itself to be OOM-killed. Bounding the search depth or capping total bytes read would guard against this.

[balasaajay]

looks good.

just a heads up: there was a requirement earlier today to surface the actual error in a different step so its easier to jump to.