fix v2 telemetry

junpuf · junpuf · commit 59635275eddb · 2025-12-03T18:11:21.000-08:00
Signed-off-by: Junpu Fan &lt;junpu@amazon.com&gt;
diff --git a/.github/workflows/pr-sglang.yml b/.github/workflows/pr-sglang.yml
@@ -13,6 +13,8 @@ permissions:
 
 env:
   # CI Image configuration
+  CONTAINER_TYPE: "general"
+  FRAMEWORK: "sglang"
   SGLANG_VERSION: "0.5.5"
   PYTHON_VERSION: "py312"
   CUDA_VERSION: "cu129"
@@ -116,6 +118,9 @@ jobs:
           docker buildx build --progress plain \
             --build-arg CACHE_REFRESH="$(date +"%Y-%m-%d")" \
             --build-arg BASE_IMAGE="lmsysorg/sglang:v${{ env.SGLANG_VERSION }}-${{ env.CUDA_VERSION }}-amd64" \
+            --build-arg CONTAINER_TYPE="${{ env.CONTAINER_TYPE }}" \
+            --build-arg FRAMEWORK="${{ env.FRAMEWORK }}" \
+            --build-arg FRAMEWORK_VERSION="${{ env.SGLANG_VERSION }}" \
             --cache-to=type=inline \
             --cache-from=type=registry,ref=${CI_IMAGE_URI} \
             --tag ${CI_IMAGE_URI} \
@@ -158,6 +163,39 @@ jobs:
           echo "AWS_ACCOUNT_ID=${AWS_ACCOUNT_ID}" >> ${GITHUB_OUTPUT}
           echo "IMAGE_URI=${IMAGE_URI}" >> ${GITHUB_OUTPUT}
 
+  sglang-telemetry-test:
+    needs: [set-test-environment, build-sglang-image]
+    if: success()
+    runs-on:
+      - codebuild-runner-${{ github.run_id }}-${{ github.run_attempt }}
+        buildspec-override:true
+    concurrency:
+      group: ${{ github.workflow }}-sglang-telemetry-test-${{ github.event.pull_request.number }}
+      cancel-in-progress: true
+    steps:
+      - name: Checkout DLC source
+        uses: actions/checkout@v5
+
+      - name: Container pull
+        uses: ./.github/actions/ecr-authenticate
+        with:
+          aws-account-id: ${{ needs.set-test-environment.outputs.aws-account-id }}
+          aws-region: ${{ vars.AWS_REGION }}
+          image-uri: ${{ needs.set-test-environment.outputs.image-uri }}
+
+      - name: Start container
+        run: |
+          CONTAINER_ID=$(docker run -d -it --rm --entrypoint /bin/bash \
+            -e TEST_MODE=1 \
+            ${{ needs.set-test-environment.outputs.image-uri }})
+          echo "CONTAINER_ID=${CONTAINER_ID}" >> ${GITHUB_ENV}
+          docker logs ${CONTAINER_ID}
+
+      - name: Run tests
+        run: |
+          docker exec ${CONTAINER_ID} cat /tmp/test_request.txt | grep "sglang"
+
+
   sglang-local-benchmark-test:
     needs: [set-test-environment, build-sglang-image]
     if: success()
diff --git a/.github/workflows/pr-vllm.yml b/.github/workflows/pr-vllm.yml
@@ -13,6 +13,8 @@ permissions:
 
 env:
   # CI Image configuration
+  CONTAINER_TYPE: "general"
+  FRAMEWORK: "vllm"
   VLLM_VERSION: 0.11.2
   VLLM_RAYSERVE_VERSION: 0.10.2
   PYTHON_VERSION: "py312"
@@ -118,6 +120,9 @@ jobs:
           docker buildx build --progress plain \
             --build-arg CACHE_REFRESH="$(date +"%Y-%m-%d")" \
             --build-arg BASE_IMAGE="vllm/vllm-openai:v${{ env.VLLM_VERSION }}" \
+            --build-arg CONTAINER_TYPE="${{ env.CONTAINER_TYPE }}" \
+            --build-arg FRAMEWORK="${{ env.FRAMEWORK }}" \
+            --build-arg FRAMEWORK_VERSION="${{ env.VLLM_VERSION }}" \
             --cache-to=type=inline \
             --cache-from=type=registry,ref=${CI_IMAGE_URI} \
             --tag ${CI_IMAGE_URI} \
@@ -160,6 +165,38 @@ jobs:
           echo "AWS_ACCOUNT_ID=${AWS_ACCOUNT_ID}" >> ${GITHUB_OUTPUT}
           echo "IMAGE_URI=${IMAGE_URI}" >> ${GITHUB_OUTPUT}
 
