interlens.runner.batched¶

def co_step(convs, turns: int | None, *, max_batch_size: int | None = None, group_seed: int = 0):
	"""Co-step ``convs`` (a shared turn schedule) in lockstep, batching each round's same-position turns.

	Each round: every conversation's current speaker is the same schedule position, so their views are gathered
	and generated in one left-padded batch (sub-batched into ``max_batch_size`` waves). The representative
	participant drives the batch — safe because same-schedule participants wrap the *same* cached model + tokenizer.
	Message hooks run per conversation before commit.

	Stop conditions are honored on the batched path too: each conversation's combined stop (its ``run_until`` /
	ambient budget, resolved once) can cap the round's ``max_new_tokens`` (the wave uses the conservative min so no
	conversation overshoots its budget) and drops a conversation from later rounds once it fires. ``turns`` may be
	``None`` for a purely stop-driven run (e.g. a matched-compute ``TokenBudget``), bounded by a large safety cap.
	"""
	if not convs:
		return convs
	stops = {}
	for c in convs:
		s = c._resolve_stop(None)
		if s is not None:
			s.reset()
		stops[id(c)] = s
	n_parts = len(convs[0].participants)
	active = list(convs)
	_SAFETY_ROUNDS = 100_000  # bound a stop-only (turns=None) run so a never-firing condition can't loop forever
	i = 0
	while active and (turns is None or i < turns) and i < _SAFETY_ROUNDS:
		spk = i % n_parts
		batch_convs = [c for c in active if _batchable(c.participants[spk])]
		other_convs = [c for c in active if not _batchable(c.participants[spk])]
		for wave in _chunks(batch_convs, max_batch_size):
			if not wave:  # this round has no batchable speaker (e.g. an all-API/non-batch group) — nothing to fuse
				continue
			rep = wave[0].participants[spk]
			caps = [_turn_cap(c, stops[id(c)], c.participants[spk]) for c in wave]
			caps = [x for x in caps if x is not None]
			mnt = min(caps) if caps else None
			views = [c._view(c.participants[spk]) for c in wave]
			msgs = rep.generate_batch(views, turn=i, group_seed=group_seed + i, max_new_tokens=mnt)
			for c, msg in zip(wave, msgs):
				msg = c._apply_hooks(msg)
				if msg is not None:
					c.transcript.messages.append(msg)
		for c in other_convs:  # tools / API / human: correctness over throughput
			c.step(c.participants[spk], max_new_tokens=_turn_cap(c, stops[id(c)], c.participants[spk]))
		for c in list(active):  # drop conversations whose stop condition has now fired
			s = stops[id(c)]
			if s is not None and len(c.transcript) and s.should_stop(c, c.transcript[-1]):
				active.remove(c)
		i += 1
	return convs

def schedule_signature(conv, turns: int):
	"""The full co-step schedule of a conversation: its turn count + the per-position participant signatures. Convs
	that share a signature have an identical speaker/model schedule, so ``co_step`` can batch each round's
	same-position turns across them safely. Grouping specs by this (instead of by turn count alone) makes batched
	execution correct for ANY mix of specs — a heterogeneous lineup simply forms its own group."""
	return (int(turns), tuple(_participant_signature(p) for p in conv.participants))