Fix the workflowId, jobId and attemp for vllm benchmark and sglang

.github/scripts/upload_benchmark_results.py

@@ -151,18 +151,23 @@ def get_git_metadata(repo_dir: str) -> Tuple[str, str]:
         # We couldn't find the branch name
         return repo_name, "main", hexsha, committed_date
 
-
 def get_benchmark_metadata(
     repo_name: str, head_branch: str, head_sha: str, timestamp: int, benchmark_name
 ) -> Dict[str, Any]:
+
+    # combine run branch name with the foreign branch name if it's not main
+    local_branch = os.getenv("RUN_LOCAL_BRANCH","")
+    branch_name = head_branch
+    if local_branch and local_branch != "main":
+        branch_name = f"{head_branch}({local_branch})"
     return {
         "timestamp": timestamp,
         "schema_version": "v3",
         "name": benchmark_name,
         "repo": repo_name,
-        "head_branch": head_branch,
+        "head_branch": branch_name,
         "head_sha": head_sha,
-        "workflow_id": os.getenv("WORKFLOW_ID", timestamp),
+        "workflow_id": os.getenv("WORKFLOW_RUN_ID", timestamp),
         "run_attempt": os.getenv("RUN_ATTEMPT", 1),
         "job_id": os.getenv("JOB_ID", timestamp),
     }
@@ -325,7 +330,6 @@ def upload(
 
 def main() -> None:
     args = parse_args()
-
     if args.repo:
         if args.head_branch or args.head_sha:
             warning("No need to set --head-branch and --head-sha when using --repo")

.github/workflows/sglang-benchmark.yml

@@ -99,6 +99,11 @@ jobs:
           python-version: '3.12'
           cache: 'pip'
 
+      - name: Get workflow job id
+        id: get-job-id
+        uses: pytorch/test-infra/.github/actions/get-workflow-job-id@36562d6c43fa914f7bdef67ce23e5c31f1387b2e
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
 
       - name: Check if the device is supported
         shell: bash
@@ -197,7 +202,7 @@ jobs:
           for tag in $(git for-each-ref --sort=-creatordate --format '%(refname:short)' refs/tags); do
             candidate_image="lmsysorg/sglang:${tag}${IMAGE_SUFFIX}"
             echo "Checking: $candidate_image"
-            
+
             if docker manifest inspect "$candidate_image" >/dev/null 2>&1; then
               SELECTED_TAG="$tag"
               DOCKER_IMAGE="$candidate_image"
@@ -206,7 +211,7 @@ jobs:
               break
             fi
           done
-          
+
           # Fallback to latest if no tagged version found
           if [[ -z "$SELECTED_TAG" ]]; then
             echo "No tagged images found, using latest"
@@ -280,12 +285,18 @@ jobs:
         env:
           BENCHMARK_RESULTS: sglang-benchmarks/benchmarks/results
           MODELS: ${{ matrix.models }}
+          WORKFLOW_RUN_ID: ${{ github.run_id }}
+          RUN_ATTEMPT: ${{ github.run_attempt }}
+          JOB_ID: ${{ steps.get-job-id.outputs.job-id }}
+          RUN_LOCAL_BRANCH: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.ref || github.ref }}
         run: |
           set -eux
 
           sudo chown -R ${UID} "${BENCHMARK_RESULTS}" || true
           ls -lah "${BENCHMARK_RESULTS}" || echo "Results directory not found"
 
+          echo "## workflow info: ${WORKFLOW_RUN_ID} ${RUN_ATTEMPT} ${JOB_ID}"
+
           SANITIZED_DEVICE_TYPE=$(echo "${DEVICE_TYPE// /_}" | sed "s/[^[:alnum:].-]/_/g")
           SANITIZED_MODELS="${MODELS//\//_}"
 

.github/workflows/vllm-benchmark.yml

@@ -101,6 +101,12 @@ jobs:
           python-version: '3.12'
           cache: 'pip'
 
+      - name: Get workflow job id
+        id: get-job-id
+        uses: pytorch/test-infra/.github/actions/get-workflow-job-id@36562d6c43fa914f7bdef67ce23e5c31f1387b2e
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+
       - name: Check if the device is supported
         shell: bash
         run: |
@@ -115,7 +121,7 @@ jobs:
           elif command -v hl-smi; then
             DEVICE_NAME=hpu
             hl-smi
-          else 
+          else
             arch=$(uname -m)
 
             case "$arch" in
@@ -348,12 +354,18 @@ jobs:
         env:
           BENCHMARK_RESULTS: vllm-benchmarks/vllm/benchmarks/results
           MODELS: ${{ matrix.models }}
+          WORKFLOW_RUN_ID: ${{ github.run_id }}
+          RUN_ATTEMPT: ${{ github.run_attempt }}
+          JOB_ID: ${{ steps.get-job-id.outputs.job-id }}
+          RUN_LOCAL_BRANCH: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.ref || github.ref }}
         run: |
           set -eux
 
           sudo chown -R ${UID} "${BENCHMARK_RESULTS}"
           ls -lah "${BENCHMARK_RESULTS}"
 
+          echo "## workflow info: ${WORKFLOW_RUN_ID} ${RUN_ATTEMPT} ${JOB_ID}"
+
           SANITIZED_DEVICE_TYPE=$(echo "${DEVICE_TYPE// /_}" | sed "s/[^[:alnum:].-]/_/g")
           SANITIZED_MODELS="${MODELS//\//_}"