diff --git a/data/cloud/launch_tracegen_iris.py b/data/cloud/launch_tracegen_iris.py index 5b7a949f9..6db1b3fcc 100644 --- a/data/cloud/launch_tracegen_iris.py +++ b/data/cloud/launch_tracegen_iris.py @@ -49,7 +49,10 @@ def _env_vars_from_datagen_yaml(path: Path) -> dict: with path.open() as f: cfg = yaml.safe_load(f) or {} except Exception as e: - print(f"[tracegen-iris] WARNING: failed to parse {path} for env_vars: {e}", file=sys.stderr) + print( + f"[tracegen-iris] WARNING: failed to parse {path} for env_vars: {e}", + file=sys.stderr, + ) return {} env = cfg.get("env_vars") or {} if not isinstance(env, dict): @@ -63,6 +66,7 @@ class TracegenIrisLauncher(IrisLauncher): task_name = "ot-tracegen-iris" job_name_prefix = "tracegen-iris" default_n_concurrent = 64 + enforce_capability_token_duration = True def add_task_specific_args(self, parser: argparse.ArgumentParser) -> None: add_harbor_args(parser, config_required=True) @@ -81,9 +85,14 @@ def add_task_specific_args(self, parser: argparse.ArgumentParser) -> None: legacy_names=["--trace-env", "--trace_env"], ) - parser.add_argument("--datagen_config", required=True, - help="Datagen config with vLLM settings (required).") - parser.add_argument("--datagen-config", dest="datagen_config", help=argparse.SUPPRESS) + parser.add_argument( + "--datagen_config", + required=True, + help="Datagen config with vLLM settings (required).", + ) + parser.add_argument( + "--datagen-config", dest="datagen_config", help=argparse.SUPPRESS + ) add_tasks_input_arg(parser, required=True) @@ -94,17 +103,25 @@ def add_task_specific_args(self, parser: argparse.ArgumentParser) -> None: # not touch). Re-declare just the health knobs here so users can # extend the startup-wait timeout for slow XLA compiles (e.g. # large MoE models on cold-cache iris workers). - parser.add_argument("--health_max_attempts", "--health-max-attempts", - type=int, default=None, - help="Override the worker's run_tracegen.py " - "--health_max_attempts (default 100 on the " - "worker; 100 × --health_retry_delay seconds " - "is the hard cap on cold-start XLA compile + " - "engine warmup). Bump for large MoE models.") - parser.add_argument("--health_retry_delay", "--health-retry-delay", - type=int, default=None, - help="Override the worker's run_tracegen.py " - "--health_retry_delay (default 30s).") + parser.add_argument( + "--health_max_attempts", + "--health-max-attempts", + type=int, + default=None, + help="Override the worker's run_tracegen.py " + "--health_max_attempts (default 100 on the " + "worker; 100 × --health_retry_delay seconds " + "is the hard cap on cold-start XLA compile + " + "engine warmup). Bump for large MoE models.", + ) + parser.add_argument( + "--health_retry_delay", + "--health-retry-delay", + type=int, + default=None, + help="Override the worker's run_tracegen.py " + "--health_retry_delay (default 30s).", + ) # NOTE: --job_name comes from add_harbor_args above. @@ -114,12 +131,16 @@ def add_task_specific_args(self, parser: argparse.ArgumentParser) -> None: add_ingress_literal_args(parser) parser.add_argument( - "--baked-venv", "--baked_venv", dest="baked_venv", action="store_true", + "--baked-venv", + "--baked_venv", + dest="baked_venv", + action="store_true", help="Run the worker from the image's baked /opt/openthoughts/.venv instead of " - "letting Iris run its default setup or rebuild /app/.venv from uv.lock. Required for " - "images whose venv carries deps NOT in the locked pin (e.g. the gpu-glm52 " - "image's fork vLLM for GLM-5.2, which a `uv sync --reinstall` would clobber " - "back to the pinned version).") + "letting Iris run its default setup or rebuild /app/.venv from uv.lock. Required for " + "images whose venv carries deps NOT in the locked pin (e.g. the gpu-glm52 " + "image's fork vLLM for GLM-5.2, which a `uv sync --reinstall` would clobber " + "back to the pinned version).", + ) def normalize_paths(self, args: argparse.Namespace) -> None: # --gpus drives vLLM tensor_parallel_size. On the GPU path it is the whole-node @@ -143,9 +164,13 @@ def normalize_paths(self, args: argparse.Namespace) -> None: # path or raise ValueError if CWD escapes the repo. if not args.tasks_input_path.startswith("/"): if (self.repo_root / args.tasks_input_path).exists(): - args.tasks_input_path = repo_relative(args.tasks_input_path, self.repo_root) + args.tasks_input_path = repo_relative( + args.tasks_input_path, self.repo_root + ) - infer_harbor_env_from_config(args, args.harbor_config, log_prefix="[tracegen-iris]") + infer_harbor_env_from_config( + args, args.harbor_config, log_prefix="[tracegen-iris]" + ) # Resolve per-model serve config from model_config/ (single source of # truth). agent_kwargs are merged + forwarded here; max_model_len / @@ -155,6 +180,7 @@ def normalize_paths(self, args: argparse.Namespace) -> None: # passed (model inferred from the datagen config on the worker), resolution # is deferred to the worker. from hpc.model_config_apply import apply_to_launcher + apply_to_launcher(args, log_prefix="[tracegen-iris]", iris=True) if args.harbor_env == "docker": @@ -170,7 +196,9 @@ def normalize_paths(self, args: argparse.Namespace) -> None: # get_daytona_api_key_override). The same file is also parsed into # the iris worker's env_vars later in run() so the worker sees the # same values. - loaded = self.load_secrets_env_into_os_environ(getattr(args, "secrets_env", None)) + loaded = self.load_secrets_env_into_os_environ( + getattr(args, "secrets_env", None) + ) if loaded: print( f"[tracegen-iris] Secrets: loaded {loaded} entries from " @@ -195,6 +223,7 @@ def normalize_paths(self, args: argparse.Namespace) -> None: maybe_prebuild_daytona_snapshots, ) from hpc.snapshot_manager import OrgConfig + api_key = get_daytona_api_key_override(vars(args)) orgs = [OrgConfig(name="cli", api_key=api_key)] if api_key else [] if not orgs: @@ -202,11 +231,15 @@ def normalize_paths(self, args: argparse.Namespace) -> None: "[tracegen-iris] WARNING: harbor_env=daytona but DAYTONA_API_KEY unset on " "launch host; skipping snapshot pre-build. Expect 'Sandbox not found' " "at trial time unless snapshots are already warm in Daytona.", - file=sys.stderr, flush=True, + file=sys.stderr, + flush=True, ) else: from hpc.hf_utils import is_raw_tasks_directory - resolved_tasks = resolve_dataset_path(args.tasks_input_path, verbose=True) + + resolved_tasks = resolve_dataset_path( + args.tasks_input_path, verbose=True + ) # HF datasets ship a single tasks.parquet; explode into task # directories so snapshot_manager can hash the environments. # hpc/launch.py:232 + hpc/datagen_launch_utils.py do this. @@ -225,24 +258,37 @@ def normalize_paths(self, args: argparse.Namespace) -> None: orgs=orgs, ) - def build_task_command(self, args: argparse.Namespace, remote_output_dir: str) -> List[str]: + def build_task_command( + self, args: argparse.Namespace, remote_output_dir: str + ) -> List[str]: cmd: List[str] = [ - "python", "data/local/run_tracegen.py", - "--harbor_config", args.harbor_config, - "--datagen_config", args.datagen_config, - "--tasks_input_path", args.tasks_input_path, + "python", + "data/local/run_tracegen.py", + "--harbor_config", + args.harbor_config, + "--datagen_config", + args.datagen_config, + "--tasks_input_path", + args.tasks_input_path, ] if args.model: cmd.extend(["--model", args.model]) - cmd.extend([ - "--agent", args.agent, - "--n_concurrent", str(args.n_concurrent), - "--n_attempts", str(args.n_attempts), - "--gpus", str(args.gpus), - "--experiments_dir", remote_output_dir, - ]) + cmd.extend( + [ + "--agent", + args.agent, + "--n_concurrent", + str(args.n_concurrent), + "--n_attempts", + str(args.n_attempts), + "--gpus", + str(args.gpus), + "--experiments_dir", + remote_output_dir, + ] + ) # Health-check passthrough (run_tracegen.py worker-side args). if args.health_max_attempts is not None: @@ -266,7 +312,9 @@ def build_task_command(self, args: argparse.Namespace, remote_output_dir: str) - # the OLD job's name so harbor's _maybe_init_existing_job picks up # the existing config.json / trials. IrisLauncher.run() stashes the # old name on args._harbor_job_name_override. - harbor_job_name = getattr(args, "_harbor_job_name_override", None) or args.job_name + harbor_job_name = ( + getattr(args, "_harbor_job_name_override", None) or args.job_name + ) if harbor_job_name: cmd.extend(["--job_name", harbor_job_name]) if args.dry_run: @@ -295,7 +343,8 @@ def build_task_command(self, args: argparse.Namespace, remote_output_dir: str) - # its editable install. bash = ( "cd /app && export PYTHONPATH=/app:${PYTHONPATH:-} && " - "export PATH=/opt/openthoughts/.venv/bin:$PATH && exec " + shlex.join(cmd) + "export PATH=/opt/openthoughts/.venv/bin:$PATH && exec " + + shlex.join(cmd) ) return ["bash", "-c", bash] diff --git a/eval/cloud/launch_eval_iris.py b/eval/cloud/launch_eval_iris.py index 14dfe3fe1..f3d7b1a04 100644 --- a/eval/cloud/launch_eval_iris.py +++ b/eval/cloud/launch_eval_iris.py @@ -112,6 +112,7 @@ class EvalIrisLauncher(IrisLauncher): task_name = "ot-eval-iris" job_name_prefix = "eval-iris" default_n_concurrent = 16 + enforce_capability_token_duration = True def add_task_specific_args(self, parser: argparse.ArgumentParser) -> None: """Mirror EvalCloudLauncher's args exactly so users don't have to relearn flags.""" diff --git a/hpc/ingress_utils.py b/hpc/ingress_utils.py index ce0ef2323..a23501b68 100644 --- a/hpc/ingress_utils.py +++ b/hpc/ingress_utils.py @@ -19,11 +19,11 @@ dropped and ``/`` -> ``.`` (the exact encoding of ``rigging.connect.capability_path`` / ``proxy_path``); our single-segment ``otagent-`` name encodes to itself. -TOKEN LIFETIME. The controller clamps a minted token to -``MAX_ENDPOINT_TOKEN_TTL_SECONDS`` = 24h (``DEFAULT`` = 1h). The endpoint +TOKEN LIFETIME. The controller clamps a minted token to its own +``MAX_ENDPOINT_TOKEN_TTL_SECONDS`` (``DEFAULT`` = 1h). The endpoint REGISTRATION is separately lease-renewed for the whole run (see :class:`ControllerEndpointRegistration`); only the token expires. So the api_base -is resolved through :func:`capability_api_base`, which mints a 24h token, caches +is resolved through :func:`capability_api_base`, which requests that maximum, caches it worker-side keyed by endpoint name, and re-mints when within ``TOKEN_REFRESH_MARGIN_SECONDS`` of expiry. @@ -34,7 +34,7 @@ refreshes the token across harbor RE-SPAWNS (resume / campaign refills), not across trials within one running harbor process. A harbor run that stays up longer than the token TTL will outlive its token — keep individual harbor runs -under 24h, or re-spawn to re-mint. There is no per-trial base_url resolution hook +within the controller maximum, or re-spawn to re-mint. There is no per-trial base_url resolution hook in the current OT-Agent->harbor plumbing. """ @@ -47,6 +47,8 @@ from dataclasses import dataclass from typing import Callable, Dict, Optional, Protocol, Tuple +from hpc.iris.capability_tokens import controller_max_endpoint_token_ttl_seconds + # The sandbox-facing api_key. The capability token rides in the URL path, so no # bearer is needed; but installed OpenAI-compatible agents refuse to start # without SOME non-empty key, so we hand them this inert placeholder. It is @@ -77,10 +79,8 @@ # The raw vLLM HTTP port the RL/datagen servers bind on the task node. DEFAULT_VLLM_PORT = 8000 -# Token TTL we request when minting (clamped server-side to the controller's -# MAX_ENDPOINT_TOKEN_TTL_SECONDS, currently 24h) and the safety margin at which a -# cached token is re-minted rather than reused. -DEFAULT_TOKEN_TTL_HOURS = 24.0 +# Safety margin at which a cached token is re-minted rather than reused. The +# request lifetime itself is resolved from the controller at runtime. TOKEN_REFRESH_MARGIN_SECONDS = 2 * 3600 # re-mint when <2h remains @@ -171,9 +171,17 @@ class CapabilityTokenCache: and lease renewers touch it from different threads. """ - def __init__(self, minter: CapabilityMinter, *, ttl_hours: float = DEFAULT_TOKEN_TTL_HOURS) -> None: + def __init__( + self, minter: CapabilityMinter, *, ttl_hours: float | None = None + ) -> None: self._minter = minter - self._ttl_hours = ttl_hours + # Read the controller-owned maximum lazily instead of maintaining an + # OT-Agent copy. Tests can still inject an explicit TTL. + self._ttl_hours = ( + ttl_hours + if ttl_hours is not None + else controller_max_endpoint_token_ttl_seconds() / 3600.0 + ) self._lock = threading.Lock() self._cache: Dict[str, _CachedToken] = {} @@ -181,7 +189,10 @@ def token_for(self, endpoint_name: str, *, now: Optional[float] = None) -> str: now = time.time() if now is None else now with self._lock: cached = self._cache.get(endpoint_name) - if cached is not None and cached.expires_at - now > TOKEN_REFRESH_MARGIN_SECONDS: + if ( + cached is not None + and cached.expires_at - now > TOKEN_REFRESH_MARGIN_SECONDS + ): return cached.token token, expires_at = self._minter.mint(endpoint_name, self._ttl_hours) if not token: @@ -189,7 +200,9 @@ def token_for(self, endpoint_name: str, *, now: Optional[float] = None) -> str: f"minting a capability token for {endpoint_name} returned an empty " "token; refusing to build an unreachable api_base." ) - self._cache[endpoint_name] = _CachedToken(token=token, expires_at=expires_at) + self._cache[endpoint_name] = _CachedToken( + token=token, expires_at=expires_at + ) return token @@ -372,7 +385,9 @@ def register( ) -> str: from iris.cluster.types import EndpointAccess - access_mode = access if access is not None else EndpointAccess.ENDPOINT_ACCESS_LINK + access_mode = ( + access if access is not None else EndpointAccess.ENDPOINT_ACCESS_LINK + ) return self._client.register( name, address, self._task_attempt, metadata or {}, access=access_mode ) @@ -596,13 +611,17 @@ def __init__( minter: CapabilityMinter, resolver: ParentEndpointResolver, *, - ttl_hours: float = DEFAULT_TOKEN_TTL_HOURS, + ttl_hours: float | None = None, mirror_timeout_s: float = DEFAULT_MIRROR_TIMEOUT_SECONDS, mirror_interval_s: float = DEFAULT_MIRROR_POLL_INTERVAL_SECONDS, ) -> None: self._minter = minter self._resolver = resolver - self._ttl_hours = ttl_hours + self._ttl_hours = ( + ttl_hours + if ttl_hours is not None + else controller_max_endpoint_token_ttl_seconds() / 3600.0 + ) self._mirror_timeout_s = mirror_timeout_s self._mirror_interval_s = mirror_interval_s self._lock = threading.Lock() @@ -613,7 +632,10 @@ def token_for(self, endpoint_name: str, *, now: Optional[float] = None) -> str: now = time.time() if now is None else now with self._lock: cached = self._cache.get(endpoint_name) - if cached is not None and cached.expires_at - now > TOKEN_REFRESH_MARGIN_SECONDS: + if ( + cached is not None + and cached.expires_at - now > TOKEN_REFRESH_MARGIN_SECONDS + ): return cached.token state = self._state.setdefault(endpoint_name, _FederatedTokenState()) if not state.mirrored: @@ -630,7 +652,9 @@ def token_for(self, endpoint_name: str, *, now: Optional[float] = None) -> str: f"parent-minting a capability token for {endpoint_name} returned an " "empty token; refusing to build an unreachable api_base." ) - self._cache[endpoint_name] = _CachedToken(token=token, expires_at=expires_at) + self._cache[endpoint_name] = _CachedToken( + token=token, expires_at=expires_at + ) return token @@ -664,7 +688,9 @@ def is_mirrored(self, endpoint_name: str) -> bool: from iris.rpc import controller_pb2 resp = self._client.list_endpoints( - controller_pb2.Controller.ListEndpointsRequest(prefix=endpoint_name, exact=True) + controller_pb2.Controller.ListEndpointsRequest( + prefix=endpoint_name, exact=True + ) ) # A mirrored (federated) row carries a non-empty peer_id; a purely-local # parent endpoint of the same name (there should be none) would not. @@ -739,7 +765,11 @@ def federated_capability_api_base( it MUST be the marin host (a peer-signed token 401s at iris.oa.dev). ``cache`` is injectable for tests; production uses the process-wide parent-authenticated cache. """ - host = ingress_host or os.environ.get(PARENT_INGRESS_HOST_ENV) or DEFAULT_PARENT_INGRESS_HOST + host = ( + ingress_host + or os.environ.get(PARENT_INGRESS_HOST_ENV) + or DEFAULT_PARENT_INGRESS_HOST + ) token_cache = cache if cache is not None else _default_federated_token_cache() token = token_cache.token_for(endpoint_name, now=now) return build_capability_api_base(host, endpoint_name, token) diff --git a/hpc/iris/capability_tokens.py b/hpc/iris/capability_tokens.py new file mode 100644 index 000000000..eacebf0f6 --- /dev/null +++ b/hpc/iris/capability_tokens.py @@ -0,0 +1,162 @@ +"""Controller capability-token policy shared by Iris datagen and eval launchers. + +The controller owns the maximum lifetime for endpoint-scoped proxy tokens. Do +not duplicate its value here: an image/controller update may safely change the +limit without an OT-Agent source edit. +""" + +from __future__ import annotations + +import importlib +import json +import os +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Callable + +import yaml + +from hpc.local_paths import PATHS, ensure as ensure_local_paths + + +_CONTROLLER_AUTH_MODULE = "iris.cluster.controller.auth" +_MAX_TTL_ATTRIBUTE = "MAX_ENDPOINT_TOKEN_TTL_SECONDS" +_POLICY_SCHEMA_VERSION = 1 + + +def is_token_required_agent(agent: str | None) -> bool: + """Whether an agent must call a capability URL from outside the Iris VPC.""" + return bool(agent) and agent.strip().lower() != "terminus-2" + + +def infer_harbor_agent( + agent: str | None, harbor_config: str | Path | None +) -> str | None: + """Use the configured first Harbor agent when the CLI omitted ``--agent``.""" + if agent: + return agent + if not harbor_config: + return None + try: + loaded = yaml.safe_load(Path(harbor_config).read_text()) or {} + except (OSError, yaml.YAMLError) as exc: + raise ValueError( + f"Could not read Harbor config {harbor_config!s}: {exc}" + ) from exc + agents = loaded.get("agents") if isinstance(loaded, dict) else None + first = agents[0] if isinstance(agents, list) and agents else None + name = first.get("name") if isinstance(first, dict) else None + if not isinstance(name, str) or not name.strip(): + raise ValueError( + "Token-duration enforcement needs --agent or Harbor config agents[0].name." + ) + return name + + +def controller_max_endpoint_token_ttl_seconds( + *, module_loader: Callable[[str], object] = importlib.import_module +) -> int: + """Read the authoritative maximum directly from the installed Marin codebase.""" + module = module_loader(_CONTROLLER_AUTH_MODULE) + value = getattr(module, _MAX_TTL_ATTRIBUTE, None) + if not isinstance(value, int) or value <= 0: + raise RuntimeError( + f"{_CONTROLLER_AUTH_MODULE}.{_MAX_TTL_ATTRIBUTE} must be a positive integer; " + f"got {value!r}." + ) + return value + + +@dataclass(frozen=True) +class CapabilityTokenDurationPolicy: + """Secret-free, persisted launch policy for a single Iris workload.""" + + schema_version: int + agent: str | None + token_required: bool + controller_max_ttl_seconds: int | None + requested_timeout_seconds: int + effective_timeout_seconds: int + requested_token_ttl_seconds: int | None + effective_token_ttl_seconds: int | None + controller_auth_module: str | None + + def to_dict(self) -> dict[str, int | str | bool | None]: + return asdict(self) + + +def resolve_token_duration_policy( + *, + agent: str | None, + timeout_seconds: int, + requested_token_ttl_seconds: int | None = None, + max_ttl_resolver: Callable[[], int] = controller_max_endpoint_token_ttl_seconds, +) -> CapabilityTokenDurationPolicy: + """Bound a token-required launch to its actual minted-token lifetime. + + ``timeout_seconds=0`` is Iris's historical "unlimited" spelling. It is + safe only for an in-process harness; token-required jobs receive the + controller maximum by default. A caller requesting a shorter token gets a + correspondingly shorter job bound. + """ + if timeout_seconds < 0: + raise ValueError("Iris timeout must be zero or a positive number of seconds.") + required = is_token_required_agent(agent) + if not required: + return CapabilityTokenDurationPolicy( + schema_version=_POLICY_SCHEMA_VERSION, + agent=agent, + token_required=False, + controller_max_ttl_seconds=None, + requested_timeout_seconds=timeout_seconds, + effective_timeout_seconds=timeout_seconds, + requested_token_ttl_seconds=None, + effective_token_ttl_seconds=None, + controller_auth_module=None, + ) + + maximum = max_ttl_resolver() + if requested_token_ttl_seconds is not None and requested_token_ttl_seconds <= 0: + raise ValueError("Requested capability-token TTL must be positive.") + if ( + requested_token_ttl_seconds is not None + and requested_token_ttl_seconds > maximum + ): + raise ValueError( + "Requested capability-token TTL " + f"({requested_token_ttl_seconds}s) exceeds Marin controller maximum ({maximum}s)." + ) + effective_token_ttl = requested_token_ttl_seconds or maximum + effective_timeout = timeout_seconds or effective_token_ttl + if effective_timeout > effective_token_ttl: + raise ValueError( + "Token-required Iris jobs cannot outlive their minted endpoint token: " + f"requested timeout={effective_timeout}s, token TTL={effective_token_ttl}s." + ) + return CapabilityTokenDurationPolicy( + schema_version=_POLICY_SCHEMA_VERSION, + agent=agent, + token_required=True, + controller_max_ttl_seconds=maximum, + requested_timeout_seconds=timeout_seconds, + effective_timeout_seconds=effective_timeout, + requested_token_ttl_seconds=requested_token_ttl_seconds, + effective_token_ttl_seconds=effective_token_ttl, + controller_auth_module=_CONTROLLER_AUTH_MODULE, + ) + + +def persist_token_duration_policy( + *, job_name: str, policy: CapabilityTokenDurationPolicy, root: Path | None = None +) -> Path: + """Write the launch policy in the managed local control-plane state tree.""" + state_root = root or (PATHS.state / "iris_capability_token_policies") + if root is None: + ensure_local_paths(PATHS.home, PATHS.state, state_root) + else: + state_root.mkdir(parents=True, exist_ok=True) + destination = state_root / f"{job_name}.json" + temporary = destination.with_suffix(".json.tmp") + temporary.write_text(json.dumps(policy.to_dict(), sort_keys=True, indent=2) + "\n") + os.replace(temporary, destination) + return destination diff --git a/hpc/iris/launcher.py b/hpc/iris/launcher.py index 8d15c47ce..3b38abbce 100644 --- a/hpc/iris/launcher.py +++ b/hpc/iris/launcher.py @@ -29,6 +29,7 @@ from __future__ import annotations import argparse +import json import os import shlex import sys @@ -39,13 +40,22 @@ from hpc.local_paths import PATHS as LOCAL_PATHS, ensure as ensure_local_paths from hpc.iris_job_registry import register_submission, get_latest_by_job_name -from hpc.iris.accelerator import DEFAULT_IRIS_JOB_API, IrisJobApi, ResolvedIrisAccelerator +from hpc.iris.accelerator import ( + DEFAULT_IRIS_JOB_API, + IrisJobApi, + ResolvedIrisAccelerator, +) from hpc.iris.bootstrap import wrap_task_command from hpc.iris.env import ( apply_iris_runtime_env, default_secrets_env, load_secrets_env_into_os_environ, ) +from hpc.iris.capability_tokens import ( + infer_harbor_agent, + persist_token_duration_policy, + resolve_token_duration_policy, +) from hpc.iris.outputs import ( DEFAULT_GCS_OUTPUT_ROOT, DEFAULT_LOCAL_OUTPUT_ROOT, @@ -91,6 +101,7 @@ class IrisLauncher: # Users may still pass --harbor_env docker; iris workers don't mount # /var/run/docker.sock so the job will fail at runtime — by design. default_harbor_env: str = "daytona" + enforce_capability_token_duration: bool = False def __init__(self, repo_root: Path, iris_api: IrisJobApi = DEFAULT_IRIS_JOB_API): self.repo_root = Path(repo_root).resolve() @@ -117,117 +128,195 @@ def create_argument_parser(self, description: str = "") -> argparse.ArgumentPars def _add_iris_common_args(self, parser: argparse.ArgumentParser) -> None: g = parser.add_argument_group("iris") - g.add_argument("--cluster-config", "--cluster_config", - default=None, - help="Path to the iris cluster YAML (default: marin for TPU, " - "cw-us-east-02a for GPU, resolved in the marin repo).") - g.add_argument("--cluster", default=None, - help="Named iris cluster: resolves lib/iris/config/.yaml in " - "the marin repo (e.g. --cluster cw-rno2a). Convenience over " - "--cluster-config; --cluster-config wins if both are given.") - g.add_argument("--task-image", "--task_image", - default=None, - help="Container image for the task (default: " - f"{DEFAULT_TASK_IMAGE} for TPU, {DEFAULT_GPU_TASK_IMAGE} for GPU).") - g.add_argument("--tpu", default=None, - help=f"TPU variant (default: {self.default_tpu} when --gpu is omitted).") - g.add_argument("--gpu", default=None, - help="GPU variant in Iris format, e.g. H100x8. Mutually " - "exclusive with --tpu (selects the CoreWeave GPU path).") - g.add_argument("--replicas", type=int, default=1, - help="Replica count passed to iris submit (default 1). " - "For a multi-host TPU slice iris REQUIRES one task " - "per VM and auto-scales replicas=1 -> vm_count " - "(iris adjust_tpu_replicas); those tasks form ONE " - "JAX device mesh (jax.distributed.initialize) that a " - "single engine spans — this is required, not " - "duplication. Use N*vm_count to request N slices. " - "NOTE: run_tracegen currently runs harbor on EVERY " - "task; the per-host duplication is the harbor layer, " - "not the replica count (fix = gate harbor to the " - "driver rank, see #69).") - g.add_argument("--cpu", type=float, default=8.0, - help="CPU cores for the entrypoint task (default 8).") - g.add_argument("--memory", default="256GB", - help="Memory for the entrypoint task (default 256GB). " - "v6e workers have 720GB total, so 256GB covers " - "HF weight loading for models up to ~120B bf16 " - "or ~400B AWQ-4-bit with comfortable headroom. " - "Bump for larger models; drop to 64GB for small " - "smokes if you want to be polite to the queue.") - g.add_argument("--disk", default=None, - help=f"Ephemeral disk (default {DEFAULT_DISK} on marin TPU, " - f"{DEFAULT_GPU_DISK} on CoreWeave GPU). marin's v6e/v5p " - "workers cap per-VM disk at 100GB; requests above that " - "queue forever waiting on an autoscaler that can't " - "provision a larger-disk worker (on TPU, for weights " - ">100GB use --load-format runai_streamer + gs://-hosted " - "weights instead of bumping disk). CoreWeave GPU nodes " - "have large local disk, so the GPU default is higher to " - "fit big model weight downloads without an ephemeral-" - "storage eviction.") - g.add_argument("--priority", default=DEFAULT_PRIORITY, - choices=["production", "interactive", "batch"], - help="Iris priority band (default interactive).") - g.add_argument("--max-retries", "--max_retries", type=int, default=0, - help="Max retries on failure (does NOT cover preemption — iris retries " - "preemptions automatically up to its own limit).") - g.add_argument("--timeout", type=int, default=0, - help="Job timeout in seconds (0 = no timeout).") - g.add_argument("--preemptible", dest="preemptible", action="store_true", default=None, - help="Force scheduling on preemptible workers (overrides iris heuristic).") - g.add_argument("--no-preemptible", dest="preemptible", action="store_false", - help="Force scheduling on non-preemptible workers.") - g.add_argument("--no-wait", dest="no_wait", action="store_true", default=False, - help="Submit and detach instead of streaming logs.") - g.add_argument("--extras", action="append", default=None, - help="OpenThoughts-Agent extras to install in the iris worker's " - "/app/.venv via `uv sync --extra `. Repeatable. " - "Default: ['datagen-tpu'] for TPU and ['datagen'] for GPU. " - "Pass --extras '' to install no extras.") + g.add_argument( + "--cluster-config", + "--cluster_config", + default=None, + help="Path to the iris cluster YAML (default: marin for TPU, " + "cw-us-east-02a for GPU, resolved in the marin repo).", + ) + g.add_argument( + "--cluster", + default=None, + help="Named iris cluster: resolves lib/iris/config/.yaml in " + "the marin repo (e.g. --cluster cw-rno2a). Convenience over " + "--cluster-config; --cluster-config wins if both are given.", + ) + g.add_argument( + "--task-image", + "--task_image", + default=None, + help="Container image for the task (default: " + f"{DEFAULT_TASK_IMAGE} for TPU, {DEFAULT_GPU_TASK_IMAGE} for GPU).", + ) + g.add_argument( + "--tpu", + default=None, + help=f"TPU variant (default: {self.default_tpu} when --gpu is omitted).", + ) + g.add_argument( + "--gpu", + default=None, + help="GPU variant in Iris format, e.g. H100x8. Mutually " + "exclusive with --tpu (selects the CoreWeave GPU path).", + ) + g.add_argument( + "--replicas", + type=int, + default=1, + help="Replica count passed to iris submit (default 1). " + "For a multi-host TPU slice iris REQUIRES one task " + "per VM and auto-scales replicas=1 -> vm_count " + "(iris adjust_tpu_replicas); those tasks form ONE " + "JAX device mesh (jax.distributed.initialize) that a " + "single engine spans — this is required, not " + "duplication. Use N*vm_count to request N slices. " + "NOTE: run_tracegen currently runs harbor on EVERY " + "task; the per-host duplication is the harbor layer, " + "not the replica count (fix = gate harbor to the " + "driver rank, see #69).", + ) + g.add_argument( + "--cpu", + type=float, + default=8.0, + help="CPU cores for the entrypoint task (default 8).", + ) + g.add_argument( + "--memory", + default="256GB", + help="Memory for the entrypoint task (default 256GB). " + "v6e workers have 720GB total, so 256GB covers " + "HF weight loading for models up to ~120B bf16 " + "or ~400B AWQ-4-bit with comfortable headroom. " + "Bump for larger models; drop to 64GB for small " + "smokes if you want to be polite to the queue.", + ) + g.add_argument( + "--disk", + default=None, + help=f"Ephemeral disk (default {DEFAULT_DISK} on marin TPU, " + f"{DEFAULT_GPU_DISK} on CoreWeave GPU). marin's v6e/v5p " + "workers cap per-VM disk at 100GB; requests above that " + "queue forever waiting on an autoscaler that can't " + "provision a larger-disk worker (on TPU, for weights " + ">100GB use --load-format runai_streamer + gs://-hosted " + "weights instead of bumping disk). CoreWeave GPU nodes " + "have large local disk, so the GPU default is higher to " + "fit big model weight downloads without an ephemeral-" + "storage eviction.", + ) + g.add_argument( + "--priority", + default=DEFAULT_PRIORITY, + choices=["production", "interactive", "batch"], + help="Iris priority band (default interactive).", + ) + g.add_argument( + "--max-retries", + "--max_retries", + type=int, + default=0, + help="Max retries on failure (does NOT cover preemption — iris retries " + "preemptions automatically up to its own limit).", + ) + g.add_argument( + "--timeout", + type=int, + default=0, + help="Job timeout in seconds (0 = no timeout).", + ) + g.add_argument( + "--preemptible", + dest="preemptible", + action="store_true", + default=None, + help="Force scheduling on preemptible workers (overrides iris heuristic).", + ) + g.add_argument( + "--no-preemptible", + dest="preemptible", + action="store_false", + help="Force scheduling on non-preemptible workers.", + ) + g.add_argument( + "--no-wait", + dest="no_wait", + action="store_true", + default=False, + help="Submit and detach instead of streaming logs.", + ) + g.add_argument( + "--extras", + action="append", + default=None, + help="OpenThoughts-Agent extras to install in the iris worker's " + "/app/.venv via `uv sync --extra `. Repeatable. " + "Default: ['datagen-tpu'] for TPU and ['datagen'] for GPU. " + "Pass --extras '' to install no extras.", + ) og = parser.add_argument_group("outputs") - og.add_argument("--output-mode", "--output_mode", - choices=["auto", "gcs", "s3", "local"], - default=os.environ.get("OT_AGENT_OUTPUT_MODE", "auto"), - help="Where workload outputs are written. 'auto' uses GCS for TPU " - "and pod-local for GPU (Harbor writes trace_jobs to local NVMe, " - "run_eval registers to Supabase/HF in-pod). 'gcs' writes to " - "--gcs-output-dir/ (TPU); 's3' writes durable Harbor " - "artifacts to --s3-output-dir/ (CoreWeave R2); 'local' " - "writes to --local-output-dir/ (GPU, default).") - og.add_argument("--gcs-output-dir", "--gcs_output_dir", - default=os.environ.get("OT_AGENT_GCS_OUTPUT_ROOT", DEFAULT_GCS_OUTPUT_ROOT), - help=f"GCS prefix for workload outputs; workload writes to " - f"//. Defaults to $OT_AGENT_GCS_OUTPUT_ROOT or " - f"{DEFAULT_GCS_OUTPUT_ROOT}. The fetch daemon " - f"(hpc.iris_fetch_daemon) pulls completed jobs from here into " - f"{LOCAL_PATHS.runs}//.") - og.add_argument("--s3-output-dir", "--s3_output_dir", - default=os.environ.get("OT_AGENT_S3_OUTPUT_ROOT", DEFAULT_S3_OUTPUT_ROOT), - help="S3-compatible prefix for durable GPU outputs, e.g. " - "s3://marin-us-east-02a/tmp/ttl=7d/ot-agent/evals/. Defaults " - "to $OT_AGENT_S3_OUTPUT_ROOT. The pod uses the cluster-injected " - "creds+endpoint (iris-task-env envFrom); launch-host storage creds " - "are withheld so they cannot clobber them. NOTE: default store moved " - "R2 (s3://marin-na) -> CW (s3://marin-us-east-02a) 2026-07-05 (marin " - "c7caecc95a); s3://marin-na is no longer reachable from pods.") - og.add_argument("--local-output-dir", "--local_output_dir", - default=os.environ.get("OT_AGENT_LOCAL_OUTPUT_ROOT", DEFAULT_LOCAL_OUTPUT_ROOT), - help="Pod-local runtime scratch root used with --output-mode local/s3. " - f"Defaults to {DEFAULT_LOCAL_OUTPUT_ROOT}.") + og.add_argument( + "--output-mode", + "--output_mode", + choices=["auto", "gcs", "s3", "local"], + default=os.environ.get("OT_AGENT_OUTPUT_MODE", "auto"), + help="Where workload outputs are written. 'auto' uses GCS for TPU " + "and pod-local for GPU (Harbor writes trace_jobs to local NVMe, " + "run_eval registers to Supabase/HF in-pod). 'gcs' writes to " + "--gcs-output-dir/ (TPU); 's3' writes durable Harbor " + "artifacts to --s3-output-dir/ (CoreWeave R2); 'local' " + "writes to --local-output-dir/ (GPU, default).", + ) + og.add_argument( + "--gcs-output-dir", + "--gcs_output_dir", + default=os.environ.get("OT_AGENT_GCS_OUTPUT_ROOT", DEFAULT_GCS_OUTPUT_ROOT), + help=f"GCS prefix for workload outputs; workload writes to " + f"//. Defaults to $OT_AGENT_GCS_OUTPUT_ROOT or " + f"{DEFAULT_GCS_OUTPUT_ROOT}. The fetch daemon " + f"(hpc.iris_fetch_daemon) pulls completed jobs from here into " + f"{LOCAL_PATHS.runs}//.", + ) + og.add_argument( + "--s3-output-dir", + "--s3_output_dir", + default=os.environ.get("OT_AGENT_S3_OUTPUT_ROOT", DEFAULT_S3_OUTPUT_ROOT), + help="S3-compatible prefix for durable GPU outputs, e.g. " + "s3://marin-us-east-02a/tmp/ttl=7d/ot-agent/evals/. Defaults " + "to $OT_AGENT_S3_OUTPUT_ROOT. The pod uses the cluster-injected " + "creds+endpoint (iris-task-env envFrom); launch-host storage creds " + "are withheld so they cannot clobber them. NOTE: default store moved " + "R2 (s3://marin-na) -> CW (s3://marin-us-east-02a) 2026-07-05 (marin " + "c7caecc95a); s3://marin-na is no longer reachable from pods.", + ) + og.add_argument( + "--local-output-dir", + "--local_output_dir", + default=os.environ.get( + "OT_AGENT_LOCAL_OUTPUT_ROOT", DEFAULT_LOCAL_OUTPUT_ROOT + ), + help="Pod-local runtime scratch root used with --output-mode local/s3. " + f"Defaults to {DEFAULT_LOCAL_OUTPUT_ROOT}.", + ) rg = parser.add_argument_group("resume") - rg.add_argument("--resume-from", "--resume_from", dest="resume_from", default=None, - help="Resume harbor state from a previously-submitted iris job " - "(by job_name; looked up in the local registry " - f"at {LOCAL_PATHS.state}/iris_jobs.db). The new iris job gets a " - "fresh timestamped name (for iris-level uniqueness), but the " - "harbor --job_name and --jobs-dir are routed at the old job's " - "GCS path so harbor's _maybe_init_existing_job picks up the " - "existing trial results and only runs the unmatched remaining " - "trials. No config gating: per the user's direction (2026-05-24), " - "OT-Agent and harbor already validate compatibility on resume.") + rg.add_argument( + "--resume-from", + "--resume_from", + dest="resume_from", + default=None, + help="Resume harbor state from a previously-submitted iris job " + "(by job_name; looked up in the local registry " + f"at {LOCAL_PATHS.state}/iris_jobs.db). The new iris job gets a " + "fresh timestamped name (for iris-level uniqueness), but the " + "harbor --job_name and --jobs-dir are routed at the old job's " + "GCS path so harbor's _maybe_init_existing_job picks up the " + "existing trial results and only runs the unmatched remaining " + "trials. No config gating: per the user's direction (2026-05-24), " + "OT-Agent and harbor already validate compatibility on resume.", + ) sg = parser.add_argument_group("secrets") # Default to $OT_AGENT_SECRETS_ENV, then ~/Documents/secrets.env if it @@ -238,21 +327,27 @@ def _add_iris_common_args(self, parser: argparse.ArgumentParser) -> None: # auto_snapshot path to fail with "Sandbox not found" even when the # snapshot is ACTIVE on the right org. _default_secrets = default_secrets_env() - sg.add_argument("--secrets-env", "--secrets_env", default=_default_secrets, - help="Path to a KEY=VALUE env file (~/Documents/secrets.env style). " - "Every entry is loaded into the iris task's env_vars at submit " - "time. Pairs with the hardcoded launcher passthrough list " - "(DAYTONA_API_KEY, OPENAI_API_KEY, etc.) — file values win on " - "conflict, explicit `-e` iris-CLI flags can't override since we " - "use IrisClient.submit() directly. Lines starting with '#' and " - "blank lines are ignored; leading 'export ' is stripped. " - "Defaults to $OT_AGENT_SECRETS_ENV, else ~/Documents/secrets.env " - "if it exists.") + sg.add_argument( + "--secrets-env", + "--secrets_env", + default=_default_secrets, + help="Path to a KEY=VALUE env file (~/Documents/secrets.env style). " + "Every entry is loaded into the iris task's env_vars at submit " + "time. Pairs with the hardcoded launcher passthrough list " + "(DAYTONA_API_KEY, OPENAI_API_KEY, etc.) — file values win on " + "conflict, explicit `-e` iris-CLI flags can't override since we " + "use IrisClient.submit() directly. Lines starting with '#' and " + "blank lines are ignored; leading 'export ' is stripped. " + "Defaults to $OT_AGENT_SECRETS_ENV, else ~/Documents/secrets.env " + "if it exists.", + ) # NOTE: --dry-run / --dry_run is provided by hpc.arg_groups.add_model_compute_args # which subclass launchers call from add_task_specific_args. We don't redeclare # it here to avoid argparse conflicts. - def _resolve_cluster_config_default(self, default_config: str = DEFAULT_CLUSTER_CONFIG) -> str: + def _resolve_cluster_config_default( + self, default_config: str = DEFAULT_CLUSTER_CONFIG + ) -> str: """Find the marin repo's cluster config relative to common locations.""" candidates = [ Path.home() / "Documents/marin" / default_config, @@ -280,7 +375,9 @@ def load_secrets_env_into_os_environ(secrets_env: Optional[str]) -> int: def normalize_paths(self, args: argparse.Namespace) -> None: """Subclass hook: validate/normalize paths and infer defaults.""" - def build_task_command(self, args: argparse.Namespace, remote_output_dir: str) -> List[str]: + def build_task_command( + self, args: argparse.Namespace, remote_output_dir: str + ) -> List[str]: """Subclass hook: build the ``python data/...py ...`` invocation.""" raise NotImplementedError @@ -293,7 +390,9 @@ def build_env(self, args: argparse.Namespace) -> dict: """ return {} - def pre_submit_precache(self, args: argparse.Namespace, *, remote_output_dir: str) -> dict: + def pre_submit_precache( + self, args: argparse.Namespace, *, remote_output_dir: str + ) -> dict: """Subclass hook: pre-cache artifacts + return extra env, after region pin. Runs in ``run()`` AFTER the region pin (so ``args._pinned_region`` is @@ -316,7 +415,9 @@ def _derive_job_name(self, args: argparse.Namespace) -> str: ts = time.strftime("%Y%m%d-%H%M%S") return f"{self.job_name_prefix}-{ts}" - def _normalize_accelerator_args(self, args: argparse.Namespace) -> ResolvedIrisAccelerator: + def _normalize_accelerator_args( + self, args: argparse.Namespace + ) -> ResolvedIrisAccelerator: """Resolve the default TPU vs explicit GPU accelerator choice.""" accelerator = ResolvedIrisAccelerator.from_args( args, @@ -335,7 +436,9 @@ def _apply_accelerator_defaults( accelerator: ResolvedIrisAccelerator, ) -> None: if args.task_image is None: - args.task_image = DEFAULT_GPU_TASK_IMAGE if accelerator.is_gpu else DEFAULT_TASK_IMAGE + args.task_image = ( + DEFAULT_GPU_TASK_IMAGE if accelerator.is_gpu else DEFAULT_TASK_IMAGE + ) elif accelerator.is_gpu and args.task_image == DEFAULT_TASK_IMAGE: raise SystemExit( "--gpu requires a GPU task image. Omit --task-image to use " @@ -348,7 +451,11 @@ def _apply_accelerator_defaults( ) if args.cluster_config is None: - config = DEFAULT_GPU_CLUSTER_CONFIG if accelerator.is_gpu else DEFAULT_CLUSTER_CONFIG + config = ( + DEFAULT_GPU_CLUSTER_CONFIG + if accelerator.is_gpu + else DEFAULT_CLUSTER_CONFIG + ) args.cluster_config = self._resolve_cluster_config_default(config) if args.disk is None: @@ -363,6 +470,19 @@ def resolved_accelerator(self, args: argparse.Namespace) -> ResolvedIrisAccelera def run(self, args: argparse.Namespace) -> int: accelerator = self._normalize_accelerator_args(args) self.normalize_paths(args) + if self.enforce_capability_token_duration: + try: + args.agent = infer_harbor_agent( + getattr(args, "agent", None), getattr(args, "harbor_config", None) + ) + args._capability_token_duration_policy = resolve_token_duration_policy( + agent=args.agent, timeout_seconds=args.timeout + ) + except ValueError as exc: + raise SystemExit(str(exc)) from exc + args.timeout = ( + args._capability_token_duration_policy.effective_timeout_seconds + ) accelerator = self.resolved_accelerator(args) output_mode = resolve_output_mode(args, accelerator_kind=accelerator.kind) @@ -411,7 +531,9 @@ def run(self, args: argparse.Namespace) -> int: and accelerator.is_tpu ): try: - region, rows = discover_region_for_tpu(args.cluster_config, accelerator.primary_tpu) + region, rows = discover_region_for_tpu( + args.cluster_config, accelerator.primary_tpu + ) except Exception as exc: raise RuntimeError( f"Region discovery failed ({exc}). Refusing to fall back to the static " @@ -435,7 +557,8 @@ def run(self, args: argparse.Namespace) -> int: summary = ", ".join( f"{r['region']}: {r.get('unassigned', 0)} warm / " f"{r.get('total', 0)} total" - for r in rows if r.get('region') + for r in rows + if r.get("region") ) print( f"[iris] Region pin: --tpu={accelerator.primary_tpu} → {region} " @@ -464,8 +587,13 @@ def run(self, args: argparse.Namespace) -> int: # We don't auto-rewrite — that masks broken mirrors. We refuse to # submit and tell the user exactly which field is wrong. if args._pinned_region: - yaml_attrs = ("datagen_config", "harbor_config", "eval_config", - "config", "harbor_yaml") + yaml_attrs = ( + "datagen_config", + "harbor_config", + "eval_config", + "config", + "harbor_yaml", + ) yaml_paths = [ Path(getattr(args, attr)) for attr in yaml_attrs @@ -483,7 +611,9 @@ def run(self, args: argparse.Namespace) -> int: if resume_target: prev = resume_prev ts = time.strftime("%Y%m%d-%H%M%S") - args.job_name = getattr(args, "job_name", None) or f"{prev.job_name}-resume-{ts}" + args.job_name = ( + getattr(args, "job_name", None) or f"{prev.job_name}-resume-{ts}" + ) args._harbor_job_name_override = prev.job_name args._resume_gcs_output_dir = prev.gcs_output_dir print( @@ -542,17 +672,29 @@ def run(self, args: argparse.Namespace) -> int: # Make sure the local managed tree exists so the daemon (and any # downstream consumers) find LOCAL_PATHS.runs/ on first run. ensure_local_paths( - LOCAL_PATHS.home, LOCAL_PATHS.state, LOCAL_PATHS.runs, LOCAL_PATHS.logs, + LOCAL_PATHS.home, + LOCAL_PATHS.state, + LOCAL_PATHS.runs, + LOCAL_PATHS.logs, ) # Pre-submit artifact pre-cache (offline staging). Runs after the region # pin so a subclass can target the region-local mirror; may rewrite # args.model / args.dataset_path and returns extra env (e.g. offline # flags). Default no-op -> byte-identical for non-eval launchers. - precache_env = self.pre_submit_precache(args, remote_output_dir=remote_output_dir) + precache_env = self.pre_submit_precache( + args, remote_output_dir=remote_output_dir + ) command = self.build_task_command(args, remote_output_dir) env_vars = self.build_env(args) + if self.enforce_capability_token_duration: + # Keep the same secret-free policy with the remote job as well as + # the local control-plane manifest. This must never contain the + # capability URL or its JWT. + env_vars["OT_AGENT_CAPABILITY_TOKEN_DURATION_POLICY"] = json.dumps( + args._capability_token_duration_policy.to_dict(), sort_keys=True + ) if precache_env: for _k, _v in precache_env.items(): env_vars.setdefault(_k, _v) @@ -574,6 +716,11 @@ def run(self, args: argparse.Namespace) -> int: ) local_dest = LOCAL_PATHS.runs / job_name + token_policy_path = None + if self.enforce_capability_token_duration: + token_policy_path = persist_token_duration_policy( + job_name=job_name, policy=args._capability_token_duration_policy + ) print(f"[iris] Job: /{user}/{job_name}", flush=True) print(f"[iris] Cluster: {args.cluster_config}", flush=True) @@ -581,12 +728,31 @@ def run(self, args: argparse.Namespace) -> int: print(f"[iris] {accelerator.label}", flush=True) print(f"[iris] Priority: {args.priority}", flush=True) print(f"[iris] Extras: {extras or '(none)'}", flush=True) - print(f"[iris] Output: {remote_output_dir} (mode={output_mode})", flush=True) + print( + f"[iris] Output: {remote_output_dir} (mode={output_mode})", flush=True + ) + if token_policy_path is not None: + policy = args._capability_token_duration_policy + print( + "[iris] Capability-token policy: " + f"required={policy.token_required} timeout={policy.effective_timeout_seconds}s " + f"manifest={token_policy_path}", + flush=True, + ) if output_mode == "gcs": - print(f"[iris] Fetch dest: {local_dest}/ (via hpc.iris_fetch_daemon)", flush=True) + print( + f"[iris] Fetch dest: {local_dest}/ (via hpc.iris_fetch_daemon)", + flush=True, + ) else: - print(f"[iris] Work dir: {work_output_dir} (pod-local runtime state)", flush=True) - print(f"[iris] Jobs dir: {args._harbor_jobs_dir} (harbor --jobs-dir)", flush=True) + print( + f"[iris] Work dir: {work_output_dir} (pod-local runtime state)", + flush=True, + ) + print( + f"[iris] Jobs dir: {args._harbor_jobs_dir} (harbor --jobs-dir)", + flush=True, + ) print(f"[iris] Command: {shlex.join(command)}", flush=True) if args.dry_run: @@ -597,7 +763,8 @@ def run(self, args: argparse.Namespace) -> int: print( "[iris] NOTE: multi-host TPU slice (vm_count > 1). Validated on v6e-8 " "(2026-05-22 smoke #10); larger slices need their own validation pass.", - file=sys.stderr, flush=True, + file=sys.stderr, + flush=True, ) # Defer the heavy iris imports so --dry-run / --help stay snappy. @@ -613,7 +780,9 @@ def run(self, args: argparse.Namespace) -> int: # Tunnel to the controller via the current pydantic config API # (mirrors iris.cli.connect.require_controller_url's SSH-tunnel branch). config = load_config(args.cluster_config) - cluster_name = resolve_cluster_name(config, None, Path(args.cluster_config).stem) + cluster_name = resolve_cluster_name( + config, None, Path(args.cluster_config).stem + ) credentials = client_credentials(config, cluster_name) bundle = provider_bundle(config) if config.controller.controller_kind() == "local": @@ -626,7 +795,9 @@ def run(self, args: argparse.Namespace) -> int: ) with bundle.controller.tunnel(address=controller_address) as controller_url: - resources = accelerator.build_resources(cpu=args.cpu, memory=args.memory, disk=args.disk) + resources = accelerator.build_resources( + cpu=args.cpu, memory=args.memory, disk=args.disk + ) tpu_variants = list(accelerator.tpu_variants) # --replicas defaults to 1; for a multi-host TPU iris's # adjust_tpu_replicas (in client.submit) auto-scales 1 -> vm_count @@ -637,7 +808,9 @@ def run(self, args: argparse.Namespace) -> int: # separate run_tracegen issue (run harbor on the driver rank only). # GPU is single-node (vm_count 1); resolve_multinode_defaults is a # no-op passthrough there. - replicas, coscheduling = accelerator.resolve_multinode_defaults(args.replicas) + replicas, coscheduling = accelerator.resolve_multinode_defaults( + args.replicas + ) resources_proto = resources.to_proto() # Pin the job to the region we discovered at submit time, so # preempt-retries land back in the same continent and our @@ -662,7 +835,9 @@ def run(self, args: argparse.Namespace) -> int: } priority_band = _PRIO.get(args.priority, priority_band) - client = IrisClient.remote(controller_url, workspace=self.repo_root, credentials=credentials) + client = IrisClient.remote( + controller_url, workspace=self.repo_root, credentials=credentials + ) wrapped = wrap_task_command( command, @@ -687,7 +862,9 @@ def run(self, args: argparse.Namespace) -> int: # Iris auto-retries on preemption; leave at default (1000). task_image=args.task_image, priority_band=priority_band, - timeout=None if args.timeout == 0 else _seconds_to_duration(args.timeout), + timeout=None + if args.timeout == 0 + else _seconds_to_duration(args.timeout), ) full_job_id = str(job.job_id) print(f"[iris] Submitted: {full_job_id}", flush=True) @@ -708,7 +885,11 @@ def run(self, args: argparse.Namespace) -> int: cluster_config=str(args.cluster_config), ) except Exception as e: - print(f"[iris] WARN: could not register job locally: {e}", file=sys.stderr, flush=True) + print( + f"[iris] WARN: could not register job locally: {e}", + file=sys.stderr, + flush=True, + ) else: print( f"[iris] {output_mode.upper()} output mode: GPU eval registers to " @@ -724,17 +905,23 @@ def run(self, args: argparse.Namespace) -> int: status = job.wait(stream_logs=True, timeout=float("inf")) exit_code = 0 if status.state == job_pb2.JOB_STATE_SUCCEEDED else 1 except KeyboardInterrupt: - print(f"[iris] Terminating job {full_job_id}...", file=sys.stderr, flush=True) + print( + f"[iris] Terminating job {full_job_id}...", + file=sys.stderr, + flush=True, + ) client.terminate_job(job.job_id) exit_code = 130 print(f"[iris] Job exit: {exit_code}", flush=True) return exit_code + # Imported lazily inside .run() to keep CLI startup fast, but tiny enough # to define here. def _seconds_to_duration(secs: int): # Duration moved from iris.cluster.types to rigging.timing on a # marin/iris refactor; iris.client imports from rigging.timing now. from rigging.timing import Duration + return Duration.from_seconds(secs) diff --git a/scripts/iris/launch_external_opencode_eval.py b/scripts/iris/launch_external_opencode_eval.py index 70886355e..8038386e8 100644 --- a/scripts/iris/launch_external_opencode_eval.py +++ b/scripts/iris/launch_external_opencode_eval.py @@ -11,6 +11,7 @@ from __future__ import annotations import argparse +import json import os import re import shlex @@ -21,6 +22,11 @@ from pathlib import Path from urllib.parse import urlsplit +from hpc.iris.capability_tokens import ( + persist_token_duration_policy, + resolve_token_duration_policy, +) + REPO_ROOT = Path(__file__).resolve().parents[2] DEFAULT_IRIS_BIN = "/Users/benjaminfeuer/miniconda3/envs/otagent/bin/iris" @@ -310,6 +316,9 @@ def build_submit_command( "OPENCODE_DUMMY_KEY": DUMMY_API_KEY, "EXTERNAL_AGENT_API_BASE": api_base, } + task_env["OT_AGENT_CAPABILITY_TOKEN_DURATION_POLICY"] = ( + args.capability_token_duration_policy_json + ) durable_jobs_dir = durable_harbor_jobs_dir( s3_output_root=args.s3_output_dir, iris_job_name=args.job_name ) @@ -336,6 +345,8 @@ def build_submit_command( args.priority, "--max-retries", "0", + "--timeout", + str(args.timeout), "--no-wait", "--job-name", args.job_name, @@ -355,7 +366,9 @@ def main() -> int: "--existing-endpoint", help="Use an already-running federated endpoint instead of submitting a new serve.", ) - parser.add_argument("--serve-name", help="Serving job name (default: -serve).") + parser.add_argument( + "--serve-name", help="Serving job name (default: -serve)." + ) parser.add_argument( "--endpoint-name", help="Endpoint name (default: /serve/)." ) @@ -385,7 +398,18 @@ def main() -> int: parser.add_argument( "--iris-bin", default=os.environ.get("IRIS_BIN", DEFAULT_IRIS_BIN) ) - parser.add_argument("--ttl-hours", type=float, default=24.0) + parser.add_argument( + "--ttl-hours", + type=float, + default=None, + help="Requested scoped-token TTL. Defaults to the controller maximum.", + ) + parser.add_argument( + "--timeout", + type=int, + default=0, + help="Iris eval timeout in seconds (0 derives the minted-token lifetime).", + ) parser.add_argument( "--mirror-timeout-seconds", type=float, default=DEFAULT_MIRROR_TIMEOUT_SECONDS ) @@ -427,6 +451,31 @@ def main() -> int: ) args.model = args.model or f"vllm/{args.serve_model}" + requested_ttl_seconds = ( + None if args.ttl_hours is None else int(args.ttl_hours * 3600) + ) + try: + token_policy = resolve_token_duration_policy( + agent="opencode", + timeout_seconds=args.timeout, + requested_token_ttl_seconds=requested_ttl_seconds, + ) + except ValueError as exc: + parser.error(str(exc)) + args.timeout = token_policy.effective_timeout_seconds + args.ttl_hours = token_policy.effective_token_ttl_seconds / 3600.0 + args.capability_token_duration_policy_json = json.dumps( + token_policy.to_dict(), sort_keys=True + ) + policy_path = persist_token_duration_policy( + job_name=args.job_name, policy=token_policy + ) + print( + "[external-eval] capability-token policy: " + f"timeout={args.timeout}s manifest={policy_path}", + flush=True, + ) + # Reject an invalid durable destination before minting a capability token or # contacting the controller. ``build_submit_command`` repeats this check # because it is also called directly by unit tests and other tooling. @@ -450,9 +499,14 @@ def main() -> int: marin_repo = Path(args.marin_repo).expanduser() if not marin_repo.is_dir(): parser.error(f"Marin checkout not found: {marin_repo}") - print(f"[external-eval] submitting federated serve {args.serve_name}", flush=True) + print( + f"[external-eval] submitting federated serve {args.serve_name}", flush=True + ) serve_result = subprocess.run( - build_serve_command(args), cwd=marin_repo, env=submit_environment, check=False + build_serve_command(args), + cwd=marin_repo, + env=submit_environment, + check=False, ) if serve_result.returncode: return serve_result.returncode diff --git a/tests/hpc/test_capability_token_duration.py b/tests/hpc/test_capability_token_duration.py new file mode 100644 index 000000000..4fa792790 --- /dev/null +++ b/tests/hpc/test_capability_token_duration.py @@ -0,0 +1,82 @@ +"""Regression coverage for Iris capability-token duration enforcement.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest + +from hpc.iris.capability_tokens import ( + controller_max_endpoint_token_ttl_seconds, + infer_harbor_agent, + persist_token_duration_policy, + resolve_token_duration_policy, +) + + +def test_resolver_reads_the_controller_owned_constant(): + module = SimpleNamespace(MAX_ENDPOINT_TOKEN_TTL_SECONDS=345) + assert ( + controller_max_endpoint_token_ttl_seconds(module_loader=lambda _: module) == 345 + ) + + +def test_terminus_two_keeps_unlimited_timeout_and_needs_no_controller_lookup(): + policy = resolve_token_duration_policy( + agent="terminus-2", + timeout_seconds=0, + max_ttl_resolver=lambda: pytest.fail("terminus-2 must not resolve token TTL"), + ) + assert not policy.token_required + assert policy.effective_timeout_seconds == 0 + assert policy.controller_max_ttl_seconds is None + + +def test_token_required_agent_defaults_timeout_to_controller_maximum(): + policy = resolve_token_duration_policy( + agent="opencode", timeout_seconds=0, max_ttl_resolver=lambda: 345 + ) + assert policy.token_required + assert policy.effective_timeout_seconds == 345 + assert policy.effective_token_ttl_seconds == 345 + + +def test_shorter_requested_token_bounds_the_job_too(): + policy = resolve_token_duration_policy( + agent="openhands", + timeout_seconds=0, + requested_token_ttl_seconds=200, + max_ttl_resolver=lambda: 345, + ) + assert policy.effective_timeout_seconds == 200 + with pytest.raises(ValueError, match="cannot outlive"): + resolve_token_duration_policy( + agent="openhands", + timeout_seconds=201, + requested_token_ttl_seconds=200, + max_ttl_resolver=lambda: 345, + ) + + +def test_rejects_token_ttl_above_controller_maximum(): + with pytest.raises(ValueError, match="exceeds Marin controller maximum"): + resolve_token_duration_policy( + agent="opencode", + timeout_seconds=0, + requested_token_ttl_seconds=346, + max_ttl_resolver=lambda: 345, + ) + + +def test_infers_harbor_agent_and_persists_only_secret_free_policy(tmp_path): + harbor = tmp_path / "harbor.yaml" + harbor.write_text("agents:\n - name: opencode\n") + assert infer_harbor_agent(None, harbor) == "opencode" + policy = resolve_token_duration_policy( + agent="opencode", timeout_seconds=0, max_ttl_resolver=lambda: 345 + ) + saved = persist_token_duration_policy(job_name="job", policy=policy, root=tmp_path) + payload = json.loads(saved.read_text()) + assert payload["effective_timeout_seconds"] == 345 + assert not {"token", "api_base", "endpoint", "url"} & set(payload) diff --git a/tests/iris/test_launch_external_opencode_eval.py b/tests/iris/test_launch_external_opencode_eval.py index 528c32540..23cd5e475 100644 --- a/tests/iris/test_launch_external_opencode_eval.py +++ b/tests/iris/test_launch_external_opencode_eval.py @@ -61,6 +61,8 @@ def test_submit_uses_env_for_url_and_fails_fast_when_missing(): memory="128GB", disk="128GB", priority="batch", + timeout=86400, + capability_token_duration_policy_json='{"token_required": true}', job_name="eval-v2", harbor_config="config.yaml", datagen_config="external.yaml", @@ -88,6 +90,7 @@ def test_submit_uses_env_for_url_and_fails_fast_when_missing(): assert "--n_concurrent 256" in shell assert "--job_name eval-v2" in shell assert "--experiments_dir /tmp/ot-agent-runs/eval-v2" in shell + assert command[command.index("--timeout") + 1] == "86400" assert ( "--harbor_extra_arg=--jobs-dir=s3://marin-us-east-02a/iris/eval-v2/trace_jobs" in shell @@ -253,3 +256,4 @@ def fake_run(command, **kwargs): assert ( "--jobs-dir=s3://marin-us-east-02a/iris/grug-r7/trace_jobs" in eval_command[-1] ) + assert eval_command[eval_command.index("--timeout") + 1] == "86400"