+  vllm-ec2-telemetry-test:
+    needs: [build-vllm-ec2-image, set-ec2-test-environment]
+    if: success()
+    runs-on:
+      - codebuild-runner-${{ github.run_id }}-${{ github.run_attempt }}
+        buildspec-override:true
+    concurrency:
+      group: ${{ github.workflow }}-vllm-ec2-telemetry-test-${{ github.event.pull_request.number }}
+      cancel-in-progress: true
+    steps:
+      - name: Checkout DLC source
+        uses: actions/checkout@v5
+
+      - name: Container pull
+        uses: ./.github/actions/ecr-authenticate
+        with:
+          aws-account-id: ${{ needs.set-ec2-test-environment.outputs.aws-account-id }}
+          aws-region: ${{ vars.AWS_REGION }}
+          image-uri: ${{ needs.set-ec2-test-environment.outputs.image-uri }}
+
+      - name: Start container
+        run: |
+          CONTAINER_ID=$(docker run -d -it --rm --entrypoint /bin/bash \
+            -e TEST_MODE=1 \
+            ${{ needs.set-ec2-test-environment.outputs.image-uri }})
+          echo "CONTAINER_ID=${CONTAINER_ID}" >> ${GITHUB_ENV}
+          docker logs ${CONTAINER_ID}
+
+      - name: Run tests
+        run: |
+          docker exec ${CONTAINER_ID} cat /tmp/test_request.txt | grep "${{ env.FRAMEWORK }}"
+
   vllm-ec2-regression-test:
     needs: [build-vllm-ec2-image, set-ec2-test-environment]
     if: success()
@@ -402,6 +439,9 @@ jobs:
           docker buildx build --progress plain \
             --build-arg CACHE_REFRESH="$(date +"%Y-%m-%d")" \
             --build-arg BASE_IMAGE="vllm/vllm-openai:v${{ env.VLLM_RAYSERVE_VERSION }}" \
+            --build-arg CONTAINER_TYPE="${{ env.CONTAINER_TYPE }}" \
+            --build-arg FRAMEWORK="${{ env.FRAMEWORK }}" \
+            --build-arg FRAMEWORK_VERSION="${{ env.VLLM_RAYSERVE_VERSION }}" \
             --cache-to=type=inline \
             --cache-from=type=registry,ref=${CI_IMAGE_URI} \
             --tag ${CI_IMAGE_URI} \
@@ -444,6 +484,38 @@ jobs:
           echo "AWS_ACCOUNT_ID=${AWS_ACCOUNT_ID}" >> ${GITHUB_OUTPUT}
           echo "IMAGE_URI=${IMAGE_URI}" >> ${GITHUB_OUTPUT}
 
+  vllm-rayserve-telemetry-test:
+    needs: [build-vllm-rayserve-image, set-rayserve-test-environment]
+    if: success()
+    runs-on:
+      - codebuild-runner-${{ github.run_id }}-${{ github.run_attempt }}
+        buildspec-override:true
+    concurrency:
+      group: ${{ github.workflow }}-vllm-rayserve-telemetry-test-${{ github.event.pull_request.number }}
+      cancel-in-progress: true
+    steps:
+      - name: Checkout DLC source
+        uses: actions/checkout@v5
+
+      - name: Container pull
+        uses: ./.github/actions/ecr-authenticate
+        with:
+          aws-account-id: ${{ needs.set-rayserve-test-environment.outputs.aws-account-id }}
+          aws-region: ${{ vars.AWS_REGION }}
+          image-uri: ${{ needs.set-rayserve-test-environment.outputs.image-uri }}
+
+      - name: Start container
+        run: |
+          CONTAINER_ID=$(docker run -d -it --rm --entrypoint /bin/bash \
+            -e TEST_MODE=1 \
+            ${{ needs.set-rayserve-test-environment.outputs.image-uri }})
+          echo "CONTAINER_ID=${CONTAINER_ID}" >> ${GITHUB_ENV}
+          docker logs ${CONTAINER_ID}
+
+      - name: Run tests
+        run: |
+          docker exec ${CONTAINER_ID} cat /tmp/test_request.txt | grep "${{ env.FRAMEWORK }}"
+
   vllm-rayserve-regression-test:
     needs: [build-vllm-rayserve-image, set-rayserve-test-environment]
     if: success()
@@ -686,6 +758,9 @@ jobs:
           docker buildx build --progress plain \
             --build-arg CACHE_REFRESH="$(date +"%Y-%m-%d")" \
             --build-arg BASE_IMAGE="vllm/vllm-openai:v${{ env.VLLM_VERSION }}" \
