[ML] Fix PyTorch allowlist validation timeout on HF download stall

.buildkite/hooks/post-checkout

@@ -36,7 +36,7 @@ if [[ "$BUILDKITE_PIPELINE_SLUG" == ml-cpp* ]]; then
   # GCS service account — inject credentials for build and Java IT steps.
   # Build steps use it for sccache; Java IT steps use it for the Gradle
   # build cache.  The key is stored in Vault.
-  if [[ "$BUILDKITE_STEP_KEY" == build_test_* || "$BUILDKITE_STEP_KEY" == java_integration_tests_* ]]; then
+  if [[ "$BUILDKITE_STEP_KEY" == build_test_* || "$BUILDKITE_STEP_KEY" == java_integration_tests_* || "$BUILDKITE_STEP_KEY" == "build_pytorch_docker_image" ]]; then
     SCCACHE_GCS_KEY_JSON=$(vault read -field=key secret/ci/elastic-ml-cpp/sccache/gcs_service_account 2>/dev/null || echo "")
     if [ -n "$SCCACHE_GCS_KEY_JSON" ]; then
       export SCCACHE_GCS_BUCKET="elastic-ml-cpp-sccache"
@@ -48,6 +48,13 @@ if [[ "$BUILDKITE_PIPELINE_SLUG" == ml-cpp* ]]; then
     fi
   fi
 
+  if [[ "$BUILDKITE_STEP_KEY" == "validate_pytorch_allowlist" ]]; then
+    HF_TOKEN=$(vault read -field=token secret/ci/elastic-ml-cpp/huggingface/hf_token 2>/dev/null || echo "")
+    if [ -n "$HF_TOKEN" ]; then
+      export HF_TOKEN
+    fi
+  fi
+
   if [[ "$BUILDKITE_STEP_KEY" == "build_pytorch_docker_image" ]]; then
     export DOCKER_REGISTRY_USERNAME=$(vault read --field=username  secret/ci/elastic-ml-cpp/prod_docker_registry_credentials)
     export DOCKER_REGISTRY_PASSWORD=$(vault read --field=password  secret/ci/elastic-ml-cpp/prod_docker_registry_credentials)

.buildkite/pipelines/validate_pytorch_allowlist.yml.sh

@@ -8,14 +8,22 @@
 # compliance with the Elastic License 2.0 and the foregoing additional
 # limitation.
 
-cat <<'EOL'
+# Use the same Docker image as the build steps — it has Python 3.12 and
+# the source-built torch package, giving exact version parity with the
+# libtorch that pytorch_inference links against.
+VALIDATION_IMAGE="${DOCKER_IMAGE:-docker.elastic.co/ml-dev/ml-linux-build:34}"
+
+cat <<EOL
 steps:
   - label: "Validate PyTorch allowlist :torch:"
     key: "validate_pytorch_allowlist"
     timeout_in_minutes: 60
+    env:
+        HF_HUB_DISABLE_XET: "1"
     command:
         - "if [ ! -f dev-tools/extract_model_ops/validate_allowlist.py ]; then echo 'validate_allowlist.py not found, skipping'; exit 0; fi"
-        - "pip install -r dev-tools/extract_model_ops/requirements.txt"
+        - "python3 -c \"import torch; print(f'PyTorch version: {torch.__version__}')\""
+        - "grep -v '^torch==' dev-tools/extract_model_ops/requirements.txt | pip3 install -r /dev/stdin"
         - "python3 dev-tools/extract_model_ops/validate_allowlist.py --config dev-tools/extract_model_ops/validation_models.json --pt-dir dev-tools/extract_model_ops/es_it_models --verbose"
 EOL
 
@@ -29,10 +37,10 @@ if [ -n "${ML_BUILD_STEP_KEYS:-}" ]; then
     done
 fi
 
-cat <<'EOL'
+cat <<EOL
     allow_dependency_failure: true
     agents:
-      image: "python:3.12"
+      image: "${VALIDATION_IMAGE}"
       memory: "32G"
       ephemeralStorage: "30G"
     notify:

dev-tools/docker/build_pytorch_linux_build_image.sh

