Skip to content

Commit f80ffc8

Browse files
KyleFromNVIDIAKyleFromNVIDIA
authored
Tag/name images with branch name (#92)
To support building multiple branches in parallel, include the branch name, either as a tag or as part of the image name. Contributes to nv-gha-runners/roadmap#127
1 parent 59dc8e4 commit f80ffc8

18 files changed

+228
-41
lines changed

‎.github/workflows/build.yaml

+13
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,9 @@ on:
66
upload_artifacts:
77
required: true
88
type: boolean
9+
branch_name:
10+
required: true
11+
type: string
912

1013
permissions:
1114
id-token: write
@@ -16,6 +19,13 @@ defaults:
1619
shell: bash
1720

1821
jobs:
22+
test:
23+
runs-on: ubuntu-latest
24+
steps:
25+
- name: Checkout
26+
uses: actions/checkout@v4
27+
- name: Run tests
28+
run: ci/image-name/tests.sh
1929
compute-constants:
2030
runs-on: ubuntu-latest
2131
outputs:
@@ -55,6 +65,7 @@ jobs:
5565
run: ci/compute-variables.sh
5666
env:
5767
ARCH: ${{ matrix.ARCH }}
68+
BRANCH_NAME: ${{ inputs.branch_name }}
5869
DRIVER_VERSION: ${{ matrix.DRIVER_VERSION }}
5970
DRIVER_FLAVOR: ${{ matrix.DRIVER_FLAVOR }}
6071
OS: ${{ matrix.OS }}
@@ -74,6 +85,7 @@ jobs:
7485
-only="*${OS}-${RUNNER_ENV}*" \
7586
-var "arch=${ARCH}" \
7687
-var "backup_aws_regions=${BACKUP_AWS_REGIONS}" \
88+
-var "branch_name=${BRANCH_NAME}" \
7789
-var "default_aws_region=${DEFAULT_AWS_REGION}" \
7890
-var "driver_version=${DRIVER_VERSION}" \
7991
-var "driver_flavor=${DRIVER_FLAVOR}" \
@@ -90,6 +102,7 @@ jobs:
90102
env:
91103
ARCH: ${{ matrix.ARCH }}
92104
BACKUP_AWS_REGIONS: ${{ needs.compute-constants.outputs.BACKUP_AWS_REGIONS }}
105+
BRANCH_NAME: ${{ inputs.branch_name }}
93106
DEFAULT_AWS_REGION: ${{ needs.compute-constants.outputs.DEFAULT_AWS_REGION }}
94107
DRIVER_VERSION: ${{ matrix.DRIVER_VERSION }}
95108
DRIVER_FLAVOR: ${{ matrix.DRIVER_FLAVOR }}

‎.github/workflows/gc.yaml

+1
Original file line numberDiff line numberDiff line change
@@ -50,3 +50,4 @@ jobs:
5050
run: python main.py --dry-run="${DRY_RUN}"
5151
env:
5252
DRY_RUN: ${{ github.event_name == 'schedule' && 'false' || inputs.dry_run }}
53+
REPOSITORY: ${{ github.repository }}

‎.github/workflows/main.yaml

+2
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ on:
44
push:
55
branches:
66
- main
7+
- test-**
78

89
permissions:
910
id-token: write
@@ -18,4 +19,5 @@ jobs:
1819
uses: ./.github/workflows/build.yaml
1920
with:
2021
upload_artifacts: true
22+
branch_name: ${{ github.ref_name }}
2123
secrets: inherit

‎.github/workflows/pr.yaml

+1
Original file line numberDiff line numberDiff line change
@@ -41,4 +41,5 @@ jobs:
4141
uses: ./.github/workflows/build.yaml
4242
with:
4343
upload_artifacts: false
44+
branch_name: ${{ github.ref_name }}
4445
secrets: inherit

‎ci/compute-variables.sh

+1-1
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
# A script that computes the variables for our GHAs workflows
33
set -euo pipefail
44

5-
IMAGE_NAME=$(ci/compute-image-name.sh)
5+
IMAGE_NAME=$(ci/image-name/serialize.sh)
66

77
cat <<EOF | tee --append "${GITHUB_ENV:-/dev/null}"
88
NV_IMAGE_NAME=${IMAGE_NAME}

‎ci/image-name/deserialize.jq

+29
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
def pick_from_list($values; $required_field):
2+
.[0] as $value |
3+
if $values | any(. == $value) then
4+
{"value": $value, "rest": .[1:]}
5+
elif $required_field != "" then
6+
"Error: \($required_field) field is required\n" | halt_error
7+
else
8+
{"value": null, "rest": .}
9+
end;
10+
11+
def deserialize_image_name($matrix):
12+
split("-") |
13+
pick_from_list($matrix.OS; "OS") | .value as $os | .rest |
14+
pick_from_list(["cpu", "gpu"]; "variant") | .value as $variant | .rest |
15+
pick_from_list($matrix.DRIVER_VERSION; "") | .value as $driver_version | .rest |
16+
pick_from_list($matrix.DRIVER_FLAVOR; "") | .value as $driver_flavor | .rest |
17+
pick_from_list($matrix.ARCH; "arch") | .value as $arch | .rest |
18+
(if any then join("-") else null end) as $branch_name |
19+
{
20+
"os": $os,
21+
"variant": $variant,
22+
"driver_version": $driver_version,
23+
"driver_flavor": $driver_flavor,
24+
"arch": $arch,
25+
"branch_name": $branch_name
26+
};
27+
28+
def deserialize_image_name:
29+
. as $input | env.MATRIX | fromjson as $matrix | $input | deserialize_image_name($matrix);

‎ci/image-name/deserialize.sh

+7
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
#!/bin/bash
2+
3+
set -euo pipefail
4+
5+
cd "$(dirname "$0")/../.."
6+
7+
MATRIX="$(yq -o json matrix.yaml)" IMAGE_NAME="$1" jq -c -n 'include "ci/image-name/deserialize"; env.IMAGE_NAME | deserialize_image_name'

‎ci/image-name/deserialize.test

+10
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
# Valid image names
2+
include "ci/image-name/deserialize"; .[] | deserialize_image_name
3+
["linux-gpu-550-open-arm64-main", "windows-cpu-amd64-pr-1234", "linux-gpu-550-open-arm64"]
4+
{"os": "linux", "variant": "gpu", "driver_version": "550", "driver_flavor": "open", "arch": "arm64", "branch_name": "main"}
5+
{"os": "windows", "variant": "cpu", "driver_version": null, "driver_flavor": null, "arch": "amd64", "branch_name": "pr-1234"}
6+
{"os": "linux", "variant": "gpu", "driver_version": "550", "driver_flavor": "open", "arch": "arm64", "branch_name": null}
7+
8+
# Invalid image names
9+
include "ci/image-name/deserialize"; .[] | deserialize_image_name
10+
["linux-550-arm64-main", "gpu-550-arm64-main", "linux-gpu-550-main"]

‎ci/compute-image-name.sh ‎ci/image-name/serialize.sh

+2-2
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,14 @@ IMAGE_NAME=$(
1515
--arg DRIVER_VERSION "${DRIVER_VERSION}" \
1616
--arg DRIVER_FLAVOR "${DRIVER_FLAVOR}" \
1717
--arg ARCH "${ARCH}" \
18-
--arg RUNNER_VERSION "${RUNNER_VERSION}" \
18+
--arg BRANCH_NAME "${BRANCH_NAME}" \
1919
'[
2020
$OS,
2121
$VARIANT,
2222
$DRIVER_VERSION,
2323
$DRIVER_FLAVOR,
2424
$ARCH,
25-
$RUNNER_VERSION
25+
$BRANCH_NAME | sub("/"; "-")
2626
] | map(select(length > 0)) | join("-")'
2727
)
2828

