Use the right Docker image for non-main commits (#45)

* Use the right Docker image for non-main commits

Signed-off-by: Huy Do <huydhn@gmail.com>

* Debug

Signed-off-by: Huy Do <huydhn@gmail.com>

* Fix the logic to detect branch name

Signed-off-by: Huy Do <huydhn@gmail.com>

---------

Signed-off-by: Huy Do <huydhn@gmail.com>

Author: huydhn

.github/scripts/upload_benchmark_results.py

@@ -131,8 +131,16 @@ def get_git_metadata(repo_dir: str) -> Tuple[str, str]:
             hexsha,
             committed_date,
         )
-    except TypeError:
-        # This is a detached HEAD, default the branch to main
+    except TypeError as e:
+        # This is a detached HEAD, try to find out where the commit comes from
+        for head in repo.heads:
+            info(f"Check commits from {head.name}")
+            for commit in repo.iter_commits(head):
+                if commit.hexsha == hexsha:
+                    return repo_name, head.name, hexsha, committed_date
+
+        warning(f"Found no branch name, default to main: {e}")
+        # We couldn't find the branch name
         return repo_name, "main", hexsha, committed_date
 
 

.github/workflows/vllm-benchmark.yml

@@ -139,8 +139,18 @@ jobs:
 
       - name: Set Docker registry
         shell: bash
+        env:
+          HEAD_BRANCH: ${{ inputs.vllm_branch || 'main' }}
         run: |
-          DOCKER_IMAGE_PREFIX=public.ecr.aws/q9t5s3a7/vllm-ci-postmerge-repo
+          set -eux
+
+          # Mimic the logic from vllm ci-infra test template
+          if [[ "${HEAD_BRANCH}" == "main" ]]; then
+            DOCKER_IMAGE_PREFIX=public.ecr.aws/q9t5s3a7/vllm-ci-postmerge-repo
+          else
+            DOCKER_IMAGE_PREFIX=public.ecr.aws/q9t5s3a7/vllm-ci-test-repo
+          fi
+
           DOCKER_IMAGE_SUFFIX=""
           if [[ "${DEVICE_NAME}" == "rocm" ]]; then
             DOCKER_IMAGE_PREFIX=docker.io/rocm/vllm-ci
@@ -284,7 +294,7 @@ jobs:
           sudo chown -R ${UID} "${BENCHMARK_RESULTS}"
           ls -lah "${BENCHMARK_RESULTS}"
 
-          SANITIZED_DEVICE_TYPE=$(echo "${DEVICE_TYPE// /_}" | sed "s/[^[:alpha:].-]/_/g")
+          SANITIZED_DEVICE_TYPE=$(echo "${DEVICE_TYPE// /_}" | sed "s/[^[:alnum:].-]/_/g")
           python3 .github/scripts/upload_benchmark_results.py \
             --repo vllm-benchmarks/vllm \
             --benchmark-name "vLLM benchmark" \