fix test setup

Signed-off-by: Yadan Wei <yadanwei@amazon.com>

Author: Yadan Wei

.github/workflows/pr-vllm.yml

@@ -160,6 +160,7 @@ jobs:
     outputs:
       aws-account-id: ${{ steps.set-env.outputs.AWS_ACCOUNT_ID }}
       image-uri: ${{ steps.set-env.outputs.IMAGE_URI }}
+      framework-version: ${{ steps.config.outputs.framework-version }}
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -211,7 +212,7 @@ jobs:
         uses: actions/checkout@v5
         with:
           repository: vllm-project/vllm
-          ref: v${{ env.VLLM_VERSION }}
+          ref: v${{ needs.set-ec2-test-environment.outputs.framework-version }}
           path: vllm_source
 
       - name: Start container
@@ -257,7 +258,7 @@ jobs:
         uses: actions/checkout@v5
         with:
           repository: vllm-project/vllm
-          ref: v${{ env.VLLM_VERSION }}
+          ref: v${{ needs.set-ec2-test-environment.outputs.framework-version }}
           path: vllm_source
 
       - name: Start container
@@ -303,7 +304,7 @@ jobs:
         uses: actions/checkout@v5
         with:
           repository: vllm-project/vllm
-          ref: v${{ env.VLLM_VERSION }}
+          ref: v${{ needs.set-ec2-test-environment.outputs.framework-version }}
           path: vllm_source
 
       - name: Start container
@@ -391,6 +392,7 @@ jobs:
     outputs:
       aws-account-id: ${{ steps.set-env.outputs.AWS_ACCOUNT_ID }}
       image-uri: ${{ steps.set-env.outputs.IMAGE_URI }}
+      framework-version: ${{ steps.config.outputs.framework-version }}
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -399,6 +401,7 @@ jobs:
         id: config
         run: |
           echo '${{ needs.load-config.outputs.rayserve-config }}' > config.json
+          echo "framework-version=$(jq -r '.common.framework_version' config.json)" >> $GITHUB_OUTPUT
           echo "prod-image=$(jq -r '.common.prod_image' config.json)" >> $GITHUB_OUTPUT
 
       - name: Set test environment
@@ -441,7 +444,7 @@ jobs:
         uses: actions/checkout@v5
         with:
           repository: vllm-project/vllm
-          ref: v${{ env.VLLM_RAYSERVE_VERSION }}
+          ref: v${{ needs.set-rayserve-test-environment.outputs.framework-version }}
           path: vllm_source
 
       - name: Start container
@@ -487,7 +490,7 @@ jobs:
         uses: actions/checkout@v5
         with:
           repository: vllm-project/vllm
-          ref: v${{ env.VLLM_RAYSERVE_VERSION }}
+          ref: v${{ needs.set-rayserve-test-environment.outputs.framework-version }}
           path: vllm_source
 
       - name: Start container
@@ -533,7 +536,7 @@ jobs:
         uses: actions/checkout@v5
         with:
           repository: vllm-project/vllm
-          ref: v${{ env.VLLM_RAYSERVE_VERSION }}
+          ref: v${{ needs.set-rayserve-test-environment.outputs.framework-version }}
           path: vllm_source
 
       - name: Start container
@@ -609,7 +612,7 @@ jobs:
           customer-type: ${{ steps.config.outputs.customer-type }}
 
   set-sagemaker-test-environment:
-    needs: [check-changes, build-vllm-sagemaker-image]
+    needs: [check-changes, build-vllm-sagemaker-image, load-config]
     if: |
       always() && !failure() && !cancelled() &&
       (needs.check-changes.outputs.build-change == 'true' || needs.check-changes.outputs.test-change == 'true')
@@ -620,10 +623,18 @@ jobs:
     outputs:
       aws-account-id: ${{ steps.set-env.outputs.AWS_ACCOUNT_ID }}
       image-uri: ${{ steps.set-env.outputs.IMAGE_URI }}
+      framework-version: ${{ steps.config.outputs.framework-version }}
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
 
+      - name: Parse SageMaker config
+        id: config
+        run: |
+          echo '${{ needs.load-config.outputs.sagemaker-config }}' > config.json
+          echo "framework-version=$(jq -r '.common.framework_version' config.json)" >> $GITHUB_OUTPUT
+          echo "prod-image=$(jq -r '.common.prod_image' config.json)" >> $GITHUB_OUTPUT
+
       - name: Set test environment
         id: set-env
         run: |
@@ -632,7 +643,7 @@ jobs:
             IMAGE_URI=${{ needs.build-vllm-sagemaker-image.outputs.ci-image }}
           else
             AWS_ACCOUNT_ID=${{ vars.PROD_AWS_ACCOUNT_ID }}
-            IMAGE_URI=${{ vars.PROD_AWS_ACCOUNT_ID }}.dkr.ecr.${{ vars.AWS_REGION }}.amazonaws.com/${{ env.PROD_SAGEMAKER_IMAGE }}
+            IMAGE_URI=${{ vars.PROD_AWS_ACCOUNT_ID }}.dkr.ecr.${{ vars.AWS_REGION }}.amazonaws.com/${{ steps.config.outputs.prod-image }}
           fi
 
           echo "Image URI to test: ${IMAGE_URI}"
@@ -664,7 +675,7 @@ jobs:
         uses: actions/checkout@v5
         with:
           repository: vllm-project/vllm
-          ref: v${{ env.VLLM_VERSION }}
+          ref: v${{ needs.set-sagemaker-test-environment.outputs.framework-version }}
           path: vllm_source
 
       - name: Start container
@@ -710,7 +721,7 @@ jobs:
         uses: actions/checkout@v5
         with:
           repository: vllm-project/vllm
-          ref: v${{ env.VLLM_VERSION }}
+          ref: v${{ needs.set-sagemaker-test-environment.outputs.framework-version }}
           path: vllm_source
 
       - name: Start container
@@ -756,7 +767,7 @@ jobs:
         uses: actions/checkout@v5
         with:
           repository: vllm-project/vllm
-          ref: v${{ env.VLLM_VERSION }}
+          ref: v${{ needs.set-sagemaker-test-environment.outputs.framework-version }}
           path: vllm_source
 
       - name: Start container