‎ci/image-name/tests.sh

+8
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
#!/bin/bash
2+
3+
set -euo pipefail
4+
5+
MATRIX="$(yq -o json matrix.yaml)"
6+
MATRIX="${MATRIX}" jq --run-tests ci/image-name/deserialize.test
7+
8+
ci/image-name/deserialize.sh linux-gpu-550-open-amd64-pr-1234 | jq -e '. == {"os": "linux", "variant": "gpu", "driver_version": "550", "driver_flavor": "open", "arch": "amd64", "branch_name": "pr-1234"}'

‎gc/collectors/amis.py

+39-12
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,6 @@
1+
from os import getenv
2+
from typing import Optional
3+
14
import boto3
25
from collectors import gc
36
from collections import defaultdict
@@ -9,10 +12,17 @@
912

1013

1114
class AMIGarbageCollector(gc.GarbageCollector):
12-
def __init__(self, current_images: list[str], region: str, dry_run: bool):
15+
def __init__(
16+
self,
17+
current_images: list[str],
18+
current_branches: list[str],
19+
region: str,
20+
dry_run: bool,
21+
):
1322
super().__init__("AMI Collector", region, dry_run)
1423
self.ec2_client = boto3.client("ec2", region_name=region)
1524
self.current_images = current_images
25+
self.current_branches = current_branches
1626
self.search_tag_name = "vm-images"
1727
self.search_tag_value = "true"
1828

@@ -44,29 +54,46 @@ def _get_amis(self) -> list[ImageTypeDef]:
4454

4555
def _find_expired_amis(self, amis: list[ImageTypeDef]) -> list[ImageTypeDef]:
4656
expired_amis = []
47-
ami_groups = defaultdict(list)
57+
ami_groups: dict[str, dict[str, list[ImageTypeDef]]] = defaultdict(
58+
lambda: defaultdict(list)
59+
)
4860

4961
# Group AMIs by "image-name" tag
5062
for ami in amis:
5163
img_tags = ami["Tags"]
5264
if img_tags:
65+
image_name: Optional[str] = None
66+
branch_name: Optional[str] = None
5367
for tag in img_tags:
5468
if tag["Key"] == "image-name":
55-
img_name = tag["Value"]
56-
ami_groups[img_name].append(ami)
69+
image_name = tag["Value"]
70+
if tag["Key"] == "branch-name":
71+
branch_name = tag["Value"]
72+
if image_name:
73+
if branch_name:
74+
ami_groups[branch_name][image_name].append(ami)
5775
break
76+
else:
77+
expired_amis.append(ami)
5878

