Skip to content

Bound Iris token-required launches to controller TTL #49

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
penfever merged 2 commits into from
Jul 23, 2026
Merged

Bound Iris token-required launches to controller TTL #49

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f8be6f4
fix(iris): bound capability token launches to controller TTL
Jul 23, 2026
b5e4d8a
style: format capability token duration guard
Jul 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
127 changes: 88 additions & 39 deletions data/cloud/launch_tracegen_iris.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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)
Expand All @@ -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)

Expand All @@ -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.

Expand All @@ -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
Expand All @@ -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 /
Expand All @@ -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":
Expand All @@ -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 "
Expand All @@ -195,18 +223,23 @@ 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:
print(
"[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.
Expand All @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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]

Expand Down
1 change: 1 addition & 0 deletions eval/cloud/launch_eval_iris.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
Loading
Loading