avoid baking in PR trial merge commit hashes

Author: olupton

Diff for: .github/container/Dockerfile.base

@@ -3,7 +3,7 @@ ARG BASE_IMAGE=nvidia/cuda:12.6.2-devel-ubuntu22.04
 ARG GIT_USER_NAME="JAX Toolbox"
 ARG [email protected]
 ARG CLANG_VERSION=18
-ARG JAX_TOOLBOX_SHA
+ARG JAX_TOOLBOX_REF
 
 ###############################################################################
 ## Obtain GCP's NCCL TCPx plugin
@@ -31,7 +31,7 @@ ARG BASE_IMAGE
 ARG GIT_USER_EMAIL
 ARG GIT_USER_NAME
 ARG CLANG_VERSION
-ARG JAX_TOOLBOX_SHA
+ARG JAX_TOOLBOX_REF
 ENV CUDA_BASE_IMAGE=${BASE_IMAGE}
 
 ###############################################################################
@@ -204,7 +204,7 @@ COPY check-shm.sh /opt/nvidia/entrypoint.d/
 ###############################################################################
 
 ADD install-nsys-jax.sh /usr/local/bin
-RUN install-nsys-jax.sh ${JAX_TOOLBOX_SHA}
+RUN install-nsys-jax.sh ${JAX_TOOLBOX_REF}
 
 ###############################################################################
 ## Copy manifest file to the container

Diff for: .github/workflows/_build_base.yaml

@@ -128,11 +128,13 @@ jobs:
           platforms: linux/${{ inputs.ARCHITECTURE }}
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
+          # head_ref is the PR source branch for pull_request pipelines, which avoids
+          # baking in the SHA of a merge commit than cannot be checked out later
           build-args: |
             GIT_USER_NAME=${{ inputs.GIT_USER_NAME }}
             GIT_USER_EMAIL=${{ inputs.GIT_USER_EMAIL }}
             BUILD_DATE=${{ inputs.BUILD_DATE }}
-            JAX_TOOLBOX_SHA=${{ github.sha }}
+            JAX_TOOLBOX_REF=${{ github.head_ref || github.sha }}
             ${{ inputs.BASE_IMAGE != 'latest' && format('BASE_IMAGE={0}', inputs.BASE_IMAGE) || '' }}
         
       - name: Generate sitrep