@@ -17,41 +17,95 @@
 # issues we may have with the latest PyTorch release as early as possible.
 # To avoid swamping the Docker registry with numerous images we re-use the
 # same tag for the image.
+#
+# Optimizations:
+#   1. Skip if viable/strict hasn't moved since the last build
+#   2. sccache using a Docker secret for the GCS backend for fast incremental rebuilds
 
-if [ `uname -m` != x86_64 ] ; then
+if [ "$(uname -m)" != x86_64 ] ; then
     echo "Native build images must be built on the correct hardware architecture"
-    echo "Required: x86_64, Current:" `uname -m`
+    echo "Required: x86_64, Current: $(uname -m)"
     exit 1
 fi
 
-DOCKER_DIR=`docker info 2>/dev/null | grep '^ *Docker Root Dir' | awk -F: '{ print $2 }' | sed 's/^ *//'`
-echo "Building this image may require up to 50GB of space for Docker"
-echo "Current space available in $DOCKER_DIR"
-df -h "$DOCKER_DIR"
-sleep 5
-
 HOST=docker.elastic.co
 ACCOUNT=ml-dev
 REPOSITORY=ml-linux-dependency-build
 VERSION=pytorch_latest
+FULL_IMAGE="$HOST/$ACCOUNT/$REPOSITORY:$VERSION"
+PYTORCH_BRANCH="${1:-viable/strict}"
 
 set -e
 
