engine: AX-task executor for kube-ax (engine-task image FROM ax-task-runner v0.3.1, presigned HTTPS I/O, ousast-engine- names, 60 s egress-ready retry, 1.7 GB cluster heap, VM OOM fallback)

Implemented with Codex. Live smoke pending the operator egress watcher.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

benchmarks/learn/engine_task_entry.py

@@ -0,0 +1,154 @@
+"""AX task entry: object-scoped transfers and stdlib only; result PUT is completion."""
+
+from __future__ import annotations
+
+import http.client
+import json
+import os
+import shutil
+import subprocess
+import sys
+import tarfile
+import time
+import urllib.error
+import urllib.request
+from pathlib import Path
+
+
+def transfer(request, destination=None, *, readiness_budget=None):
+    """Wait for actor egress on startup; bound later transfers to three attempts."""
+    deadline = time.monotonic() + readiness_budget if readiness_budget is not None else None
+    delay = 0.5
+    attempt = 0
+    while True:
+        remaining = deadline - time.monotonic() if deadline is not None else 60
+        if remaining <= 0:
+            raise RuntimeError("egress not open after 60 s")
+        attempt += 1
+        try:
+            with urllib.request.urlopen(request, timeout=min(60, remaining)) as response:
+                if destination is not None:
+                    with destination.open("wb") as output:
+                        shutil.copyfileobj(response, output)
+            return
+        except urllib.error.HTTPError as exc:
+            exc.close()
+            if exc.code != 403 and not 500 <= exc.code < 600:
+                raise
+            if deadline is None and attempt >= 3:
+                raise
+        except (urllib.error.URLError, OSError, http.client.HTTPException):
+            if deadline is None and attempt >= 3:
+                raise
+        if deadline is not None:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                raise RuntimeError("egress not open after 60 s") from None
+            time.sleep(min(delay, remaining))
+            delay = min(delay * 2, 8)
+        else:
+            time.sleep(attempt)
+
+
+def download(url, destination, *, readiness_budget=None):
+    transfer(url, destination, readiness_budget=readiness_budget)
+
+
+def unpack(source, destination):
+    """Reject links and traversal; source archives contain only regular files/directories."""
+    with tarfile.open(source) as archive:
+        for member in archive.getmembers():
+            target = (destination / member.name).resolve()
+            if not target.is_relative_to(destination.resolve()) or not (member.isfile() or member.isdir()):
+                raise ValueError("unsafe source archive member")
+        archive.extractall(destination, filter="data")
+
+
+def failure_record(pin, reason):
+    return {
+        **pin,
+        "done": True,
+        "status": "failed",
+        "reason": reason,
+        "units": [
+            {
+                **unit,
+                "status": "failed" if unit.get("supported", True) else "unsupported",
+                "reason": reason,
+                "instrument": {"files": [], "bytes": 0, "jvm": []},
+                "questions": {},
+                "questions_asked": [],
+                "questions_completed": [],
+                "witness_rows": [],
+                "traces": [],
+            }
+            for unit in pin.get("units", [])
+        ],
+    }
+
+
+def main(*, work_root=Path("/tmp"), worker_script=None, readiness_budget=60):
+    pin = {}
+    failed = False
+    output = work_root / "out"
+    result = output / "result.json"
+    try:
+        output.mkdir(parents=True, exist_ok=True)
+        result.unlink(missing_ok=True)
+        questions = work_root / "questions.json"
+        download(os.environ["QUESTIONS_URL"], questions, readiness_budget=readiness_budget)
+        pin = json.loads(questions.read_text())
+        if not isinstance(pin, dict):
+            pin = {}
+            raise ValueError("questions object must contain a prepared pin")
+        source = work_root / "source.tar"
+        download(os.environ["SOURCE_URL"], source)
+        case = work_root / "case"
+        case.mkdir(parents=True, exist_ok=True)
+        unpack(source, case)
+        deadline = int(os.environ.get("DEADLINE", "900"))
+        command = [
+            sys.executable,
+            str(worker_script or Path(__file__).with_name("engine_trace_worker.py")),
+            "--worker",
+            str(questions),
+            "--root",
+            str(case),
+            "--output",
+            str(result),
+            "--heap-profile",
+            "cluster",
+            "--deadline",
+            str(deadline),
+            "--question-deadline",
+            os.environ.get("QUESTION_DEADLINE", "120"),
+        ]
+        completed = subprocess.run(command, timeout=deadline + 5, check=False)
+        if completed.returncode:
+            raise RuntimeError(f"worker exited {completed.returncode}")
+        if not result.is_file() or not json.loads(result.read_text()).get("done"):
+            raise RuntimeError("worker did not write a completed result")
+    except Exception as exc:
+        failed = True
+        # Avoid including exception text containing a presigned URL in durable output.
+        reason = str(exc) if isinstance(exc, (RuntimeError, ValueError)) else f"task entry failed: {type(exc).__name__}"
+        try:
+            output.mkdir(parents=True, exist_ok=True)
+            result.write_text(json.dumps(failure_record(pin, reason)) + "\n")
+        except Exception:
+            print("task entry could not write failure result", file=sys.stderr)
+    try:
+        archive = work_root / "result.tar"
+        with tarfile.open(archive, "w") as tar:
+            for path in sorted(output.iterdir()):
+                tar.add(path, arcname=path.name)
+        data = archive.read_bytes()
+        request = urllib.request.Request(os.environ["RESULT_URL"], data=data, method="PUT", headers={"Content-Type": "application/x-tar"})
+        transfer(request, readiness_budget=readiness_budget if failed else None)
+    except Exception:
+        print("task entry could not upload result; dispatcher deadline remains authoritative", file=sys.stderr)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())