ci: make CI resilient to pip/uv network timeouts

.github/actions/action.yml

@@ -127,6 +127,9 @@ runs:
         export PYTHONPATH=$(pwd)
         export NEMORUN_HOME=$(pwd)
         export NCCL_DEBUG=INFO
+        export PIP_DEFAULT_TIMEOUT=120
+        export PIP_RETRIES=5
+        export UV_HTTP_TIMEOUT=120
         uv venv .venv
         uv cache clean
         uv sync --no-cache --only-group test
@@ -169,6 +172,9 @@ runs:
 
         export PYTHONPATH=$(pwd)
         export NEMORUN_HOME=$(pwd)
+        export PIP_DEFAULT_TIMEOUT=120
+        export PIP_RETRIES=5
+        export UV_HTTP_TIMEOUT=120
         uv venv .venv
         uv cache clean
         uv sync --no-cache --only-group test

.github/workflows/_build_test_publish_wheel.yml

@@ -72,6 +72,11 @@ jobs:
 
           pushd $BUILD_DIR
             rm LICENSE || true
+            for i in 1 2 3; do
+              docker pull "$IMAGE" && break
+              echo "docker pull attempt $i failed, retrying..."
+              sleep 10
+            done
             docker run --rm -e NO_VCS_VERSION=1 -v $(pwd):/workspace -w /workspace $IMAGE bash -c '\
               for python_version in cp311 cp312 cp313; do \
                 /opt/python/${python_version}-${python_version}/bin/pip install --upgrade "setuptools>=80" build; \

.github/workflows/cicd-main.yml

@@ -716,6 +716,8 @@ jobs:
       PIP_DISABLE_PIP_VERSION_CHECK: 1
       PIP_NO_PYTHON_VERSION_WARNING: 1
       PIP_ROOT_USER_ACTION: ignore
+      PIP_DEFAULT_TIMEOUT: 120
+      PIP_RETRIES: 5
     steps:
       - name: Checkout
         uses: actions/checkout@v6
@@ -865,6 +867,8 @@ jobs:
       PIP_DISABLE_PIP_VERSION_CHECK: 1
       PIP_NO_PYTHON_VERSION_WARNING: 1
       PIP_ROOT_USER_ACTION: ignore
+      PIP_DEFAULT_TIMEOUT: 120
+      PIP_RETRIES: 5
     if: |
       !cancelled()
       && needs.cicd-integration-gate.outputs.should_run == 'true'
@@ -962,6 +966,8 @@ jobs:
       PIP_DISABLE_PIP_VERSION_CHECK: 1
       PIP_NO_PYTHON_VERSION_WARNING: 1
       PIP_ROOT_USER_ACTION: ignore
+      PIP_DEFAULT_TIMEOUT: 120
+      PIP_RETRIES: 5
     if: |
       !cancelled()
       && needs.cicd-integration-gate.outputs.should_run == 'true'

tests/test_utils/python_scripts/launch_nemo_run_workload.py

@@ -38,13 +38,50 @@ def is_flaky_failure(concat_allranks_logs: str) -> bool:
         or "zmq.error.ZMQError: Address already in use" in concat_allranks_logs
         or "We couldn't connect to 'https://huggingface.co'" in concat_allranks_logs
         or "Unpack failed: incomplete input" in concat_allranks_logs
+        or "The read operation timed out" in concat_allranks_logs
+        or "Read timed out" in concat_allranks_logs
+        or "TimeoutError" in concat_allranks_logs
+        or "Connection broken" in concat_allranks_logs
+        or "Temporary failure in name resolution" in concat_allranks_logs
         or "unspecified launch failure" in concat_allranks_logs
         or "free(): corrupted unsorted chunks" in concat_allranks_logs
         or "Segfault encountered" in concat_allranks_logs
         or "The following metrics failed" in concat_allranks_logs
     )
 
 
+def _collect_failure_logs(workdir: pathlib.Path) -> list[str]:
+    """Reads every log file that may carry a flaky-failure signature.
+
+    The per-rank ``attempt_0/*/std*.log`` files only contain torchrun training
+    output. The golden-value comparison runs in ``run_ci_test.sh`` and emits its
+    assertion (e.g. ``The following metrics failed``) to the nemo-run task log,
+    which lives outside that per-rank tree. Globbing both ensures harness-side
+    failures are seen by ``is_flaky_failure`` and therefore eligible for retry.
+
+    Args:
+        workdir: Working directory under which nemo-run writes all log files.
+
+    Returns:
+        The concatenated lines of every discovered log file, deduplicated by
+        resolved path to avoid double-counting overlapping globs.
+    """
+    seen_paths = set()
+    collected_lines: list[str] = []
+    for pattern in ("**/attempt_0/*/std*.log", "**/*.log"):
+        for log_file_path in workdir.glob(pattern):
+            resolved = log_file_path.resolve()
+            if resolved in seen_paths or not log_file_path.is_file():
+                continue
+            seen_paths.add(resolved)
+            try:
+                with open(log_file_path, "r", errors="replace") as f:
+                    collected_lines.extend(f.readlines())
+            except OSError as error:
+                logger.warning("Could not read log file %s: %s", log_file_path, error)
+    return collected_lines
+
+
 @click.command()
 @click.option("--scope", required=True, type=str, help="Scope of the workload")
 @click.option("--model", required=True, type=str, help="Model of the workload")
@@ -194,12 +231,8 @@ def __getattr__(self, name):
             sys.exit(0)
 
         logger.error(f"Job failed with status: {job_dict['status']}")
-        log_file_paths = pathlib.Path(os.getcwd()).glob("**/attempt_0/*/std*.log")
         all_ranks_all_logs = [tee_buffer.getvalue()]
-        for log_file_path in log_file_paths:
-            with open(log_file_path, "r") as f:
-                all_logs = f.readlines()
-            all_ranks_all_logs.extend(all_logs)
+        all_ranks_all_logs.extend(_collect_failure_logs(pathlib.Path(os.getcwd())))
         all_ranks_all_logs_string = "\n".join(all_ranks_all_logs)
         if is_flaky_failure(all_ranks_all_logs_string):
             logger.warning("Detected flaky failure, attempt restart.")