-cd `dirname $0`
+cd "$(dirname "$0")"
+
+# ---- Skip-if-unchanged check ----
+# Compare the current viable/strict HEAD with the commit baked into the
+# last published image.  If they match, there's nothing to rebuild.
+echo "--- Checking if PyTorch ${PYTORCH_BRANCH} has changed since last build"
+CURRENT_SHA=$(git ls-remote https://github.com/pytorch/pytorch.git "refs/heads/${PYTORCH_BRANCH}" 2>/dev/null | cut -f1)
+
+if [ -n "$CURRENT_SHA" ]; then
+    echo "Current ${PYTORCH_BRANCH} SHA: ${CURRENT_SHA}"
+
+    # Try to read the label from the existing image (pull first if not local)
+    docker pull "$FULL_IMAGE" 2>/dev/null || true
+    PREVIOUS_SHA=$(docker inspect --format '{{index .Config.Labels "pytorch.commit"}}' "$FULL_IMAGE" 2>/dev/null || echo "")
+
+    if [ -n "$PREVIOUS_SHA" ]; then
+        echo "Previous build SHA:           ${PREVIOUS_SHA}"
+        if [ "$CURRENT_SHA" = "$PREVIOUS_SHA" ]; then
+            echo "PyTorch ${PYTORCH_BRANCH} unchanged — skipping rebuild"
+            exit 0
+        fi
+        echo "SHA changed — rebuild needed"
+    else
+        echo "No previous build SHA found — full build required"
+    fi
+else
+    echo "WARNING: could not fetch ${PYTORCH_BRANCH} SHA, proceeding with build"
+fi
+
+# ---- Build ----
+DOCKER_DIR=$(docker info 2>/dev/null | grep '^ *Docker Root Dir' | awk -F: '{ print $2 }' | sed 's/^ *//')
+echo "Building this image may require up to 50GB of space for Docker"
+echo "Current space available in $DOCKER_DIR"
+df -h "$DOCKER_DIR"
 
 . ./prefetch_docker_image.sh
 CONTEXT=pytorch_linux_image
 prefetch_docker_base_image $CONTEXT/Dockerfile
-if [ $# -gt 0 ]; then
-  VERSION=pytorch_latest
-  echo "VERSION = $VERSION"
-  docker build --progress=plain --no-cache -t $HOST/$ACCOUNT/$REPOSITORY:$VERSION --build-arg pytorch_branch=viable/strict $CONTEXT
+
+echo "--- Building $FULL_IMAGE (branch: ${PYTORCH_BRANCH})"
+
+export DOCKER_BUILDKIT=1
+
+BUILD_ARGS="--build-arg pytorch_branch=${PYTORCH_BRANCH}"
+if [ -n "$CURRENT_SHA" ]; then
+    BUILD_ARGS="$BUILD_ARGS --build-arg pytorch_commit=${CURRENT_SHA}"
+fi
+
+# Pass GCS credentials as a Docker secret for sccache.
+# The key is injected by the Buildkite post-checkout hook into
+# SCCACHE_GCS_KEY_FILE (a temp file with the service account JSON).
+SECRET_ARGS=""
+if [ -n "${SCCACHE_GCS_KEY_FILE:-}" ] && [ -f "$SCCACHE_GCS_KEY_FILE" ]; then
+    SECRET_ARGS="--secret id=gcs_key,src=${SCCACHE_GCS_KEY_FILE}"
+    echo "sccache: GCS credentials available — incremental build enabled"
 else
-  echo "VERSION = $VERSION"
-  docker build --progress=plain --no-cache -t $HOST/$ACCOUNT/$REPOSITORY:$VERSION $CONTEXT
+    echo "sccache: no GCS credentials — full build (no cache)"
 fi
 
-# We use a special machine user account to authenticate to docker.elastic.co from within our Buildkite pipelines
-echo "Pushing $HOST/$ACCOUNT/$REPOSITORY:$VERSION"
+docker build \
+    --progress=plain \
+    $BUILD_ARGS \
+    $SECRET_ARGS \
+    -t "$FULL_IMAGE" \
+    "$CONTEXT"
+
+# ---- Push ----
+echo "--- Pushing $FULL_IMAGE"
 echo "$DOCKER_REGISTRY_PASSWORD" | docker login -u "$DOCKER_REGISTRY_USERNAME" --password-stdin docker.elastic.co
-docker push $HOST/$ACCOUNT/$REPOSITORY:$VERSION
+docker push "$FULL_IMAGE"
+
+echo "Build complete. PyTorch commit: ${CURRENT_SHA:-unknown}"

dev-tools/docker/pytorch_linux_image/Dockerfile

@@ -18,7 +18,13 @@ LABEL maintainer="Ed Savage <ed.savage@elastic.co>"
 RUN dnf -y update && \
     dnf install -y dnf-plugins-core && \
     dnf config-manager --set-enabled powertools && \
-    dnf install -y bzip2 gcc gcc-c++ git libffi-devel make texinfo unzip wget which xz zip zlib-devel findutils
+    dnf install -y bzip2 curl gcc gcc-c++ git libffi-devel make texinfo unzip wget which xz zip zlib-devel findutils
+
+# Install sccache for GCS-backed compilation caching across builds
+ARG SCCACHE_VERSION=v0.14.0
+RUN curl -fsSL "https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-x86_64-unknown-linux-musl.tar.gz" \
+    | tar xz -C /usr/local/bin --strip-components=1 "sccache-${SCCACHE_VERSION}-x86_64-unknown-linux-musl/sccache" && \
+    chmod +x /usr/local/bin/sccache
 
 # For compiling with hardening and optimisation
 ENV CFLAGS="-g -O3 -fstack-protector -D_FORTIFY_SOURCE=2 -msse4.2 -mfpmath=sse"
@@ -37,17 +43,24 @@ ENV CXX="g++ -std=gnu++17"
 # PYTORCH_BUILD_VERSION is only set for tagged branches (e.g. v2.7.1);
 # for main/viable/strict PyTorch derives the version from version.txt.
 ARG pytorch_branch=viable/strict
+ARG pytorch_commit=unknown
 
 ENV PYTORCH_BRANCH=${pytorch_branch}
 
+# Split the build into clone + build so sccache (GCS-backed) can accelerate
+# incremental rebuilds. The GCS key is passed via --mount=type=secret to
+# avoid baking credentials into the image.
 RUN \
   cd ${build_dir} && \
   git -c advice.detachedHead=false clone --depth=1 --branch=${PYTORCH_BRANCH} https://github.com/pytorch/pytorch.git && \
   cd pytorch && \
   git submodule sync && \
   git submodule update --init --recursive && \
   sed -i -e 's/system(/strlen(/' torch/csrc/jit/codegen/fuser/cpu/fused_kernel.cpp && \
-  sed -i -e '104 i set(PYTHON_EXECUTABLE "/usr/local/bin/python3.12")' ./third_party/onnx/CMakeLists.txt && \
+  sed -i -e '104 i set(PYTHON_EXECUTABLE "/usr/local/bin/python3.12")' ./third_party/onnx/CMakeLists.txt
+
+RUN --mount=type=secret,id=gcs_key \
+  cd ${build_dir}/pytorch && \
   if [[ "$PYTORCH_BRANCH" != "main" && "$PYTORCH_BRANCH" != "viable/strict" ]]; then export PYTORCH_BUILD_VERSION=$(expr "$PYTORCH_BRANCH" : 'v\(.*\)'); fi && \
   export BLAS=MKL && \
   export BUILD_TEST=OFF && \
@@ -60,8 +73,21 @@ RUN \
   export USE_XNNPACK=OFF && \
   export PYTORCH_BUILD_NUMBER=1 && \
   export MAX_JOBS=10 && \
+  if [ -f /run/secrets/gcs_key ]; then \
+    export SCCACHE_GCS_BUCKET=elastic-ml-cpp-sccache && \
+    export SCCACHE_GCS_KEY_PREFIX=pytorch-build && \
+    export SCCACHE_GCS_RW_MODE=READ_WRITE && \
+    export SCCACHE_GCS_KEY_PATH=/run/secrets/gcs_key && \
+    export CMAKE_C_COMPILER_LAUNCHER=sccache && \
+    export CMAKE_CXX_COMPILER_LAUNCHER=sccache && \
+    sccache --start-server && \
+    echo "sccache: using GCS backend (bucket=$SCCACHE_GCS_BUCKET prefix=$SCCACHE_GCS_KEY_PREFIX)"; \
+  else \
+    echo "sccache: no GCS credentials — building without cache"; \
+  fi && \
   /usr/local/bin/python3.12 -m pip install --trusted-host pypi.python.org --trusted-host files.pythonhosted.org --trusted-host pypi.org -r requirements.txt && \
   /usr/local/bin/python3.12 setup.py install && \
+  if command -v sccache &>/dev/null && sccache --show-stats &>/dev/null; then sccache --show-stats; sccache --stop-server; fi && \
   mkdir -p /usr/local/gcc133/include/pytorch && \
   /bin/cp -rf torch/include/* /usr/local/gcc133/include/pytorch/ && \
   /bin/cp -f torch/lib/libtorch_cpu.so /usr/local/gcc133/lib && \
@@ -71,8 +97,19 @@ RUN \
 
 FROM rockylinux:8
 COPY --from=builder /usr/local/gcc133 /usr/local/gcc133
+# Python 3.12 + torch site-packages for allowlist validation
+COPY --from=builder /usr/local/bin/python3.12 /usr/local/bin/python3.12
+COPY --from=builder /usr/local/bin/pip3.12 /usr/local/bin/pip3.12
+COPY --from=builder /usr/local/lib/python3.12 /usr/local/lib/python3.12
+COPY --from=builder /usr/local/lib/libpython3.12.a /usr/local/lib/libpython3.12.a
+RUN ln -sf /usr/local/bin/python3.12 /usr/local/bin/python3 && \
+    ln -sf /usr/local/bin/pip3.12 /usr/local/bin/pip3
+
+ARG pytorch_commit=unknown
+LABEL pytorch.commit=${pytorch_commit}
+
 RUN \
   dnf -y update && \
-  dnf install -y bzip2 gcc git make unzip which zip zlib-devel findutils && \
+  dnf install -y bzip2 gcc git make openssl-libs unzip which zip zlib-devel findutils && \
   dnf clean all && \
   rm -rf /var/cache/dnf /tmp/*

dev-tools/extract_model_ops/requirements.txt

@@ -1,4 +1,4 @@
 torch==2.7.1
-transformers>=4.40.0
+transformers>=4.40.0,<5.0.0
 sentencepiece>=0.2.0
 protobuf>=5.0.0

dev-tools/extract_model_ops/torchscript_utils.py

@@ -111,7 +111,7 @@ def load_and_trace_hf_model(model_name: str, quantize: bool = False,
 
     Returns the traced module, or None if the model could not be loaded or traced.
     """
-    token = os.environ.get("HF_TOKEN")
+    token = os.environ.get("HF_TOKEN") or None
     model_cls = _resolve_auto_class(auto_class)
     overrides = config_overrides or {}
 

dev-tools/extract_model_ops/validate_allowlist.py

@@ -16,6 +16,10 @@
 checks every operation against the ALLOWED_OPERATIONS and FORBIDDEN_OPERATIONS
 sets parsed from CSupportedOperations.cc.
 
+Each model download/trace is subject to a timeout (default 10 minutes,
+configurable via --model-timeout) to prevent stalled HuggingFace downloads
+from consuming the entire step timeout.
+
 This is the Python-side equivalent of the C++ CModelGraphValidator and is
 intended as an integration test: if any legitimate model produces an
 operation that the C++ code would reject, this script exits non-zero.
@@ -31,10 +35,21 @@
 import argparse
 import gc
 import re
+import signal
 import sys
 from pathlib import Path
 from typing import Optional
 
+MODEL_TIMEOUT_SECONDS = 600  # 10 minutes per model
+
+
+class ModelTimeoutError(Exception):
+    pass
+
+
+def _timeout_handler(signum, frame):
+    raise ModelTimeoutError("Model download/trace timed out")
+
 import torch
 
 from torchscript_utils import (
@@ -107,17 +122,32 @@ def validate_model(model_name: str,
                    verbose: bool,
                    quantize: bool = False,
                    auto_class: str | None = None,
-                   config_overrides: dict | None = None) -> str:
+                   config_overrides: dict | None = None,
+                   timeout: int = MODEL_TIMEOUT_SECONDS) -> str:
     """Validate one HuggingFace model.
 
     Returns "pass", "fail" (op validation failed), or "skip" (could not
-    load/trace — e.g. private model without HF_TOKEN).
+    load/trace — e.g. private model without HF_TOKEN, or download timeout).
     """
     label = f"{model_name} (quantized)" if quantize else model_name
     print(f"  {label}...", file=sys.stderr)
-    traced = load_and_trace_hf_model(model_name, quantize=quantize,
-                                     auto_class=auto_class,
-                                     config_overrides=config_overrides)
+
+    has_alarm = hasattr(signal, "SIGALRM")
+    if has_alarm:
+        old_handler = signal.signal(signal.SIGALRM, _timeout_handler)
+        signal.alarm(timeout)
+    try:
+        traced = load_and_trace_hf_model(model_name, quantize=quantize,
+                                         auto_class=auto_class,
+                                         config_overrides=config_overrides)
+    except ModelTimeoutError:
+        print(f"    SKIPPED (timed out after {timeout}s)", file=sys.stderr)
+        return "skip"
+    finally:
+        if has_alarm:
+            signal.alarm(0)
+            signal.signal(signal.SIGALRM, old_handler)
+
     if traced is None:
         print(f"    SKIPPED (could not load/trace)", file=sys.stderr)
         return "skip"
@@ -158,6 +188,9 @@ def main():
     parser.add_argument(
         "--verbose", action="store_true",
         help="Print per-model op counts")
+    parser.add_argument(
+        "--model-timeout", type=int, default=MODEL_TIMEOUT_SECONDS,
+        help=f"Timeout in seconds for each model download/trace (default: {MODEL_TIMEOUT_SECONDS})")
     args = parser.parse_args()
 
     print(f"PyTorch version: {torch.__version__}", file=sys.stderr)
@@ -178,7 +211,8 @@ def main():
             spec["model_id"], allowed, forbidden, args.verbose,
             quantize=spec["quantized"],
             auto_class=spec.get("auto_class"),
-            config_overrides=spec.get("config_overrides"))
+            config_overrides=spec.get("config_overrides"),
+            timeout=args.model_timeout)
 
     if args.pt_dir and args.pt_dir.is_dir():
         pt_files = sorted(args.pt_dir.glob("*.pt"))