feat(validator): add AICR_VALIDATOR_IMAGE_TAG env-var override

yuanchen8911 · yuanchen8911 · commit 0bea785f9884 · 2026-04-23T19:31:51.000-07:00
Running `aicr validate` from a feature-branch dev build fails with ImagePullBackOff on every validator pod because on-push.yaml only pushes `:sha-<commit>` images for commits merged to `main`, while PR #655 (rightly) made non-release builds resolve to `:sha-<commit>` instead of `:latest`. Contributors dog-fooding a PR before merge hit a "NotFound" dead-end with no escape hatch. This adds AICR_VALIDATOR_IMAGE_TAG as an opt-in override. When set, the resolved tag is replaced on every validator image — including explicit catalog tags like `:v1.2.3` — so a feature-branch build can point at a known-published tag: AICR_VALIDATOR_IMAGE_TAG=latest aicr validate --phase performance ... Default behavior is unchanged: release builds still resolve to `:v<version>`, main-branch dev builds still resolve to `:sha-<commit>`, and reproducibility for CI paths is preserved. The override is strictly additive and strictly opt-in. The override is forwarded from the CLI invocation into the validator container (alongside AICR_CLI_VERSION, AICR_CLI_COMMIT, and AICR_VALIDATOR_IMAGE_REGISTRY), so validators that resolve inner workload images at runtime (inference-perf's AIPerf benchmark Job) apply the same semantics as catalog.Load. Without this, the outer validator pod would get `:latest` while the inner benchmark pod would still resolve to the same unpublished `:sha-<commit>` and ImagePullBackOff — defeating the motivating feature-branch workflow. Digest-pinned references (`name@sha256:…`) are preserved verbatim. A tag override is meaningless against a content-addressable pin, and naive last-colon splitting would corrupt the digest hash into the tag slot. replaceTag detects the `@` separator and returns the image unchanged; the registry override still applies to digest refs. Tests: - catalog_test.go: 10 new cases covering the tag-override escape hatch, composition with registry override, the no-tag-append case, the empty-env-var no-op, the localhost:5001 port-preservation edge case, digest-only and mixed `name:tag@digest` forms - deployer_test.go: 2 new cases asserting the env var is forwarded into the validator container when set, and strictly omitted when unset (so the default release / main-branch paths are untouched) - docs/contributor/validator.md updated with the resolution order, digest behavior, and env-var forwarding
diff --git a/docs/contributor/validator.md b/docs/contributor/validator.md
@@ -136,6 +136,7 @@ The validator engine mounts snapshot and recipe data as ConfigMaps:
 | `AICR_SNAPSHOT_PATH` | Override snapshot mount path |
 | `AICR_RECIPE_PATH` | Override recipe mount path |
 | `AICR_VALIDATOR_IMAGE_REGISTRY` | Override image registry prefix (set by user) |
+| `AICR_VALIDATOR_IMAGE_TAG` | Override the resolved image tag (e.g. `latest`). Bypasses the default `:v<version>` / `:sha-<commit>` resolution for feature-branch dev builds whose commit has no published image. |
 | `AICR_NODE_SELECTOR` | User-provided node selector override for inner workloads (comma-separated `key=value` pairs). Set by the `--node-selector` CLI flag. Use `ctx.NodeSelector` to access the parsed value. |
 | `AICR_TOLERATIONS` | User-provided toleration override for inner workloads (comma-separated `key=value:effect` entries). Set by the `--toleration` CLI flag. Use `ctx.Tolerations` to access the parsed value. |
 
@@ -227,8 +228,16 @@ Each entry in `recipes/validators/catalog.yaml`:
 **Image tag resolution** (applied by `catalog.Load`):
 
 1. `:latest` tags are replaced with the CLI version (e.g., `:v0.9.5`) for release builds
