interlens.runner¶

def available_devices() -> list[str]:
	"""List the devices to spread conversations across: every CUDA GPU, else a single mps/cpu fallback.

	Multi-GPU parallelism lives *across* conversations (they're independent); within one conversation turns are
	sequential, so more GPUs never speed up a single conversation — only throughput over many."""
	if torch.cuda.is_available():
		return [f"cuda:{i}" for i in range(torch.cuda.device_count())]
	if torch.backends.mps.is_available():
		return ["mps"]
	return ["cpu"]

def register_worker_init(fn) -> object:
	"""Register a zero-arg callable to run once at worker startup (e.g. to populate the tool/analyzer registries)."""
	_WORKER_INIT_HOOKS.append(fn)
	return fn

def resolve_analyzer(analyze):
	"""Accept either a callable (in-process) or a registered name (spawn-safe) and return the callable."""
	if analyze is None or callable(analyze):
		return analyze
	if analyze in _ANALYZERS:
		return _ANALYZERS[analyze]
	raise KeyError(f"analyzer {analyze!r} not registered (have: {sorted(_ANALYZERS)})")

def run(conversations, devices=None, out_dir=None, resume=False, batched=True, max_batch_size=None) -> RunReport:
	"""Run several conversation lineups in ONE pool — the multi-lineup entry point (e.g. a ladder of model pairs ×
	conditions in one overnight job).

	Each conversation is expanded to its jobs (one per ``data()`` row if it has data, else a single conversation),
	with job ids namespaced by that conversation's ``name`` (default ``conv{i}``) so ids stay unique and resumable.
	All jobs go into one pool, so GPUs stay packed across lineups (no idle tail between sequential rollouts) under a
	single ``out_dir``/resume namespace, returning one merged ``RunReport``. Mixing lineups is safe: batched
	co-stepping groups by schedule signature, so each distinct lineup forms its own batch group.

	``Conversation.rollout`` is the single-lineup sugar for this (``run([conv], ...)``-equivalent via ``run_jobs``).
	"""
	jobs = []
	for i, conv in enumerate(conversations):
		jobs.extend(conv._jobs_for_run(i))
	return run_jobs(jobs, devices=devices, out_dir=out_dir, resume=resume, batched=batched,
	                max_batch_size=max_batch_size)

def run_jobs(jobs, devices=None, out_dir=None, resume=False, batched=True, max_batch_size=None) -> RunReport:
	"""Run ``(job_id, Conversation)`` jobs across devices, with checkpointing, resume, and per-job failure isolation.

	Parallel by default on **two** axes: one worker **process per device** (jobs round-robined; multi-GPU spawns via
	``torch.multiprocessing`` since fork+CUDA is broken — a single device runs in-process), and within each device
	**batched co-stepping** (``batched=True``). Jobs are grouped by co-step schedule signature so same-schedule jobs
	batch into one ``model.generate`` — correct for ANY mix. ``batched=False`` gives the DETERMINISTIC path.
	"""
	devices = devices or available_devices()
	pending = [(jid, conv) for jid, conv in jobs if not (resume and _already_done(out_dir, jid))]
	skipped = [jid for jid, conv in jobs if resume and _already_done(out_dir, jid)]

	if len(devices) > 1:
		results = _run_spawn(pending, devices, out_dir, batched, max_batch_size)
	elif batched:
		results = _run_group(pending, devices[0], out_dir, max_batch_size)
	else:
		results = {}
		for i, (jid, conv) in enumerate(pending):
			results[jid] = _run_one(jid, conv, devices[i % len(devices)], out_dir)
	return RunReport(results=results, skipped=skipped)