refactor: version PyTorch directory structure for multi-version support

.github/scripts/efa/ec2_helpers.py

@@ -386,6 +386,44 @@ def release_eip(aws_session, alloc_id):
         LOGGER.warning(f"Failed to release EIP {alloc_id}: {e}")
 
 
+def cleanup_stale_efa_instances(aws_session, max_age_hours=4):
+    """Terminate EFA test instances older than max_age_hours and release their EIPs.
+
+    Prevents resource leaks from cancelled/crashed workflow runs that didn't reach cleanup.
+    """
+    from datetime import datetime, timezone
+
+    cutoff = datetime.now(timezone.utc).timestamp() - (max_age_hours * 3600)
+
+    try:
+        resp = aws_session.ec2.describe_instances(
+            Filters=[
+                {"Name": "tag:Name", "Values": ["CI-CD EFA efa-test"]},
+                {"Name": "instance-state-name", "Values": ["running", "stopped"]},
+            ]
+        )
+        for reservation in resp.get("Reservations", []):
+            for instance in reservation.get("Instances", []):
+                launch_time = instance["LaunchTime"].timestamp()
+                if launch_time < cutoff:
+                    instance_id = instance["InstanceId"]
+                    LOGGER.warning(
+                        f"Terminating stale EFA instance {instance_id} (launched {instance['LaunchTime']})"
+                    )
+                    aws_session.ec2.terminate_instances(InstanceIds=[instance_id])
+
+        # Release unassociated EIPs (leaked from terminated instances)
+        addresses = aws_session.ec2.describe_addresses().get("Addresses", [])
+        for addr in addresses:
+            if not addr.get("AssociationId") and addr.get("AllocationId"):
+                LOGGER.warning(
+                    f"Releasing orphaned EIP {addr['AllocationId']} ({addr.get('PublicIp', '')})"
+                )
+                release_eip(aws_session, addr["AllocationId"])
+    except Exception as e:
+        LOGGER.warning(f"Stale resource cleanup failed (non-fatal): {e}")
+
+
 @contextmanager
 def efa_instances(image_uri, instance_type="p4d.24xlarge", region=DEFAULT_REGION):
     """Context manager that launches 2 EFA instances, sets up containers + SSH, and cleans up.
@@ -396,6 +434,9 @@ def efa_instances(image_uri, instance_type="p4d.24xlarge", region=DEFAULT_REGION
     ami_id = aws_session.get_latest_ami()
     sg_id = get_efa_security_group_id(aws_session)
 
+    # Clean up leaked resources from previous cancelled/crashed runs
+    cleanup_stale_efa_instances(aws_session)
+
     key_name = None
     key_path = None
     runner_ip = None

.github/workflows/autorelease-pytorch-ec2-cpu.yml

@@ -2,9 +2,14 @@ name: Auto Release - PyTorch EC2 CPU
 
 on:
   schedule:
-    - cron: '00 17 * * 1,3'
-
+    - cron: '00 17 * * 1,3'   # PyTorch 2.11 — Mon/Wed 9:00 AM PST
+    # - cron: '10 17 * * 1,3' # PyTorch 2.12 — Mon/Wed 9:10 AM PST (future)
   workflow_dispatch:
+    inputs:
+      config-file:
+        description: "Config file path (e.g., .github/config/image/pytorch-2.11-ec2-cpu.yml)"
+        required: true
+        type: string
 
 concurrency:
   group: ${{ github.workflow }}
@@ -15,10 +20,29 @@ permissions:
 
 env:
   FORCE_COLOR: "1"
-  CONFIG_FILE: ".github/config/image/pytorch-ec2-cpu.yml"
 
 jobs:
+  determine-config:
+    runs-on: ubuntu-latest
+    outputs:
+      config-file: ${{ steps.config.outputs.config-file }}
+    steps:
+      - name: Determine config file
+        id: config
+        run: |
+          if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then
+            echo "config-file=${{ inputs.config-file }}" >> $GITHUB_OUTPUT
+          else
+            CRON="${{ github.event.schedule }}"
+            case "$CRON" in
+              "00 17 * * 1,3") echo "config-file=.github/config/image/pytorch-2.11-ec2-cpu.yml" >> $GITHUB_OUTPUT ;;
+              # "10 17 * * 1,3") echo "config-file=.github/config/image/pytorch-2.12-ec2-cpu.yml" >> $GITHUB_OUTPUT ;;
+              *) echo "::error::Unknown cron: $CRON"; exit 1 ;;
+            esac
+          fi
+
   load-config:
+    needs: [determine-config]
     runs-on: ubuntu-latest
     outputs:
       config: ${{ steps.load.outputs.config }}
@@ -41,7 +65,7 @@ jobs:
         id: load
         uses: ./.github/actions/load-config
         with:
-          config-file: ${{ env.CONFIG_FILE }}
+          config-file: ${{ needs.determine-config.outputs.config-file }}
 
       - name: Parse configuration
         id: parse
@@ -83,10 +107,12 @@ jobs:
       - name: Build runtime image
         id: build-runtime
         run: |
-          source docker/pytorch/versions-cpu.env
+          VERSION=$(echo "${{ needs.load-config.outputs.framework-version }}" | cut -d. -f1,2)
+          source docker/pytorch/${VERSION}/versions-cpu.env
           CI_IMAGE_URI="${{ vars.CI_AWS_ACCOUNT_ID }}.dkr.ecr.${{ vars.AWS_REGION }}.amazonaws.com/ci:pytorch-cpu-runtime-${{ needs.load-config.outputs.framework-version }}-${{ needs.load-config.outputs.python-version }}-${{ needs.load-config.outputs.os-version }}-ec2-${{ github.run_id }}"
 
           docker buildx build --progress plain \
+            --build-arg DLC_PYTORCH_VERSION=${VERSION} \
             --build-arg FRAMEWORK=${{ needs.load-config.outputs.framework }} \
             --build-arg PYTHON_VERSION=${PYTHON_VERSION} \
             --build-arg TORCH_VERSION=${TORCH_VERSION} \
@@ -98,7 +124,7 @@ jobs:
             --tag ${CI_IMAGE_URI} \
             --push \
             --target runtime \
-            -f docker/pytorch/Dockerfile.cpu .
+            -f docker/pytorch/${VERSION}/Dockerfile.cpu .
 
           echo "image-uri=${CI_IMAGE_URI}" >> $GITHUB_OUTPUT
 
@@ -142,7 +168,7 @@ jobs:
       container-type: ${{ needs.load-config.outputs.container-type }}
 
   unit-test:
-    needs: [build-image]
+    needs: [load-config, build-image]
     runs-on:
       - codebuild-runner-${{ github.run_id }}-${{ github.run_attempt }}
         fleet:default-runner
@@ -159,10 +185,11 @@ jobs:
 
       - name: Run unit tests
         run: |
+          VERSION=$(echo "${{ needs.load-config.outputs.framework-version }}" | cut -d. -f1,2)
           IMAGE="${{ needs.build-image.outputs.ci-image }}"
           docker pull ${IMAGE}
           CONTAINER_ID=$(docker run -d --rm --entrypoint /bin/bash \
-            -e DLC_WORKDIR=/workdir \
+            -e DLC_WORKDIR=/workdir -e DLC_PYTORCH_VERSION=${VERSION} \
             -v $(pwd):/workdir --workdir /workdir \
             ${IMAGE} -c 'sleep infinity')
           docker exec ${CONTAINER_ID} pip install pytest -q

.github/workflows/autorelease-pytorch-ec2-cuda.yml

@@ -2,10 +2,14 @@ name: Auto Release - PyTorch EC2 CUDA
 
 on:
   schedule:
-    # Runs at 9AM/10AM PST/PDT on Mondays and Wednesdays
-    - cron: '00 17 * * 1,3'
-
+    - cron: '00 17 * * 1,3'   # PyTorch 2.11 — Mon/Wed 9:00 AM PST
+    # - cron: '10 17 * * 1,3' # PyTorch 2.12 — Mon/Wed 9:10 AM PST (future)
   workflow_dispatch:
+    inputs:
+      config-file:
+        description: "Config file path (e.g., .github/config/image/pytorch-2.11-ec2-cuda.yml)"
+        required: true
+        type: string
 
 concurrency:
   group: ${{ github.workflow }}
@@ -16,10 +20,29 @@ permissions:
 
 env:
   FORCE_COLOR: "1"
-  CONFIG_FILE: ".github/config/image/pytorch-ec2-cuda.yml"
 
 jobs:
+  determine-config:
+    runs-on: ubuntu-latest
+    outputs:
+      config-file: ${{ steps.config.outputs.config-file }}
+    steps:
+      - name: Determine config file
+        id: config
+        run: |
+          if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then
+            echo "config-file=${{ inputs.config-file }}" >> $GITHUB_OUTPUT
+          else
+            CRON="${{ github.event.schedule }}"
+            case "$CRON" in
+              "00 17 * * 1,3") echo "config-file=.github/config/image/pytorch-2.11-ec2-cuda.yml" >> $GITHUB_OUTPUT ;;
+              # "10 17 * * 1,3") echo "config-file=.github/config/image/pytorch-2.12-ec2-cuda.yml" >> $GITHUB_OUTPUT ;;
+              *) echo "::error::Unknown cron: $CRON"; exit 1 ;;
+            esac
+          fi
+
   load-config:
+    needs: [determine-config]
     runs-on: ubuntu-latest
     outputs:
       config: ${{ steps.load.outputs.config }}
@@ -42,7 +65,7 @@ jobs:
         id: load
         uses: ./.github/actions/load-config
         with:
-          config-file: ${{ env.CONFIG_FILE }}
+          config-file: ${{ needs.determine-config.outputs.config-file }}
 
       - name: Parse configuration
         id: parse
@@ -85,17 +108,19 @@ jobs:
       - name: Source versions
         id: versions
         run: |
-          source docker/pytorch/versions-cuda.env
+          VERSION=$(echo "${{ needs.load-config.outputs.framework-version }}" | cut -d. -f1,2)
+          source docker/pytorch/${VERSION}/versions-cuda.env
           echo "torch-version=${TORCH_VERSION}" >> $GITHUB_OUTPUT
           echo "cuda-version=${CUDA_VERSION}" >> $GITHUB_OUTPUT
           echo "python-version=${PYTHON_VERSION}" >> $GITHUB_OUTPUT
 
       - name: Fetch cached wheels
         run: |
-          source docker/pytorch/versions-cuda.env
-          mkdir -p docker/pytorch/wheels
+          VERSION=$(echo "${{ needs.load-config.outputs.framework-version }}" | cut -d. -f1,2)
+          source docker/pytorch/${VERSION}/versions-cuda.env
+          mkdir -p docker/pytorch/${VERSION}/wheels
           bash scripts/pytorch/fetch_cached_wheels.sh \
-            docker/pytorch/wheels \
+            docker/pytorch/${VERSION}/wheels \
             "${{ vars.WHEEL_CACHE_BUCKET }}" \
             "${CUDA_VERSION}" "${TORCH_VERSION}" "${PYTHON_VERSION}" \
             "flash-attn:${FLASH_ATTN_VERSION}" \
@@ -105,10 +130,12 @@ jobs:
       - name: Build runtime image
         id: build-runtime
         run: |
-          source docker/pytorch/versions-cuda.env
+          VERSION=$(echo "${{ needs.load-config.outputs.framework-version }}" | cut -d. -f1,2)
+          source docker/pytorch/${VERSION}/versions-cuda.env
           CI_IMAGE_URI="${{ vars.CI_AWS_ACCOUNT_ID }}.dkr.ecr.${{ vars.AWS_REGION }}.amazonaws.com/ci:pytorch-runtime-${{ needs.load-config.outputs.framework-version }}-gpu-${{ needs.load-config.outputs.python-version }}-${{ needs.load-config.outputs.cuda-version }}-${{ needs.load-config.outputs.os-version }}-ec2-${{ github.run_id }}"
 
           docker buildx build --progress plain \
+            --build-arg DLC_PYTORCH_VERSION=${VERSION} \
             --build-arg FRAMEWORK=${{ needs.load-config.outputs.framework }} \
             --build-arg CUDA_VERSION=${CUDA_VERSION} \
             --build-arg PYTHON_VERSION=${PYTHON_VERSION} \
@@ -125,17 +152,19 @@ jobs:
             --tag ${CI_IMAGE_URI} \
             --push \
             --target runtime \
-            -f docker/pytorch/Dockerfile.cuda .
+            -f docker/pytorch/${VERSION}/Dockerfile.cuda .
 
           echo "image-uri=${CI_IMAGE_URI}" >> $GITHUB_OUTPUT
 
       - name: Upload built wheels to cache
         run: |
-          source docker/pytorch/versions-cuda.env
+          VERSION=$(echo "${{ needs.load-config.outputs.framework-version }}" | cut -d. -f1,2)
+          source docker/pytorch/${VERSION}/versions-cuda.env
           bash scripts/pytorch/upload_cached_wheels.sh \
             "${{ vars.WHEEL_CACHE_BUCKET }}" \
             "${CUDA_VERSION}" "${TORCH_VERSION}" "${PYTHON_VERSION}" \
             "${{ steps.build-runtime.outputs.image-uri }}" \
+            "docker/pytorch/${VERSION}/Dockerfile.cuda" \
             "flash-attn:${FLASH_ATTN_VERSION}" \
             "transformer-engine-torch:${TRANSFORMER_ENGINE_VERSION}" \
         continue-on-error: true
@@ -180,7 +209,7 @@ jobs:
       container-type: ${{ needs.load-config.outputs.container-type }}
 
   unit-test:
-    needs: [build-image]
+    needs: [load-config, build-image]
     runs-on:
       - codebuild-runner-${{ github.run_id }}-${{ github.run_attempt }}
         fleet:default-runner
@@ -197,10 +226,11 @@ jobs:
 
       - name: Run unit tests
         run: |
+          VERSION=$(echo "${{ needs.load-config.outputs.framework-version }}" | cut -d. -f1,2)
           IMAGE="${{ needs.build-image.outputs.ci-image }}"
           docker pull ${IMAGE}
           CONTAINER_ID=$(docker run -d --rm --entrypoint /bin/bash \
-            -e DLC_WORKDIR=/workdir \
+            -e DLC_WORKDIR=/workdir -e DLC_PYTORCH_VERSION=${VERSION} \
             -v $(pwd):/workdir --workdir /workdir \
             ${IMAGE} -c 'sleep infinity')
           docker exec ${CONTAINER_ID} pip install pytest -q