From b44fb56d8db1de23a9d7a3aa2461987f8552b772 Mon Sep 17 00:00:00 2001
From: Marius Ciepluch <11855163+norandom@users.noreply.github.com>
Date: Mon, 5 Oct 2026 20:33:35 +0200
Subject: [PATCH] ax batch: workload-kind prefixes in task names (pass-,
 replay-) [skip ci]

The operator keys per-host egress rules on the task-name prefix; engine passes and
replays had none. Names stay within 49 bytes.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 benchmarks/ax/batch.py              | 8 ++++++--
 benchmarks/learn/engine_trace_ax.py | 2 +-
 benchmarks/unseen/replay.py         | 2 +-
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/benchmarks/ax/batch.py b/benchmarks/ax/batch.py
index 47eeb6d..ecf4fcd 100644
--- a/benchmarks/ax/batch.py
+++ b/benchmarks/ax/batch.py
@@ -262,8 +262,10 @@ class AXLane:
     def _execute_attempt(self, item, output, attempt, attempts):
         pin_id = item.id
         digest = hashlib.sha256(f"{self.run}/{pin_id}/attempt-{attempt}".encode()).hexdigest()[:12]
-        slug = re.sub(r"[^a-z0-9]+", "-", str(pin_id).lower()).strip("-")[:22]
-        name = f"ousast-engine-{slug or 'pin'}-{digest}"
+        kind = re.sub(r"[^a-z0-9]+", "-", self.workload.kind.lower()).strip("-") or "pin"
+        # Names are capped at 49 bytes: 14 (prefix) + kind + 1 + slug + 1 + 12 (digest).
+        slug = re.sub(r"[^a-z0-9]+", "-", str(pin_id).lower()).strip("-")[: max(4, 21 - len(kind))]
+        name = f"ousast-engine-{kind}-{slug or 'pin'}-{digest}"
         key_id = hashlib.sha256(str(pin_id).encode()).hexdigest()
         prefix = f"engine-queue/{self.run}/{key_id}/attempt-{attempt}/"
         keys = [prefix + item.object_names.get(key, key.lower() + PRESIGNED_ARCHIVE_SUFFIX) for key in item.inputs]
@@ -361,6 +363,8 @@ class Workload:
     url_env: frozenset[str]
     deadline: float
     validate: object
+    # Task-name prefix after "ousast-engine-"; the operator's per-host egress rules match on it.
+    kind: str = "pin"
 
 
 def validate_image(image):
diff --git a/benchmarks/learn/engine_trace_ax.py b/benchmarks/learn/engine_trace_ax.py
index 6046f8b..46e8258 100644
--- a/benchmarks/learn/engine_trace_ax.py
+++ b/benchmarks/learn/engine_trace_ax.py
@@ -27,7 +27,7 @@ def task_manifest(name, args, urls):
 
 
 def engine_workload(args, validate):
-    return Workload(args.image, frozenset({"SOURCE_URL", "QUESTIONS_URL", "RESULT_URL"}), args.deadline, validate)
+    return Workload(args.image, frozenset({"SOURCE_URL", "QUESTIONS_URL", "RESULT_URL"}), args.deadline, validate, kind="pass")
 
 
 def engine_item(pin, pin_id, source, args):
diff --git a/benchmarks/unseen/replay.py b/benchmarks/unseen/replay.py
index b5830e5..570abde 100644
--- a/benchmarks/unseen/replay.py
+++ b/benchmarks/unseen/replay.py
@@ -220,7 +220,7 @@ def run_replays(changes, args, out, lane, *, deadline_scale=1.0, deadline=None):
     maximum = max(
         max(json.loads(i.inputs["SPEC_URL"])["deadline_seconds"], json.loads(i.inputs["SPEC_URL"])["base_deadline_seconds"]) for i in items
     )
-    workload = Workload(args.image, frozenset({"SPEC_URL", "RESULT_URL"}), maximum + 600, validate_result)
+    workload = Workload(args.image, frozenset({"SPEC_URL", "RESULT_URL"}), maximum + 600, validate_result, kind="replay")
     lanes = ("ax", "ax", "docker") if lane == "mixed" else (lane,) * (2 if lane in {"ax", "kind"} else 1)
     original = {item.id: change for item, change in zip(items, changes, strict=True)}
     executors = {}