Improve update_ci_for_release.sh (#1433)

Author: NicolasHug

packaging/update_ci_for_release.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 set -euo pipefail
 
-VERSION="${1:?Usage: $0 <release-version> (e.g. 0.12)}"
+BRANCH="${1:?Usage: $0 <test-infra-release-branch> (e.g. release/2.12)}"
 
 SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
 REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
@@ -10,8 +10,8 @@ WORKFLOW_DIR="${REPO_ROOT}/.github/workflows"
 # 1. Update test-infra @main refs and test-infra-ref on ALL workflow files
 for f in "${WORKFLOW_DIR}"/*.yaml; do
     echo "Updating test-infra refs in $f ..."
-    sed -i "s|pytorch/test-infra/\(.*\)@main|pytorch/test-infra/\1@release/${VERSION}|g" "$f"
-    sed -i "s|test-infra-ref: main|test-infra-ref: release/${VERSION}|g" "$f"
+    sed -i "s|pytorch/test-infra/\(.*\)@main|pytorch/test-infra/\1@${BRANCH}|g" "$f"
+    sed -i "s|test-infra-ref: main|test-infra-ref: ${BRANCH}|g" "$f"
 done
 
 # 2. Update python/ffmpeg/cuda versions in wheel install-and-test jobs only.
@@ -37,4 +37,4 @@ for f in "${WHEEL_FILES[@]}"; do
     sed -i "s|cuda-version: \[.*\]|cuda-version: ['12.6', '13.0', '13.2']|g" "$f"
 done
 
-echo "Done! Updated workflow files for release/${VERSION}"
+echo "Done! Updated workflow files to point to test-infra branch '${BRANCH}'"