5979
# Sort AMIs by creation date.
6080
# If image is currently supported, keep only the newest AMI. Expire the rest.
6181
# If image is not currently supported, expire all AMIs.
62-
for img_name, amis in ami_groups.items():
63-
amis = sorted(
64-
amis, key=lambda x: parser.parse(x["CreationDate"]), reverse=True
65-
)
66-
if img_name in self.current_images:
67-
expired_amis.extend(amis[1:])
68-
else:
69-
expired_amis.extend(amis)
82+
for branch_name, images in ami_groups.items():
83+
if branch_name == gc.DEFAULT_BRANCH_NAME:
84+
for image_name, amis in images.items():
85+
if image_name in self.current_images:
86+
amis = sorted(
87+
amis,
88+
key=lambda x: parser.parse(x["CreationDate"]),
89+
reverse=True,
90+
)
91+
expired_amis.extend(amis[1:])
92+
else:
93+
expired_amis.extend(amis)
94+
elif branch_name not in self.current_branches:
95+
for image_name, amis in images.items():
96+
expired_amis.extend(amis)
7097

7198
return expired_amis
7299

‎gc/collectors/ecr.py

+20-5
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,21 @@
44
from mypy_boto3_ecr_public.paginator import (
55
DescribeImagesPaginator as ECRPublicDescribeImagesPaginator,
66
)
7+
from .image_name import deserialize_image_name
78

89

910
class ECRGarbageCollector(gc.GarbageCollector):
10-
def __init__(self, current_images: list[str], region: str, dry_run: bool):
11+
def __init__(
12+
self,
13+
current_images: list[str],
14+
current_branches: list[str],
15+
region: str,
16+
dry_run: bool,
17+
):
1118
super().__init__("ECR Collector", region, dry_run)
1219
self.ecr_client = boto3.client("ecr-public", region_name=region)
1320
self.current_images = current_images
21+
self.current_branches = current_branches
1422
self.repository_name = "kubevirt-images"
1523

1624
def _run(self) -> None:
@@ -34,19 +42,26 @@ def _get_ecr_images(self) -> list[ImageDetailTypeDef]:
3442
def _find_expired_ecr_images(
3543
self, images: list[ImageDetailTypeDef]
3644
) -> list[ImageDetailTypeDef]:
45+
branches = {b.replace("/", "-") for b in self.current_branches}
3746
expired_images = []
3847

3948
for image in images:
4049
image_tags = image.get("imageTags")
4150
hasSupportedTags = False
4251
if image_tags:
4352
for tag in image_tags:
44-
if tag in self.current_images:
45-
hasSupportedTags = True
46-
break
53+
parsed_image_name = deserialize_image_name(tag)
54+
if parsed_image_name:
55+
if parsed_image_name.branch_name == gc.DEFAULT_BRANCH_NAME:
56+
if tag in self.current_images:
57+
hasSupportedTags = True
58+
break
59+
elif parsed_image_name.branch_name in branches:
60+
hasSupportedTags = True
61+
break
4762

4863
# Remove images that don't have any tags or don't have any supported tags
49-
if not image_tags or not hasSupportedTags:
64+
if not hasSupportedTags:
5065
expired_images.append(image)
5166
continue
5267
return expired_images

‎gc/collectors/gc.py

+3
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
from abc import ABC, abstractmethod
22

33

4+
DEFAULT_BRANCH_NAME = "main"
5+
6+
47
class GarbageCollector(ABC):
58
def __init__(self, collector_name: str, region: str, dry_run: bool):
69
self.dry_run = dry_run

‎gc/collectors/image_name.py

+35
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
import dataclasses
2+
import json
3+
import os.path
4+
import subprocess
5+
from typing import Optional
6+
7+
8+
DESERIALIZE_PATH = os.path.join(
9+
os.path.dirname(__file__), "..", "..", "ci", "image-name", "deserialize.sh"
10+
)
11+
12+
13+
@dataclasses.dataclass
14+
class ImageName:
15+
os: str
16+
variant: str
17+
driver_version: Optional[str]
18+
driver_flavor: Optional[str]
19+
arch: str
20+
branch_name: Optional[str]
21+
22+
23+
def deserialize_image_name(image_name: str) -> Optional[ImageName]:
24+
result = subprocess.run(
25+
[DESERIALIZE_PATH, image_name], stdout=subprocess.PIPE, text=True, check=True
26+
)
27+
parsed_json = json.loads(result.stdout)
28+
return ImageName(
29+
os=parsed_json["os"],
30+
variant=parsed_json["variant"],
31+
driver_version=parsed_json["driver_version"],
32+
driver_flavor=parsed_json["driver_flavor"],
33+
arch=parsed_json["arch"],
34+
branch_name=parsed_json["branch_name"],
35+
)

0 commit comments

Comments
 (0)