fix(pytorch): opt 2.13 GPU configs out of wheel-cache upload

Add a per-config `build.skip_wheel_cache_upload` toggle and enable it on the
two 2.13 GPU configs. The wheel-cache upload path in
scripts/ci/build/pytorch_runtime/lib/upload_wheels.sh runs
`docker buildx build --target wheel-export` which re-executes the entire
Dockerfile from scratch to reach a scratch stage that COPYs the built wheels
out of the flash-attn and TE builder stages. On 2.13 that concurrent recompile
(flash-attn 2.8.3 source build + transformer-engine 2.16.1) exceeds the CI
runner cgroup memory limit and gets OOM-killed with SIGKILL / exit code 137,
before post_build.sh's `|| true` can catch it.

Evidence from PR CI attempt 2 (both GPU builds, same signature):

  upload_wheels.sh: line 45: NNN Killed  docker buildx build \
    --progress=plain --target wheel-export \
    --output type=local,dest=${EXPORT_DIR} -f ${DOCKERFILE} .
  ##[error]Process completed with exit code 137.

The primary image build+push succeeded before this step in both jobs; the
wheel cache is a build-time optimization that has no active consumer for 2.13
yet (WHEEL_CACHE_HIT was already false). Skipping keeps GPU builds green until
a larger CI runner class is provisioned.

- resolve_build_args.py: add skip_wheel_cache_upload to RESERVED_KEYS so it
  doesn't leak into Docker as a --build-arg.
- post_build.sh: short-circuit before invoking upload_wheels.sh when the flag
  is set.
- 2.13-{ec2,sagemaker}-cuda.yml: set build.skip_wheel_cache_upload: "true"
  with an inline comment explaining the opt-out. CPU configs unaffected.

Author: Bhanu Teja Goshikonda

.github/actions/build-image/resolve_build_args.py

@@ -24,7 +24,7 @@
 import subprocess
 import sys
 
-RESERVED_KEYS = {"dockerfile", "target"}
+RESERVED_KEYS = {"dockerfile", "target", "skip_wheel_cache_upload"}
 
 
 def load_yaml(path):

.github/config/image/pytorch/2.13-ec2-cuda.yml

@@ -29,6 +29,10 @@ build:
   max_jobs: "8"
   dlc_major_version: "1"
   dlc_minor_version: "0"
+  # wheel-export re-runs the CUDA-devel builder stages and OOM-kills on the CI
+  # runner with flash-attn + transformer-engine 2.16 wheels. Opt out until a
+  # larger runner class is available; the primary image still builds/pushes.
+  skip_wheel_cache_upload: "true"
 
 release:
   release: true

.github/config/image/pytorch/2.13-sagemaker-cuda.yml

@@ -30,6 +30,10 @@ build:
   max_jobs: "8"
   dlc_major_version: "1"
   dlc_minor_version: "0"
+  # wheel-export re-runs the CUDA-devel builder stages and OOM-kills on the CI
+  # runner with flash-attn + transformer-engine 2.16 wheels. Opt out until a
+  # larger runner class is available; the primary image still builds/pushes.
+  skip_wheel_cache_upload: "true"
 
 release:
   release: true

scripts/ci/build/pytorch_runtime/post_build.sh

@@ -33,6 +33,21 @@ done
 [[ -n "$CONFIG_FILE" ]] || { echo "ERROR: --config-file is required" >&2; exit 1; }
 [[ -f "$CONFIG_FILE" ]] || { echo "ERROR: Config file not found: $CONFIG_FILE" >&2; exit 1; }
 
+# Per-config opt-out. When `build.skip_wheel_cache_upload: true` is set on an
+# image config, skip lib/upload_wheels.sh entirely. Rationale: upload_wheels.sh
+# invokes `docker buildx build --target wheel-export` which re-executes the
+# CUDA-devel builder stages from scratch to produce a scratch-image containing
+# built wheels. On heavy configs (flash-attn + transformer-engine >=2.16), the
+# concurrent compile drives the runner over its cgroup memory limit and gets
+# OOM-killed (SIGKILL, exit 137). The wheel cache is a build-time optimization
+# — skipping it means future PR builds cannot reuse wheels from this run, but
+# the primary image already built and pushed successfully.
+SKIP_WHEEL_CACHE_UPLOAD=$(yq '.build.skip_wheel_cache_upload // "false"' "$CONFIG_FILE")
+if [[ "${SKIP_WHEEL_CACHE_UPLOAD}" == "true" ]]; then
+  echo "build.skip_wheel_cache_upload=true in ${CONFIG_FILE} — skipping wheel-cache upload"
+  exit 0
+fi
+
 CUDA_VERSION=$(yq '.build.cuda_version' "$CONFIG_FILE")
 FLASH_ATTN_VERSION=$(yq '.build.flash_attn_version // ""' "$CONFIG_FILE")
 TRANSFORMER_ENGINE_VERSION=$(yq '.build.transformer_engine_version // ""' "$CONFIG_FILE")