Tag/name images with branch name

.github/workflows/build.yaml

@@ -6,6 +6,9 @@ on:
       upload_artifacts:
         required: true
         type: boolean
+      branch_name:
+        required: true
+        type: string
 
 permissions:
   id-token: write
@@ -55,6 +58,7 @@ jobs:
         run: ci/compute-variables.sh
         env:
           ARCH: ${{ matrix.ARCH }}
+          BRANCH_NAME: ${{ inputs.branch_name }}
           DRIVER_VERSION: ${{ matrix.DRIVER_VERSION }}
           OS: ${{ matrix.OS }}
           RUNNER_ENV: ${{ matrix.ENV }}
@@ -73,6 +77,7 @@ jobs:
             -only="*${OS}-${RUNNER_ENV}*" \
             -var "arch=${ARCH}" \
             -var "backup_aws_regions=${BACKUP_AWS_REGIONS}" \
+            -var "branch_name=${BRANCH_NAME}" \
             -var "default_aws_region=${DEFAULT_AWS_REGION}" \
             -var "driver_version=${DRIVER_VERSION}" \
             -var "gh_run_id=${NV_RUN_ID}" \
@@ -88,6 +93,7 @@ jobs:
         env:
           ARCH: ${{ matrix.ARCH }}
           BACKUP_AWS_REGIONS: ${{ needs.compute-constants.outputs.BACKUP_AWS_REGIONS }}
+          BRANCH_NAME: ${{ inputs.branch_name }}
           DEFAULT_AWS_REGION: ${{ needs.compute-constants.outputs.DEFAULT_AWS_REGION }}
           DRIVER_VERSION: ${{ matrix.DRIVER_VERSION }}
           GH_TOKEN: ${{ github.token }}

.github/workflows/gc.yaml

@@ -6,11 +6,17 @@ on:
       dry_run:
         default: true
         type: boolean
+      branch_name:
+        required: true
+        type: string
   workflow_dispatch:
     inputs:
       dry_run:
         default: true
         type: boolean
+      branch_name:
+        required: true
+        type: string
   schedule:
     # Run every Monday at 1pm UTC / 8am ET
     - cron: "0 13 * * 1"
@@ -49,4 +55,5 @@ jobs:
       - name: Run garbage collection
         run: python main.py --dry-run="${DRY_RUN}"
         env:
+          BRANCH_NAME: ${{ inputs.branch_name }}
           DRY_RUN: ${{ github.event_name == 'schedule' && 'false' || inputs.dry_run }}

.github/workflows/main.yaml

@@ -18,4 +18,5 @@ jobs:
     uses: ./.github/workflows/build.yaml
     with:
       upload_artifacts: true
+      branch_name: ${{ github.ref_name }}
     secrets: inherit

.github/workflows/pr.yaml

@@ -36,9 +36,11 @@ jobs:
     uses: ./.github/workflows/gc.yaml
     with:
       dry_run: true
+      branch_name: ${{ github.ref_name }}
     secrets: inherit
   build:
     uses: ./.github/workflows/build.yaml
     with:
       upload_artifacts: false
+      branch_name: ${{ github.ref_name }}
     secrets: inherit

ci/compute-image-name.sh

@@ -15,12 +15,15 @@ IMAGE_NAME=$(
     --arg DRIVER_VERSION "${DRIVER_VERSION}" \
     --arg ARCH "${ARCH}" \
     --arg RUNNER_VERSION "${RUNNER_VERSION}" \
+    --arg RUNNER_ENV "${RUNNER_ENV}" \
+    --arg BRANCH_NAME "${BRANCH_NAME}" \
   '[
     $OS,
     $VARIANT,
     $DRIVER_VERSION,
     $ARCH,
-    $RUNNER_VERSION
+    $RUNNER_VERSION,
+    if $RUNNER_ENV != "aws" then ($BRANCH_NAME | sub("/"; "-")) else empty end
   ] | map(select(length > 0)) | join("-")'
 )
 

gc/main.py

@@ -2,7 +2,7 @@
 import subprocess
 import json
 import yaml
-from os import path
+from os import path, getenv
 from collectors.ecr import ECRGarbageCollector
 from collectors.amis import AMIGarbageCollector
 from collectors.gc import GarbageCollector
@@ -27,8 +27,11 @@ def load_current_images() -> list[str]:
         result = subprocess.run(
             compute_image_name_path,
             cwd="..",
-            capture_output=True,
-            env=entry,
+            stdout=subprocess.PIPE,
+            env={
+                **{"RUNNER_ENV" if k == "ENV" else k: v for k, v in entry.items()},
+                "BRANCH_NAME": getenv("BRANCH_NAME"),
+            },
             check=True,
         )
         images.append(result.stdout.decode("utf-8").strip())

locals.pkr.hcl

@@ -5,6 +5,7 @@ locals {
       "driver-version" = var.driver_version
       "os"             = var.os
       "runner-version" = var.runner_version
+      "branch-name"    = var.branch_name
       "variant"        = local.variant
       "Name"           = local.image_id
     } : k => v if v != ""

matrix.yaml

@@ -25,7 +25,8 @@ ARCH:
 
 ENV:
   - aws
-  - qemu
+  # TODO: Uncomment this before merging
+  #- qemu
 
 exclude:
   # only use amd64 for windows

variables.pkr.hcl

@@ -14,6 +14,11 @@ variable "backup_aws_regions" {
   description = "A comma-separated list of the AWS regions to copy the AMIs to."
 }
 
+variable "branch_name" {
+  type        = string
+  description = "The name of the GitHub branch the image originated from."
+}
+
 variable "default_aws_region" {
   type        = string
   description = "The AWS region to provision EC2 instances in."