+            --build-arg CONTAINER_TYPE="${{ env.CONTAINER_TYPE }}" \
+            --build-arg FRAMEWORK="${{ env.FRAMEWORK }}" \
+            --build-arg FRAMEWORK_VERSION="${{ env.VLLM_VERSION }}" \
             --cache-to=type=inline \
             --cache-from=type=registry,ref=${CI_IMAGE_URI} \
             --tag ${CI_IMAGE_URI} \
@@ -728,6 +803,38 @@ jobs:
           echo "AWS_ACCOUNT_ID=${AWS_ACCOUNT_ID}" >> ${GITHUB_OUTPUT}
           echo "IMAGE_URI=${IMAGE_URI}" >> ${GITHUB_OUTPUT}
 
+  vllm-sagemaker-telemetry-test:
+    needs: [build-vllm-sagemaker-image, set-sagemaker-test-environment]
+    if: success()
+    runs-on:
+      - codebuild-runner-${{ github.run_id }}-${{ github.run_attempt }}
+        buildspec-override:true
+    concurrency:
+      group: ${{ github.workflow }}-vllm-sagemaker-telemetry-test-${{ github.event.pull_request.number }}
+      cancel-in-progress: true
+    steps:
+      - name: Checkout DLC source
+        uses: actions/checkout@v5
+
+      - name: Container pull
+        uses: ./.github/actions/ecr-authenticate
+        with:
+          aws-account-id: ${{ needs.set-sagemaker-test-environment.outputs.aws-account-id }}
+          aws-region: ${{ vars.AWS_REGION }}
+          image-uri: ${{ needs.set-sagemaker-test-environment.outputs.image-uri }}
+
+      - name: Start container
+        run: |
+          CONTAINER_ID=$(docker run -d -it --rm --entrypoint /bin/bash \
+            -e TEST_MODE=1 \
+            ${{ needs.set-sagemaker-test-environment.outputs.image-uri }})
+          echo "CONTAINER_ID=${CONTAINER_ID}" >> ${GITHUB_ENV}
+          docker logs ${CONTAINER_ID}
+
+      - name: Run tests
+        run: |
+          docker exec ${CONTAINER_ID} cat /tmp/test_request.txt | grep "${{ env.FRAMEWORK }}"
+
   vllm-sagemaker-regression-test:
     needs: [build-vllm-sagemaker-image, set-sagemaker-test-environment]
     if: success()
diff --git a/docker/sglang/Dockerfile b/docker/sglang/Dockerfile
@@ -1,5 +1,8 @@
 # Declare the argument as default to use as input
 # base image: https://hub.docker.com/r/lmsysorg/sglang/tags
+ARG FRAMEWORK=sglang
+ARG FRAMEWORK_VERSION=0.5.5
+ARG CONTAINER_TYPE=general
 ARG BASE_IMAGE=lmsysorg/sglang:v0.5.5-cu129-amd64
 
 # Use input argument as base image
@@ -32,13 +35,21 @@ WORKDIR /
 # Copy artifacts
 # ===============
 COPY ./scripts/telemetry/deep_learning_container.py /usr/local/bin/deep_learning_container.py
-COPY ./scripts/telemetry/bash_telemetry.sh /usr/local/bin/bash_telemetry.sh
+COPY ./scripts/telemetry/bash_telemetry.sh.template /tmp/bash_telemetry.sh.template
 COPY ./scripts/common/install_efa.sh install_efa.sh
 COPY ./scripts/common/start_cuda_compat.sh /usr/local/bin/start_cuda_compat.sh
 
+ARG FRAMEWORK
+ARG FRAMEWORK_VERSION
+ARG CONTAINER_TYPE
 RUN chmod +x /usr/local/bin/deep_learning_container.py \
+  && sed -e "s/{{FRAMEWORK}}/${FRAMEWORK}/g" \
+    -e "s/{{FRAMEWORK_VERSION}}/${FRAMEWORK_VERSION}/g" \
+    -e "s/{{CONTAINER_TYPE}}/${CONTAINER_TYPE}/g" \
+    /tmp/bash_telemetry.sh.template >/usr/local/bin/bash_telemetry.sh \
   && chmod +x /usr/local/bin/bash_telemetry.sh \
-  && chmod +x /usr/local/bin/start_cuda_compat.sh
+  && chmod +x /usr/local/bin/start_cuda_compat.sh \
+  && rm /tmp/bash_telemetry.sh.template
 
 # Install EFA and remove vulnerable nvjpeg
 # =========================================
diff --git a/docker/vllm/Dockerfile b/docker/vllm/Dockerfile
@@ -1,5 +1,8 @@
 # Declare the argument as default to use as input
 # base image: https://hub.docker.com/r/vllm/vllm-openai/tags