-2. Explicit version tags (e.g., `:v1.2.3`) are never modified
-3. `AICR_VALIDATOR_IMAGE_REGISTRY` overrides the registry prefix
+2. On non-release dev builds with a valid commit, `:latest` becomes `:sha-<commit>` (matches the tags `on-push.yaml` pushes for merges to `main`)
+3. Explicit version tags (e.g., `:v1.2.3`) are not modified by steps 1-2
+4. `AICR_VALIDATOR_IMAGE_TAG` overrides the resolved tag on every validator image, including explicit catalog tags. Use this when running `aicr validate` from a feature-branch dev build whose commit has not been merged to `main` (no `:sha-<commit>` image has been published). Typical value: `latest`. Example: `AICR_VALIDATOR_IMAGE_TAG=latest aicr validate --phase performance ...`
+5. `AICR_VALIDATOR_IMAGE_REGISTRY` overrides the registry prefix
+
+**Digest-pinned references** (`name@sha256:…`) are not rewritten by step 4. A tag override is meaningless against a content-addressable pin, and naive rewriting would corrupt the digest. Step 5's registry override still applies — only the registry prefix changes, the digest is preserved verbatim.
+
+**Env-var forwarding to the validator pod:** `AICR_CLI_VERSION`, `AICR_CLI_COMMIT`, `AICR_VALIDATOR_IMAGE_REGISTRY`, and `AICR_VALIDATOR_IMAGE_TAG` are forwarded from the CLI invocation into the validator container so that validators resolving inner workload images at runtime (e.g. `inference-perf`'s AIPerf benchmark Job) apply the same semantics as `catalog.Load`. If you set `AICR_VALIDATOR_IMAGE_TAG=latest` on the CLI, the override reaches both the outer validator Job and the inner benchmark Job — they always travel together.
+
+**Pull-policy behavior when the override is set:** when `AICR_VALIDATOR_IMAGE_TAG` is set, the outer validator Job uses `imagePullPolicy: Always` regardless of the resolved tag. Override values are typically mutable (`latest`, `edge`, `main`, or any tag `on-push.yaml` recreates on every merge), so a default `IfNotPresent` would let a node's previously cached image win over the tag's current target. The inner workload Job in each validator (e.g. the AIPerf benchmark pod dispatched by `inference-perf`, implemented in the `NVIDIA/aicr-validators` repo) should mirror this behavior when consuming `AICR_VALIDATOR_IMAGE_TAG`; see that repo for the current state.
 
 **Performance phase example — inference perf:**
 
diff --git a/pkg/validator/catalog/catalog.go b/pkg/validator/catalog/catalog.go
@@ -103,9 +103,14 @@ type EnvVar struct {
 //     the tag is replaced with the CLI version for reproducibility.
 //  2. If version is a non-release dev build and commit is a valid short SHA,
 //     the tag is replaced with :sha-<commit> to match on-push.yaml image tags.
-//  3. If AICR_VALIDATOR_IMAGE_REGISTRY is set, the registry prefix is replaced.
+//  3. If AICR_VALIDATOR_IMAGE_TAG is set, the resolved tag is overridden.
+//     Useful for feature-branch dev builds whose commit SHA has no published
+//     image (on-push.yaml only pushes SHA tags for commits merged to main).
+//     Common value: `latest`.
+//  4. If AICR_VALIDATOR_IMAGE_REGISTRY is set, the registry prefix is replaced.
 //
-// Entries with explicit version tags (e.g., :v1.2.3) are never modified.
+// Entries with explicit version tags (e.g., :v1.2.3) are never modified by
+// steps 1-2 but are replaced by step 3 if that env var is set.
 func Load(version, commit string) (*ValidatorCatalog, error) {
 	data, err := recipe.GetDataProvider().ReadFile("validators/catalog.yaml")
 	if err != nil {
@@ -131,7 +136,11 @@ func Load(version, commit string) (*ValidatorCatalog, error) {
 //
 //  1. :latest tag replacement with version if version is a release (vX.Y.Z).
 //  2. If non-release and commit is a valid SHA, :latest → :sha-<commit>.
-//  3. Registry prefix override if AICR_VALIDATOR_IMAGE_REGISTRY is set.
+//  3. Tag override if AICR_VALIDATOR_IMAGE_TAG is set (overrides steps 1-2
+//     AND explicit catalog tags). Intended for feature-branch dev builds
+//     where no :sha-<commit> image has been published; typical value:
+//     `latest`.
+//  4. Registry prefix override if AICR_VALIDATOR_IMAGE_REGISTRY is set.
 //
 // Images with explicit version tags are not modified by steps 1-2.
 func ResolveImage(image, version, commit string) string {
@@ -141,6 +150,9 @@ func ResolveImage(image, version, commit string) string {
 	} else if isValidCommit(commit) {
 		image = replaceLatestWithSHA(image, commit)
 	}
+	if tag := os.Getenv("AICR_VALIDATOR_IMAGE_TAG"); tag != "" {
+		image = replaceTag(image, tag)
+	}
 	if override := os.Getenv("AICR_VALIDATOR_IMAGE_REGISTRY"); override != "" {
 		image = replaceRegistry(image, override)
 	}
@@ -192,6 +204,35 @@ func isValidCommit(commit string) bool {
 	return true
 }
 
+// replaceTag forces the image's tag to newTag, regardless of what tag (if
+// any) the image currently carries. Unlike replaceLatestTag / replaceLatestWithSHA,
+// which only rewrite :latest, this helper supports the AICR_VALIDATOR_IMAGE_TAG
+// env-var escape hatch: a user running a feature-branch dev build (where no
+// :sha-<commit> image was published by on-push.yaml) can set the env var
+// to `latest` and force every validator image to a published tag.
+//
+// Digest-pinned references (`name@sha256:…`) are cryptographic pins and are
+// intentionally left untouched — a tag override is meaningless against a
+// content-addressable ref, and naively rewriting would corrupt the digest.
+// For non-digest refs, the tag separator is found as the last ':' that sits
+// after the last '/' to avoid colliding with the registry port (`:5001` in
+// `localhost:5001/...`).
+func replaceTag(image, newTag string) string {
+	if strings.Contains(image, "@") {
+		// Digest-pinned ref (e.g. ghcr.io/foo/bar@sha256:deadbeef, or the
+		// mixed form name:tag@sha256:…). The digest is the authoritative
+		// pin; preserve it verbatim.
+		return image
+	}
+	slash := strings.LastIndex(image, "/")
+	colon := strings.LastIndex(image, ":")
+	if colon <= slash {
+		// No tag on the image (just an image reference) — append one.
+		return image + ":" + newTag
+	}
+	return image[:colon] + ":" + newTag
+}
+
 // replaceLatestWithSHA replaces :latest with :sha-<commit> to match the
 // image tags pushed by the on-push CI workflow.
 // Images with explicit version tags are not modified.
diff --git a/pkg/validator/catalog/catalog_test.go b/pkg/validator/catalog/catalog_test.go
@@ -492,6 +492,7 @@ func TestResolveImage(t *testing.T) {
 		version  string
 		commit   string
 		registry string // if non-empty, sets AICR_VALIDATOR_IMAGE_REGISTRY for the test
+		tag      string // if non-empty, sets AICR_VALIDATOR_IMAGE_TAG for the test
 		want     string
 	}{
 		{
@@ -642,6 +643,107 @@ func TestResolveImage(t *testing.T) {
 			commit:  "ABC1234",
 			want:    "ghcr.io/nvidia/aicr-validators/aiperf-bench:sha-abc1234",
 		},
+		// --- AICR_VALIDATOR_IMAGE_TAG escape hatch -----------------------
+		// Motivating scenario: a contributor building aicr from an
+		// un-merged feature-branch checkout. The commit isn't on main,
+		// so on-push.yaml never pushed :sha-<commit> to ghcr, and
+		// `aicr validate` pod ImagePullBackOffs. Setting
+		// AICR_VALIDATOR_IMAGE_TAG=latest (or any published tag) forces
+		// every validator image to a reachable tag without losing the
+		// registry / version-based resolution as the default.
+		{
+			name:    "tag override rewrites :latest on dev build (feature-branch dogfooding)",
+			image:   imgLatest,
+			version: "dev",
+			commit:  "abc1234",
+			tag:     "latest",
+			want:    "ghcr.io/nvidia/aicr-validators/aiperf-bench:latest",
+		},
+		{
+			name:    "tag override replaces :sha-<commit> when both resolve",
+			image:   imgLatest,
+			version: "v0.11.1-next",
+			commit:  "abc1234",
+			tag:     "latest",
+			want:    "ghcr.io/nvidia/aicr-validators/aiperf-bench:latest",
+		},
+		{
+			name:    "tag override replaces the release :vX.Y.Z tag",
+			image:   imgLatest,
+			version: "v0.11.1",
+			tag:     "latest",
+			want:    "ghcr.io/nvidia/aicr-validators/aiperf-bench:latest",
+		},
+		{
+			name:    "tag override replaces an explicit catalog tag",
+			image:   imgPinned, // :v1.2.3 is normally untouched
+			version: "v0.11.1",
+			tag:     "latest",
+			want:    "ghcr.io/nvidia/aicr-validators/aiperf-bench:latest",
+		},
+		{
+			name:    "tag override with no tag on image appends the tag",
+			image:   "ghcr.io/nvidia/aicr-validators/aiperf-bench",
+			version: "dev",
+			tag:     "latest",
+			want:    "ghcr.io/nvidia/aicr-validators/aiperf-bench:latest",
+		},
+		{
+			name:    "empty tag env var leaves image untouched (no-op)",
+			image:   imgPinned,
+			version: "v0.11.1",
+			tag:     "", // explicitly empty — should behave like unset
+			want:    imgPinned,
+		},
+		{
+			name:    "tag override preserves registry port (localhost:5001 edge case)",
+			image:   "localhost:5001/aicr-validators/aiperf-bench:sha-abc1234",
+			version: "dev",
+			commit:  "abc1234",
+			tag:     "latest",
+			want:    "localhost:5001/aicr-validators/aiperf-bench:latest",
+		},
+		{
+			name:     "tag override and registry override compose",
+			image:    imgLatest,
+			version:  "dev",
+			commit:   "abc1234",
+			tag:      "v0.11.0",
+			registry: "localhost:5001",
+			want:     "localhost:5001/aicr-validators/aiperf-bench:v0.11.0",
+		},
+		// --- digest-pinned refs must never be tag-rewritten --------------
+		// A tag override is incompatible with a content-addressable digest
+		// pin. Naive last-colon splitting would corrupt `@sha256:<hash>`
+		// into `@sha256:<newTag>`, emitting an invalid reference.
+		{
+			name:    "digest-pinned image is not rewritten by tag override",
+			image:   "ghcr.io/foo/bar@sha256:deadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef",
+			version: "dev",
+			commit:  "abc1234",
+			tag:     "latest",
+			want:    "ghcr.io/foo/bar@sha256:deadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef",
+		},
+		{
+			name:    "mixed ref (name:tag@digest) is not rewritten by tag override",
+			image:   "ghcr.io/foo/bar:v1.0.0@sha256:deadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef",
+			version: "dev",
+			commit:  "abc1234",
+			tag:     "latest",
+			want:    "ghcr.io/foo/bar:v1.0.0@sha256:deadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef",
+		},
+		{
+			// Registry override still applies to digest refs (existing
+			// replaceRegistry behavior), so compose test verifies the two
+			// env vars don't step on each other on a digest-pinned image.
+			name:     "digest ref + registry override: digest preserved, prefix replaced",
+			image:    "ghcr.io/nvidia/aicr-validators/aiperf-bench@sha256:deadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef",
+			version:  "dev",
+			commit:   "abc1234",
+			tag:      "latest",
+			registry: "localhost:5001",
+			want:     "localhost:5001/aicr-validators/aiperf-bench@sha256:deadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef",
+		},
 	}
 
 	for _, tt := range tests {
@@ -651,6 +753,7 @@ func TestResolveImage(t *testing.T) {
 			} else {
 				t.Setenv("AICR_VALIDATOR_IMAGE_REGISTRY", "")
 			}
+			t.Setenv("AICR_VALIDATOR_IMAGE_TAG", tt.tag)
 			got := ResolveImage(tt.image, tt.version, tt.commit)
 			if got != tt.want {
 				t.Errorf("ResolveImage(%q, %q, %q) = %q, want %q", tt.image, tt.version, tt.commit, got, tt.want)
diff --git a/pkg/validator/job/deployer.go b/pkg/validator/job/deployer.go
@@ -221,10 +221,15 @@ func (d *Deployer) buildEnvApply() []*applycorev1.EnvVarApplyConfiguration {
 	if len(d.tolerations) > 0 {
 		env = append(env, applycorev1.EnvVar().WithName("AICR_TOLERATIONS").WithValue(serializeTolerations(d.tolerations)))
 	}
-	// Forward CLI version and image-registry override so the validator can
-	// resolve images it references outside the catalog (e.g. inference-perf's
-	// aiperf-bench benchmark image) with the same :latest→version pinning and
-	// registry override that catalog.Load applies to catalog entries.
+	// Forward CLI version, image-registry override, and image-tag override so
+	// the validator can resolve images it references outside the catalog
+	// (e.g. inference-perf's aiperf-bench benchmark image) with the same
+	// resolution semantics that catalog.Load applies to catalog entries.
+	// All three must travel together: if a feature-branch dev build set
+	// AICR_VALIDATOR_IMAGE_TAG=latest on the CLI side to get a published
+	// outer validator image, the inner benchmark pod needs the same
+	// override or it will resolve to the same unpublished :sha-<commit>
+	// the outer pod would have hit without the override.
 	if d.cliVersion != "" {
 		env = append(env, applycorev1.EnvVar().WithName("AICR_CLI_VERSION").WithValue(d.cliVersion))
 	}
@@ -234,6 +239,9 @@ func (d *Deployer) buildEnvApply() []*applycorev1.EnvVarApplyConfiguration {
 	if override := os.Getenv("AICR_VALIDATOR_IMAGE_REGISTRY"); override != "" {
 		env = append(env, applycorev1.EnvVar().WithName("AICR_VALIDATOR_IMAGE_REGISTRY").WithValue(override))
 	}
+	if tag := os.Getenv("AICR_VALIDATOR_IMAGE_TAG"); tag != "" {
+		env = append(env, applycorev1.EnvVar().WithName("AICR_VALIDATOR_IMAGE_TAG").WithValue(tag))
+	}
 	for _, e := range d.entry.Env {
 		env = append(env, applycorev1.EnvVar().WithName(e.Name).WithValue(e.Value))
 	}
@@ -281,6 +289,12 @@ func serializeTolerations(tols []corev1.Toleration) string {
 // via `kind load docker-image` and no registry exists to pull from.
 // All other images (including localhost registry) follow the standard policy:
 // :latest tag uses Always to ensure fresh images, versioned tags use IfNotPresent.
+//
+// When AICR_VALIDATOR_IMAGE_TAG is set, the resolved tag is user-supplied and
+// is typically mutable (e.g. `latest`, `edge`, `main` — tags that on-push.yaml
+// recreates on every merge). Treat override-pinned images as mutable and use
+// PullAlways so node-local caches can't serve a stale image instead of the
+// tag's current target.
 func (d *Deployer) imagePullPolicy() corev1.PullPolicy {
 	img := d.entry.Image
 	// Side-loaded images via kind load — no registry exists, never pull.
@@ -289,6 +303,10 @@ func (d *Deployer) imagePullPolicy() corev1.PullPolicy {
 
 		return corev1.PullNever
 	}
+	// Override-pinned tags are assumed mutable — always re-pull.
+	if os.Getenv("AICR_VALIDATOR_IMAGE_TAG") != "" {
+		return corev1.PullAlways
+	}
 	if strings.HasSuffix(img, ":latest") {
 		return corev1.PullAlways
 	}
diff --git a/pkg/validator/job/deployer_test.go b/pkg/validator/job/deployer_test.go