+ARG FRAMEWORK=vllm
+ARG FRAMEWORK_VERSION=0.11.0
+ARG CONTAINER_TYPE=general
 ARG BASE_IMAGE=vllm/vllm-openai:v0.11.0
 
 # Use input argument as base image
@@ -25,11 +28,19 @@ ENV DEBIAN_FRONTEND=noninteractive \
 WORKDIR /
 
 COPY ./scripts/telemetry/deep_learning_container.py /usr/local/bin/deep_learning_container.py
-COPY ./scripts/telemetry/bash_telemetry.sh /usr/local/bin/bash_telemetry.sh
+COPY ./scripts/telemetry/bash_telemetry.sh.template /tmp/bash_telemetry.sh.template
 COPY ./scripts/common/setup_oss_compliance.sh setup_oss_compliance.sh
 
+ARG FRAMEWORK
+ARG FRAMEWORK_VERSION
+ARG CONTAINER_TYPE
 RUN chmod +x /usr/local/bin/deep_learning_container.py \
+  && sed -e "s/{{FRAMEWORK}}/${FRAMEWORK}/g" \
+    -e "s/{{FRAMEWORK_VERSION}}/${FRAMEWORK_VERSION}/g" \
+    -e "s/{{CONTAINER_TYPE}}/${CONTAINER_TYPE}/g" \
+    /tmp/bash_telemetry.sh.template >/usr/local/bin/bash_telemetry.sh \
   && chmod +x /usr/local/bin/bash_telemetry.sh \
+  && rm /tmp/bash_telemetry.sh.template \
   && echo 'source /usr/local/bin/bash_telemetry.sh' >>/etc/bash.bashrc \
   && bash setup_oss_compliance.sh ${PYTHON} && rm setup_oss_compliance.sh \
   # create symlink for python
diff --git a/scripts/telemetry/bash_telemetry.sh.template b/scripts/telemetry/bash_telemetry.sh.template
@@ -3,9 +3,9 @@
 if [ -f /usr/local/bin/deep_learning_container.py ] && [[ -z "${OPT_OUT_TRACKING}" || "${OPT_OUT_TRACKING,,}" != "true" ]]; then
     (
         python /usr/local/bin/deep_learning_container.py \
-            --framework "${FRAMEWORK}" \
-            --framework-version "${FRAMEWORK_VERSION}" \
-            --container-type "${CONTAINER_TYPE}" \
+            --framework "{{FRAMEWORK}}" \
+            --framework-version "{{FRAMEWORK_VERSION}}" \
+            --container-type "{{CONTAINER_TYPE}}" \
             &>/dev/null &
     )
-fi
+fi
diff --git a/scripts/telemetry/deep_learning_container.py b/scripts/telemetry/deep_learning_container.py
@@ -240,7 +240,7 @@ def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--framework",
-        choices=["tensorflow", "mxnet", "pytorch", "base", "vllm"],
+        choices=["tensorflow", "mxnet", "pytorch", "base", "vllm", "sglang"],
         help="framework of container image.",
         required=True,
     )

Original file line number	Diff line number	Diff line change
`@@ -3,9 +3,9 @@`
`3`	`3`	`if [ -f /usr/local/bin/deep_learning_container.py ] && [[ -z "${OPT_OUT_TRACKING}" \|\| "${OPT_OUT_TRACKING,,}" != "true" ]]; then`
`4`	`4`	`(`
`5`	`5`	`python /usr/local/bin/deep_learning_container.py \`
`6`		`- --framework "${FRAMEWORK}" \`
`7`		`- --framework-version "${FRAMEWORK_VERSION}" \`
`8`		`- --container-type "${CONTAINER_TYPE}" \`
	`6`	`+ --framework "{{FRAMEWORK}}" \`
	`7`	`+ --framework-version "{{FRAMEWORK_VERSION}}" \`
	`8`	`+ --container-type "{{CONTAINER_TYPE}}" \`
`9`	`9`	`&>/dev/null &`
`10`	`10`	`)`
`11`		`-fi`
	`11`	`+fi`
Original file line number	Diff line number	Diff line change
`@@ -240,7 +240,7 @@ def parse_args():`
`240`	`240`	`parser = argparse.ArgumentParser()`
`241`	`241`	`parser.add_argument(`
`242`	`242`	`"--framework",`
`243`		`- choices=["tensorflow", "mxnet", "pytorch", "base", "vllm"],`
	`243`	`+ choices=["tensorflow", "mxnet", "pytorch", "base", "vllm", "sglang"],`
`244`	`244`	`help="framework of container image.",`
`245`	`245`	`required=True,`
`246`	`246`	`)`