from __future__ import annotations import heapq import math import random from dataclasses import asdict, dataclass from statistics import mean from typing import Any from .latency import AnalyticalLatencyModel from .metrics import percentile from .profiles import get_accelerator, get_model AGENT_ROLES = ("planner", "retriever", "reasoner", "verifier", "writer") END = "END" ROLE_PREFIX_TOKENS = { "planner": 512, "retriever": 768, "reasoner": 896, "verifier": 640, "writer": 576, } # Structured workflow transitions used to make online execution learning # meaningful. Rows sum to one. The shifted matrix changes the mix of external # retrieval/verification work without changing the model or hardware profile. BASE_TRANSITIONS: dict[str, dict[str, float]] = { "planner": {"retriever": 0.52, "reasoner": 0.28, "writer": 0.12, END: 0.08}, "retriever": {"reasoner": 0.50, "verifier": 0.20, "retriever": 0.10, "writer": 0.12, END: 0.08}, "reasoner": {"verifier": 0.32, "retriever": 0.18, "writer": 0.30, "reasoner": 0.10, END: 0.10}, "verifier": {"reasoner": 0.24, "retriever": 0.12, "writer": 0.34, END: 0.30}, "writer": {"verifier": 0.22, "reasoner": 0.10, "retriever": 0.05, END: 0.63}, } SHIFTED_TRANSITIONS: dict[str, dict[str, float]] = { "planner": {"retriever": 0.72, "reasoner": 0.14, "writer": 0.06, END: 0.08}, "retriever": {"reasoner": 0.30, "verifier": 0.30, "retriever": 0.24, "writer": 0.08, END: 0.08}, "reasoner": {"verifier": 0.24, "retriever": 0.36, "writer": 0.18, "reasoner": 0.12, END: 0.10}, "verifier": {"reasoner": 0.18, "retriever": 0.28, "writer": 0.24, END: 0.30}, "writer": {"verifier": 0.34, "reasoner": 0.08, "retriever": 0.10, END: 0.48}, } ROLE_GAP_MULTIPLIERS = { "planner": 0.35, "retriever": 1.45, "reasoner": 0.55, "verifier": 0.75, "writer": 0.30, } PREFETCH_POLICIES = {"none", "cumulative", "decayed", "multistep", "utility", "oracle", "oracle_horizon"} @dataclass class ExecutionLearningConfig: model: str = "Qwen2.5-3B" accelerator: str = "L4" quantization: str = "int8" seed: int = 7 duration_s: float = 120.0 workflow_rate_rps: float = 0.18 dynamic_prompt_tokens_mean: int = 320 dynamic_prompt_cv: float = 0.40 output_tokens_mean: int = 64 output_tokens_cv: float = 0.45 tool_gap_mean_s: float = 1.2 tool_gap_cv: float = 0.65 max_steps: int = 8 shift_fraction: float = 0.55 prefix_cache_budget_fraction: float = 0.62 host_bandwidth_gbps: float = 32.0 transfer_base_ms: float = 0.15 prefetch_policy: str = "decayed" transition_decay: float = 0.85 confidence_threshold: float = 0.55 prior_strength: float = 0.35 forecast_horizon: int = 3 prefetch_top_k: int = 2 forecast_discount: float = 0.75 forecast_min_score: float = 0.10 utility_threshold_ms: float = 0.0 timeline_points: int = 240 @classmethod def from_dict(cls, data: dict[str, Any]) -> "ExecutionLearningConfig": allowed = cls.__dataclass_fields__.keys() cfg = cls(**{key: data[key] for key in allowed if key in data}) cfg.duration_s = max(float(cfg.duration_s), 20.0) cfg.workflow_rate_rps = max(float(cfg.workflow_rate_rps), 0.01) cfg.max_steps = max(2, min(int(cfg.max_steps), 16)) cfg.shift_fraction = min(max(float(cfg.shift_fraction), 0.05), 0.95) cfg.prefix_cache_budget_fraction = min(max(float(cfg.prefix_cache_budget_fraction), 0.10), 1.50) cfg.host_bandwidth_gbps = max(float(cfg.host_bandwidth_gbps), 0.1) cfg.transfer_base_ms = max(float(cfg.transfer_base_ms), 0.0) cfg.transition_decay = min(max(float(cfg.transition_decay), 0.20), 1.0) cfg.confidence_threshold = min(max(float(cfg.confidence_threshold), 0.0), 1.0) cfg.prior_strength = max(float(cfg.prior_strength), 0.01) cfg.forecast_horizon = max(1, min(int(cfg.forecast_horizon), 6)) cfg.prefetch_top_k = max(1, min(int(cfg.prefetch_top_k), len(AGENT_ROLES))) cfg.forecast_discount = min(max(float(cfg.forecast_discount), 0.05), 1.0) cfg.forecast_min_score = min(max(float(cfg.forecast_min_score), 0.0), 1.0) cfg.utility_threshold_ms = float(cfg.utility_threshold_ms) cfg.timeline_points = max(40, min(int(cfg.timeline_points), 1000)) if cfg.prefetch_policy not in PREFETCH_POLICIES: raise ValueError(f"Unsupported execution prefetch policy: {cfg.prefetch_policy}") return cfg def to_dict(self) -> dict[str, Any]: return asdict(self) @dataclass class WorkflowStep: role: str dynamic_prompt_tokens: int output_tokens: int gap_after_s: float shifted_regime: bool @dataclass class WorkflowSpec: workflow_id: int arrival_time: float steps: list[WorkflowStep] @dataclass class PrefixEntry: role: str size_gb: float last_access: float available_time: float source: str generation: int used: bool = False class OnlineTransitionPredictor: """First-order online execution model with optional exponential forgetting. The predictor never observes a transition before the next role is revealed. ``decay=1`` is cumulative counting. Smaller values forget older transitions and can adapt after a workflow-regime shift. """ def __init__(self, *, decay: float = 0.85, prior_strength: float = 0.35) -> None: self.decay = min(max(float(decay), 0.20), 1.0) self.prior_strength = max(float(prior_strength), 0.01) self.targets = (*AGENT_ROLES, END) self.rows: dict[str, dict[str, float]] = { role: {target: self.prior_strength for target in self.targets} for role in AGENT_ROLES } self.observations = 0 def predict(self, current_role: str) -> tuple[str, float, dict[str, float]]: row = self.rows[current_role] total = sum(row.values()) or 1.0 probabilities = {target: value / total for target, value in row.items()} best = max(self.targets, key=lambda target: (probabilities[target], -self.targets.index(target))) return best, probabilities[best], probabilities def forecast( self, current_role: str, *, horizon: int = 3, discount: float = 0.75, ) -> dict[str, Any]: """Roll the learned first-order matrix forward without future trace access. ``role_scores`` are discounted expected future visits, not calibrated probabilities of at-least-one reuse. They are therefore used as a ranking signal and are normalized separately for UI/decision thresholds. """ horizon = max(1, min(int(horizon), 8)) discount = min(max(float(discount), 0.0), 1.0) distributions: list[dict[str, float]] = [] current: dict[str, float] = {target: 0.0 for target in self.targets} current[current_role] = 1.0 role_scores = {role: 0.0 for role in AGENT_ROLES} for depth in range(1, horizon + 1): nxt = {target: 0.0 for target in self.targets} for source, mass in current.items(): if mass <= 0.0: continue if source == END: nxt[END] += mass continue _, _, probs = self.predict(source) for target, probability in probs.items(): nxt[target] += mass * probability distributions.append(nxt) weight = discount ** (depth - 1) for role in AGENT_ROLES: role_scores[role] += weight * nxt.get(role, 0.0) current = nxt total_score = sum(role_scores.values()) normalized = { role: (score / total_score if total_score > 1e-12 else 0.0) for role, score in role_scores.items() } ranked = sorted(AGENT_ROLES, key=lambda role: (-role_scores[role], role)) return { "horizon": horizon, "discount": discount, "distributions": distributions, "role_scores": role_scores, "normalized_scores": normalized, "ranked_roles": ranked, } def observe(self, current_role: str, next_role: str) -> None: row = self.rows[current_role] if self.decay < 1.0: for target in self.targets: row[target] = max(self.prior_strength * 0.05, row[target] * self.decay) row[next_role] = row.get(next_role, 0.0) + 1.0 self.observations += 1 def snapshot(self) -> dict[str, Any]: matrix = {} for role in AGENT_ROLES: _, _, probs = self.predict(role) matrix[role] = probs return { "decay": self.decay, "prior_strength": self.prior_strength, "observations": self.observations, "probabilities": matrix, } def _sample_positive_lognormal(rng: random.Random, mean_value: float, cv: float, minimum: float) -> float: mean_value = max(float(mean_value), minimum) cv = max(float(cv), 0.0) if cv <= 1e-12: return mean_value sigma2 = math.log1p(cv * cv) sigma = math.sqrt(sigma2) mu = math.log(max(mean_value, 1e-9)) - sigma2 / 2.0 return max(minimum, rng.lognormvariate(mu, sigma)) def _sample_int(rng: random.Random, mean_value: float, cv: float, minimum: int = 1) -> int: return max(minimum, int(round(_sample_positive_lognormal(rng, mean_value, cv, float(minimum))))) def _sample_transition(rng: random.Random, role: str, shifted: bool) -> str: matrix = SHIFTED_TRANSITIONS if shifted else BASE_TRANSITIONS row = matrix[role] targets = tuple(row) return rng.choices(targets, weights=[row[target] for target in targets], k=1)[0] def generate_workflows(cfg: ExecutionLearningConfig) -> list[WorkflowSpec]: rng = random.Random(cfg.seed ^ 0xE91A7) workflows: list[WorkflowSpec] = [] now = 0.0 workflow_id = 0 while True: now += rng.expovariate(cfg.workflow_rate_rps) if now > cfg.duration_s: break shifted = now >= cfg.duration_s * cfg.shift_fraction role = "planner" steps: list[WorkflowStep] = [] for step_idx in range(cfg.max_steps): dynamic_prompt = _sample_int(rng, cfg.dynamic_prompt_tokens_mean, cfg.dynamic_prompt_cv, 32) output = _sample_int(rng, cfg.output_tokens_mean, cfg.output_tokens_cv, 1) next_role = _sample_transition(rng, role, shifted) if next_role == END or step_idx == cfg.max_steps - 1: gap = 0.0 else: mean_gap = cfg.tool_gap_mean_s * ROLE_GAP_MULTIPLIERS.get(next_role, 1.0) if shifted and next_role in {"retriever", "verifier"}: mean_gap *= 1.35 gap = _sample_positive_lognormal(rng, mean_gap, cfg.tool_gap_cv, 0.0) steps.append(WorkflowStep(role, dynamic_prompt, output, gap, shifted)) if next_role == END or step_idx == cfg.max_steps - 1: break role = next_role workflows.append(WorkflowSpec(workflow_id, now, steps)) workflow_id += 1 return workflows def _prediction_calibration(rows: list[dict[str, Any]], bins: int = 10) -> dict[str, Any]: if not rows: return {"count": 0, "brier": 0.0, "log_loss": 0.0, "ece": 0.0, "bins": []} targets = (*AGENT_ROLES, END) brier_total = 0.0 log_total = 0.0 buckets: list[list[tuple[float, float]]] = [[] for _ in range(max(2, bins))] for row in rows: probs = row.get("probabilities", {}) actual = str(row.get("actual_next_role", END)) brier_total += sum((float(probs.get(target, 0.0)) - (1.0 if target == actual else 0.0)) ** 2 for target in targets) actual_p = max(float(probs.get(actual, 0.0)), 1e-12) log_total += -math.log(actual_p) confidence = min(max(float(row.get("confidence", 0.0)), 0.0), 1.0) correct = 1.0 if row.get("prediction_correct") else 0.0 index = min(int(confidence * len(buckets)), len(buckets) - 1) buckets[index].append((confidence, correct)) ece = 0.0 rendered = [] for index, bucket in enumerate(buckets): if not bucket: continue avg_conf = mean(value[0] for value in bucket) accuracy = mean(value[1] for value in bucket) weight = len(bucket) / len(rows) ece += weight * abs(avg_conf - accuracy) rendered.append( { "lower": index / len(buckets), "upper": (index + 1) / len(buckets), "count": len(bucket), "mean_confidence": avg_conf, "accuracy": accuracy, } ) return { "count": len(rows), "brier": brier_total / len(rows), "log_loss": log_total / len(rows), "ece": ece, "bins": rendered, } class ExecutionLearningSimulator: """Reference simulator for online workflow learning and prefix prefetch. This lab intentionally isolates cross-workflow *agent-prefix* caching from the session-KV simulator. It uses one serial compute station plus one serialized prefix-transfer link and a bounded shared HBM prefix cache. Static agent prefixes are always available in the modeled host tier; policy decisions determine whether likely next prefixes are prefetched into HBM during gaps. """ def __init__(self, cfg: ExecutionLearningConfig, workflows: list[WorkflowSpec] | None = None): self.cfg = cfg self.model = get_model(cfg.model) self.accelerator = get_accelerator(cfg.accelerator) self.latency = AnalyticalLatencyModel(self.model, self.accelerator, cfg.quantization) self.workflows = workflows if workflows is not None else generate_workflows(cfg) self.predictor = OnlineTransitionPredictor(decay=cfg.transition_decay, prior_strength=cfg.prior_strength) self.events: list[tuple[float, int, str, tuple[Any, ...]]] = [] self.event_seq = 0 self.now = 0.0 self.server_busy_until = 0.0 self.transfer_busy_until = 0.0 self.queue: list[tuple[int, int, float]] = [] self.cache: dict[str, PrefixEntry] = {} self.prefix_source: dict[str, str] = {} self.pending_predictions: dict[tuple[int, int], dict[str, Any]] = {} self.step_rows: list[dict[str, Any]] = [] self.prediction_rows: list[dict[str, Any]] = [] self.timeline: list[dict[str, Any]] = [] self.last_memory_time = 0.0 self.hbm_gb_seconds = 0.0 self.peak_hbm_gb = 0.0 self.pressure_evictions = 0 self.prefetch_attempts = 0 self.prefetch_correct = 0 self.prefetch_useful = 0 self.prefetch_bytes = 0.0 self.prefetch_useful_bytes = 0.0 self.unused_prefetch_bytes = 0.0 self.wrong_step_prefetch_bytes = 0.0 self.transfer_latencies_ms: list[float] = [] self.prefix_hits = 0 self.prefetch_hits = 0 self.prefill_tokens_saved = 0 self.workflow_completion: dict[int, float] = {} kv_bytes = self.latency.kv_bytes_per_token() self.role_size_gb = {role: ROLE_PREFIX_TOKENS[role] * kv_bytes / 1e9 for role in AGENT_ROLES} reference_working_set = sum(self.role_size_gb.values()) self.cache_capacity_gb = max(reference_working_set * cfg.prefix_cache_budget_fraction, min(self.role_size_gb.values())) for workflow in self.workflows: if workflow.steps: self._push(workflow.arrival_time, "step_ready", workflow.workflow_id, 0) def _push(self, when: float, kind: str, *payload: Any) -> None: self.event_seq += 1 heapq.heappush(self.events, (float(when), self.event_seq, kind, tuple(payload))) def _used_gb(self) -> float: return sum(entry.size_gb for entry in self.cache.values()) def _integrate_memory(self, new_time: float) -> None: new_time = max(float(new_time), self.last_memory_time) elapsed = new_time - self.last_memory_time self.hbm_gb_seconds += self._used_gb() * elapsed self.last_memory_time = new_time self.peak_hbm_gb = max(self.peak_hbm_gb, self._used_gb()) def _record_timeline(self) -> None: min_gap = max(self.cfg.duration_s / self.cfg.timeline_points, 0.08) if self.timeline and self.now - self.timeline[-1]["time_s"] < min_gap: return self.timeline.append( { "time_s": self.now, "queued_steps": len(self.queue), "server_busy": 1 if self.server_busy_until > self.now + 1e-12 else 0, "prefix_hbm_gb": self._used_gb(), "cached_roles": len(self.cache), "transfer_busy": 1 if self.transfer_busy_until > self.now + 1e-12 else 0, } ) if len(self.timeline) > self.cfg.timeline_points * 2: self.timeline = self.timeline[::2] def _mark_unused_prefetch(self, entry: PrefixEntry) -> None: if entry.source == "prefetch" and not entry.used: self.unused_prefetch_bytes += entry.size_gb * 1e9 def _candidate_eviction_cost_ms(self, role: str, forecast_scores: dict[str, float]) -> float: """Approximate downstream cost of evicting forecast-relevant cached prefixes.""" size_gb = self.role_size_gb[role] existing = self.cache.get(role) current_without = self._used_gb() - (existing.size_gb if existing else 0.0) overflow = max(0.0, current_without + size_gb - self.cache_capacity_gb) if overflow <= 1e-12: return 0.0 cost_ms = 0.0 reclaimed = 0.0 candidates = sorted( [entry for key, entry in self.cache.items() if key != role], key=lambda entry: (entry.last_access, entry.role), ) for victim in candidates: if reclaimed >= overflow - 1e-12: break reuse_score = max(float(forecast_scores.get(victim.role, 0.0)), 0.0) recompute_ms = self.latency.prefill_seconds([ROLE_PREFIX_TOKENS[victim.role]]) * 1000.0 cost_ms += reuse_score * recompute_ms reclaimed += victim.size_gb return cost_ms def _prefetch_utility_ms(self, role: str, forecast_scores: dict[str, float]) -> float: score = max(float(forecast_scores.get(role, 0.0)), 0.0) saved_ms = self.latency.prefill_seconds([ROLE_PREFIX_TOKENS[role]]) * 1000.0 transfer_ms = self.cfg.transfer_base_ms + self.role_size_gb[role] / self.cfg.host_bandwidth_gbps * 1000.0 eviction_ms = self._candidate_eviction_cost_ms(role, forecast_scores) return score * saved_ms - transfer_ms - eviction_ms def _evict_for(self, role: str, size_gb: float) -> bool: existing = self.cache.get(role) current_without_role = self._used_gb() - (existing.size_gb if existing else 0.0) if size_gb > self.cache_capacity_gb + 1e-12: return False while current_without_role + size_gb > self.cache_capacity_gb + 1e-12: candidates = [entry for key, entry in self.cache.items() if key != role] if not candidates: return False victim = min(candidates, key=lambda entry: (entry.last_access, entry.role)) self._integrate_memory(self.now) self._mark_unused_prefetch(victim) self.cache.pop(victim.role, None) self.prefix_source.pop(victim.role, None) self.pressure_evictions += 1 current_without_role = self._used_gb() - (existing.size_gb if existing and role in self.cache else 0.0) return True def _install_prefix(self, role: str, *, available_time: float, source: str) -> int | None: size_gb = self.role_size_gb[role] if not self._evict_for(role, size_gb): return None self._integrate_memory(self.now) previous = self.cache.get(role) if previous is not None: self._mark_unused_prefetch(previous) generation = (previous.generation + 1) if previous else 1 self.cache[role] = PrefixEntry(role, size_gb, self.now, available_time, source, generation) self.prefix_source[role] = source self.peak_hbm_gb = max(self.peak_hbm_gb, self._used_gb()) self._record_timeline() return generation def _prefetch(self, role: str) -> dict[str, Any]: entry = self.cache.get(role) if entry is not None: return {"attempted": False, "reason": "already_cached", "role": role, "available_time": entry.available_time} size_gb = self.role_size_gb[role] transfer_s = self.cfg.transfer_base_ms / 1000.0 + size_gb / self.cfg.host_bandwidth_gbps start = max(self.now, self.transfer_busy_until) available = start + transfer_s generation = self._install_prefix(role, available_time=available, source="prefetch") if generation is None: return {"attempted": False, "reason": "capacity_reject", "role": role, "available_time": available} self.transfer_busy_until = available self.prefetch_attempts += 1 self.prefetch_bytes += size_gb * 1e9 self.transfer_latencies_ms.append((available - self.now) * 1000.0) self._push(available, "prefetch_complete", role, generation) return {"attempted": True, "reason": "scheduled", "role": role, "available_time": available, "size_gb": size_gb} def _actual_next_role(self, workflow_id: int, step_index: int) -> str: workflow = self.workflows[workflow_id] if step_index + 1 >= len(workflow.steps): return END return workflow.steps[step_index + 1].role def _actual_future_roles(self, workflow_id: int, step_index: int, horizon: int) -> list[str]: workflow = self.workflows[workflow_id] roles = [step.role for step in workflow.steps[step_index + 1 : step_index + 1 + max(1, horizon)]] return roles def _planned_prefetch_roles( self, workflow_id: int, step_index: int, current_role: str, ) -> tuple[str, float, dict[str, float], dict[str, Any], list[dict[str, Any]]]: actual_next = self._actual_next_role(workflow_id, step_index) predicted_role, confidence, probabilities = self.predictor.predict(current_role) forecast = self.predictor.forecast( current_role, horizon=self.cfg.forecast_horizon, discount=self.cfg.forecast_discount, ) plans: list[dict[str, Any]] = [] if self.cfg.prefetch_policy == "oracle": predicted_role = actual_next confidence = 1.0 probabilities = {target: 0.0 for target in (*AGENT_ROLES, END)} probabilities[predicted_role] = 1.0 if actual_next != END: plans = [{"role": actual_next, "score": 1.0, "utility_ms": None, "oracle": True}] return predicted_role, confidence, probabilities, forecast, plans if self.cfg.prefetch_policy == "oracle_horizon": actual_future = self._actual_future_roles(workflow_id, step_index, self.cfg.forecast_horizon) predicted_role = actual_next confidence = 1.0 probabilities = {target: 0.0 for target in (*AGENT_ROLES, END)} probabilities[predicted_role] = 1.0 seen: set[str] = set() for role in actual_future: if role not in seen: plans.append({"role": role, "score": 1.0, "utility_ms": None, "oracle": True}) seen.add(role) if len(plans) >= self.cfg.prefetch_top_k: break return predicted_role, confidence, probabilities, forecast, plans if self.cfg.prefetch_policy in {"none", "cumulative", "decayed"}: if self.cfg.prefetch_policy != "none" and predicted_role != END and confidence >= self.cfg.confidence_threshold: plans = [{"role": predicted_role, "score": confidence, "utility_ms": None, "oracle": False}] return predicted_role, confidence, probabilities, forecast, plans ranked = forecast["ranked_roles"] normalized = forecast["normalized_scores"] raw_scores = forecast["role_scores"] for role in ranked: if len(plans) >= self.cfg.prefetch_top_k: break normalized_score = float(normalized.get(role, 0.0)) if normalized_score < self.cfg.forecast_min_score: continue utility_ms = self._prefetch_utility_ms(role, raw_scores) if self.cfg.prefetch_policy == "utility" and utility_ms < self.cfg.utility_threshold_ms: continue plans.append( { "role": role, "score": normalized_score, "raw_score": float(raw_scores.get(role, 0.0)), "utility_ms": utility_ms, "oracle": False, } ) return predicted_role, confidence, probabilities, forecast, plans def _schedule_prediction(self, workflow_id: int, step_index: int, current_role: str) -> None: actual_next = self._actual_next_role(workflow_id, step_index) predicted_role, confidence, probabilities, forecast, plans = self._planned_prefetch_roles( workflow_id, step_index, current_role ) attempted_roles: list[str] = [] attempt_sizes: dict[str, float] = {} attempt_reasons: dict[str, str] = {} for plan in plans: role = str(plan["role"]) result = self._prefetch(role) attempt_reasons[role] = str(result.get("reason", "unknown")) if result.get("attempted"): attempted_roles.append(role) attempt_sizes[role] = float(result.get("size_gb", 0.0)) if self.cfg.prefetch_policy == "none": decision_reason = "learn_only" elif not plans: decision_reason = "no_candidate" elif attempted_roles: decision_reason = "scheduled" else: decision_reason = ",".join(sorted(set(attempt_reasons.values()))) or "not_scheduled" self.pending_predictions[(workflow_id, step_index)] = { "workflow_id": workflow_id, "step_index": step_index + 1, "current_role": current_role, "predicted_role": predicted_role, "confidence": confidence, "probabilities": probabilities, "forecast_horizon": self.cfg.forecast_horizon, "forecast_discount": self.cfg.forecast_discount, "forecast_role_scores": forecast["role_scores"], "forecast_normalized_scores": forecast["normalized_scores"], "forecast_roles": [str(plan["role"]) for plan in plans], "forecast_plans": plans, "prefetch_attempted": bool(attempted_roles), "prefetch_roles": attempted_roles, "prefetch_reasons": attempt_reasons, "prefetch_reason": decision_reason, "prefetch_sizes_gb": attempt_sizes, "prefetch_size_gb": sum(attempt_sizes.values()), "shifted_regime": self.workflows[workflow_id].steps[step_index].shifted_regime, } if actual_next == END: self._observe_transition(workflow_id, step_index, actual_next) def _observe_transition(self, workflow_id: int, step_index: int, actual_next: str) -> None: key = (workflow_id, step_index) pending = self.pending_predictions.pop(key, None) workflow = self.workflows[workflow_id] current_role = workflow.steps[step_index].role if pending is None: return predicted_role = str(pending["predicted_role"]) attempted_roles = [str(role) for role in pending.get("prefetch_roles", [])] correct = predicted_role == actual_next if actual_next != END and actual_next in attempted_roles: self.prefetch_correct += 1 immediate_wrong_bytes = 0.0 for role in attempted_roles: if role != actual_next: immediate_wrong_bytes += float(pending.get("prefetch_sizes_gb", {}).get(role, 0.0)) * 1e9 self.wrong_step_prefetch_bytes += immediate_wrong_bytes future_roles = self._actual_future_roles(workflow_id, step_index, int(pending.get("forecast_horizon", 1))) actual_future_set = set(future_roles) forecast_roles = [str(role) for role in pending.get("forecast_roles", [])] forecast_set = set(forecast_roles) forecast_recall = ( len(actual_future_set & forecast_set) / len(actual_future_set) if actual_future_set else None ) self.prediction_rows.append( pending | { "actual_next_role": actual_next, "actual_future_roles": future_roles, "prediction_correct": correct, "forecast_next_hit": actual_next in forecast_set if actual_next != END else predicted_role == END, "forecast_recall": forecast_recall, "immediate_wrong_prefetch_gb": immediate_wrong_bytes / 1e9, "post_shift": bool(pending["shifted_regime"]), } ) if self.cfg.prefetch_policy not in {"oracle", "oracle_horizon"}: self.predictor.observe(current_role, actual_next) def _start_next(self) -> None: if self.server_busy_until > self.now + 1e-12 or not self.queue: return workflow_id, step_index, ready_time = self.queue.pop(0) workflow = self.workflows[workflow_id] step = workflow.steps[step_index] role = step.role entry = self.cache.get(role) prefix_tokens = ROLE_PREFIX_TOKENS[role] prefix_wait_s = 0.0 cache_hit = False cache_source = "miss" if entry is not None and entry.available_time <= self.now + 1e-12: cache_hit = True cache_source = self.prefix_source.get(role, entry.source) elif entry is not None: # If a prefetch is almost complete, waiting can be cheaper than # recomputing the whole static prefix. Otherwise the step proceeds # immediately with a reactive prefill and the pending prefetch is # treated as not useful for this invocation. remaining = max(0.0, entry.available_time - self.now) full_prefix_s = self.latency.prefill_seconds([prefix_tokens]) if remaining < full_prefix_s: prefix_wait_s = remaining cache_hit = True cache_source = "prefetch_wait" if cache_hit: self.prefix_hits += 1 if cache_source.startswith("prefetch"): self.prefetch_hits += 1 self.prefill_tokens_saved += prefix_tokens if entry is not None: if entry.source == "prefetch" and not entry.used: entry.used = True self.prefetch_useful += 1 self.prefetch_useful_bytes += entry.size_gb * 1e9 entry.last_access = self.now + prefix_wait_s prefill_tokens = max(step.dynamic_prompt_tokens, 1) else: prefill_tokens = max(prefix_tokens + step.dynamic_prompt_tokens, 1) prefill_s = self.latency.prefill_seconds([prefill_tokens]) context = prefill_tokens first_decode_s = self.latency.decode_step_seconds([context]) midpoint = context + max(step.output_tokens // 2, 1) decode_step_s = self.latency.decode_step_seconds([midpoint]) service_s = prefix_wait_s + prefill_s + step.output_tokens * decode_step_s ttft_s = (self.now - ready_time) + prefix_wait_s + prefill_s + first_decode_s self.server_busy_until = self.now + service_s self._push( self.server_busy_until, "step_complete", workflow_id, step_index, ready_time, ttft_s, service_s, cache_hit, cache_source, prefill_tokens, ) def _finish_step( self, workflow_id: int, step_index: int, ready_time: float, ttft_s: float, service_s: float, cache_hit: bool, cache_source: str, prefill_tokens: int, ) -> None: workflow = self.workflows[workflow_id] step = workflow.steps[step_index] role = step.role # Reactive execution materializes the role prefix after a miss. if not cache_hit: self._install_prefix(role, available_time=self.now, source="reactive") elif role in self.cache: self.cache[role].last_access = self.now self.step_rows.append( { "workflow_id": workflow_id, "step_index": step_index + 1, "role": role, "shifted_regime": step.shifted_regime, "ready_time": ready_time, "completion_time": self.now, "cache_hit": cache_hit, "cache_source": cache_source, "prefix_tokens": ROLE_PREFIX_TOKENS[role], "dynamic_prompt_tokens": step.dynamic_prompt_tokens, "prefill_tokens": prefill_tokens, "output_tokens": step.output_tokens, "ttft_ms": ttft_s * 1000.0, "e2e_ms": (self.now - ready_time) * 1000.0, "queue_ms": max(0.0, (self.now - ready_time - service_s) * 1000.0), } ) self._schedule_prediction(workflow_id, step_index, role) actual_next = self._actual_next_role(workflow_id, step_index) if actual_next == END: self.workflow_completion[workflow_id] = self.now else: gap = step.gap_after_s ready_at = self.now + gap self._push(ready_at, "transition_observed", workflow_id, step_index, actual_next) self._push(ready_at, "step_ready", workflow_id, step_index + 1) self.server_busy_until = self.now self._start_next() def run(self) -> dict[str, Any]: self._record_timeline() while self.events: event_time, _, kind, payload = heapq.heappop(self.events) self._integrate_memory(event_time) self.now = event_time if kind == "step_ready": workflow_id, step_index = int(payload[0]), int(payload[1]) self.queue.append((workflow_id, step_index, self.now)) self._start_next() elif kind == "step_complete": self._finish_step( int(payload[0]), int(payload[1]), float(payload[2]), float(payload[3]), float(payload[4]), bool(payload[5]), str(payload[6]), int(payload[7]), ) elif kind == "transition_observed": self._observe_transition(int(payload[0]), int(payload[1]), str(payload[2])) elif kind == "prefetch_complete": role, generation = str(payload[0]), int(payload[1]) entry = self.cache.get(role) if entry is not None and entry.generation == generation and entry.available_time <= self.now + 1e-12: entry.last_access = self.now self._record_timeline() self._integrate_memory(self.now) self._record_timeline() step_ttft = [float(row["ttft_ms"]) for row in self.step_rows] step_e2e = [float(row["e2e_ms"]) for row in self.step_rows] workflow_e2e = [ (self.workflow_completion[wf.workflow_id] - wf.arrival_time) * 1000.0 for wf in self.workflows if wf.workflow_id in self.workflow_completion ] horizon = max([*self.workflow_completion.values(), self.cfg.duration_s, 1e-9]) predictions = self.prediction_rows pre = [row for row in predictions if not row["post_shift"]] post = [row for row in predictions if row["post_shift"]] def accuracy(rows: list[dict[str, Any]]) -> float: return sum(1 for row in rows if row["prediction_correct"]) / len(rows) if rows else 0.0 # Account for prefetched entries that remained unused through the end of the run. for entry in list(self.cache.values()): self._mark_unused_prefetch(entry) prefetch_precision = self.prefetch_correct / self.prefetch_attempts if self.prefetch_attempts else 0.0 prefetch_coverage = ( sum(1 for row in predictions if row.get("prefetch_attempted")) / len(predictions) if predictions else 0.0 ) prefetch_utilization = self.prefetch_useful / self.prefetch_attempts if self.prefetch_attempts else 0.0 forecast_recall_values = [ float(row["forecast_recall"]) for row in predictions if row.get("forecast_recall") is not None ] forecast_recall = mean(forecast_recall_values) if forecast_recall_values else 0.0 workflow_count = len(self.workflows) completed_count = len(self.workflow_completion) summary = { "workflows_generated": workflow_count, "workflows_completed": completed_count, "steps_completed": len(self.step_rows), "workflow_completion_rate": completed_count / workflow_count if workflow_count else 0.0, "workflow_throughput_rps": completed_count / horizon, "simulated_makespan_s": horizon, } resource = { "prefix_cache_capacity_gb": self.cache_capacity_gb, "mean_prefix_hbm_gb": self.hbm_gb_seconds / horizon, "peak_prefix_hbm_gb": self.peak_hbm_gb, "hbm_gb_seconds": self.hbm_gb_seconds, "prefix_hit_rate": self.prefix_hits / len(self.step_rows) if self.step_rows else 0.0, "prefetch_hit_rate": self.prefetch_hits / len(self.step_rows) if self.step_rows else 0.0, "prefill_tokens_saved": self.prefill_tokens_saved, "pressure_evictions": self.pressure_evictions, "prefetch_attempts": self.prefetch_attempts, "prefetch_correct": self.prefetch_correct, "prefetch_precision": prefetch_precision, "prefetch_coverage": prefetch_coverage, "prefetch_utilization": prefetch_utilization, "prefetch_attempts_per_prediction": self.prefetch_attempts / len(predictions) if predictions else 0.0, "forecast_recall": forecast_recall, "prefetch_gb": self.prefetch_bytes / 1e9, "prefetch_useful_gb": self.prefetch_useful_bytes / 1e9, "unused_prefetch_gb": self.unused_prefetch_bytes / 1e9, "wrong_step_prefetch_gb": self.wrong_step_prefetch_bytes / 1e9, "p95_prefetch_transfer_ms": percentile(self.transfer_latencies_ms, 0.95), } prediction = { "count": len(predictions), "top1_accuracy": accuracy(predictions), "pre_shift_accuracy": accuracy(pre), "post_shift_accuracy": accuracy(post), "calibration": _prediction_calibration(predictions), "pre_shift_calibration": _prediction_calibration(pre), "post_shift_calibration": _prediction_calibration(post), "rows": predictions[:4000], "predictor": self.predictor.snapshot(), } latency = { "step_ttft_ms": {"p50": percentile(step_ttft, 0.50), "p95": percentile(step_ttft, 0.95), "p99": percentile(step_ttft, 0.99)}, "step_e2e_ms": {"p50": percentile(step_e2e, 0.50), "p95": percentile(step_e2e, 0.95), "p99": percentile(step_e2e, 0.99)}, "workflow_e2e_ms": {"p50": percentile(workflow_e2e, 0.50), "p95": percentile(workflow_e2e, 0.95), "p99": percentile(workflow_e2e, 0.99)}, } return { "config": self.cfg.to_dict(), "provenance": { "simulator": "InferScale-Sim", "mode": "online-agent-execution-learning", "latency_profile_type": "analytical-reference", "execution_model": "online-transition-multistep-prefix-prefetch-reference", "prefetch_policy": self.cfg.prefetch_policy, "lookahead": "oracle-upper-bound" if self.cfg.prefetch_policy in {"oracle", "oracle_horizon"} else "no-future-transition-lookahead", "warning": "Execution Learning isolates cross-workflow static-prefix prediction/prefetch from session-KV retention and dynamic batching.", }, "summary": summary, "latency": latency, "resource": resource, "prediction": prediction, "steps": self.step_rows[:4000], "timeline": self.timeline, "role_prefix_tokens": ROLE_PREFIX_TOKENS, } def run_execution_learning(config: dict[str, Any], workflows: list[WorkflowSpec] | None = None) -> dict[str, Any]: cfg = ExecutionLearningConfig.from_dict(config) return ExecutionLearningSimulator(cfg, workflows=workflows).run() def _common_workflows(cfg: ExecutionLearningConfig) -> list[WorkflowSpec]: return generate_workflows(cfg) def _row(label: str, result: dict[str, Any]) -> dict[str, Any]: return { "label": label, "prefetch_policy": result["config"]["prefetch_policy"], "p95_step_ttft_ms": result["latency"]["step_ttft_ms"]["p95"], "p95_workflow_e2e_ms": result["latency"]["workflow_e2e_ms"]["p95"], "workflow_throughput_rps": result["summary"]["workflow_throughput_rps"], "prefix_hit_rate": result["resource"]["prefix_hit_rate"], "prefetch_precision": result["resource"]["prefetch_precision"], "prefetch_coverage": result["resource"]["prefetch_coverage"], "prefetch_utilization": result["resource"].get("prefetch_utilization", 0.0), "forecast_recall": result["resource"].get("forecast_recall", 0.0), "mean_hbm_gb": result["resource"]["mean_prefix_hbm_gb"], "pressure_evictions": result["resource"].get("pressure_evictions", 0), "wrong_step_prefetch_gb": result["resource"]["wrong_step_prefetch_gb"], "unused_prefetch_gb": result["resource"].get("unused_prefetch_gb", 0.0), "prefill_tokens_saved": result["resource"]["prefill_tokens_saved"], "top1_accuracy": result["prediction"]["top1_accuracy"], "brier": result["prediction"].get("calibration", {}).get("brier", 0.0), "ece": result["prediction"].get("calibration", {}).get("ece", 0.0), "post_shift_ece": result["prediction"].get("post_shift_calibration", {}).get("ece", 0.0), "pre_shift_accuracy": result["prediction"]["pre_shift_accuracy"], "post_shift_accuracy": result["prediction"]["post_shift_accuracy"], } def execution_prefetch_study(config: dict[str, Any]) -> dict[str, Any]: base = ExecutionLearningConfig.from_dict(config) workflows = _common_workflows(base) candidates = [ ("Learn only / no prefetch", "none", base.transition_decay), ("Cumulative transitions", "cumulative", 1.0), ("Decayed transitions", "decayed", base.transition_decay), ("Oracle next-role", "oracle", 1.0), ] rows: list[dict[str, Any]] = [] results: dict[str, dict[str, Any]] = {} for label, policy, decay in candidates: cfg = ExecutionLearningConfig.from_dict(base.to_dict()) cfg.prefetch_policy = policy cfg.transition_decay = decay result = run_execution_learning(cfg.to_dict(), workflows) results[label] = result rows.append(_row(label, result)) return { "protocol": "common-shifted-agent-workflow-trace", "study": "online-transition-prefetch-policy-comparison", "rows": rows, "results": results, "shift_fraction": base.shift_fraction, "note": "Oracle next-role sees the realized next agent and is an upper bound; online policies update only after a transition is observed.", } def _pearson(xs: list[float], ys: list[float]) -> float: if len(xs) < 2 or len(xs) != len(ys): return 0.0 mx = mean(xs) my = mean(ys) numerator = sum((x - mx) * (y - my) for x, y in zip(xs, ys, strict=True)) sx = math.sqrt(sum((x - mx) ** 2 for x in xs)) sy = math.sqrt(sum((y - my) ** 2 for y in ys)) if sx <= 1e-12 or sy <= 1e-12: return 0.0 return numerator / (sx * sy) def execution_threshold_sweep( config: dict[str, Any], thresholds: list[float] | None = None ) -> dict[str, Any]: base = ExecutionLearningConfig.from_dict(config) base.prefetch_policy = "decayed" workflows = _common_workflows(base) values = thresholds or [0.0, 0.35, 0.50, 0.60, 0.70, 0.80, 0.90] cleaned = sorted({min(max(float(value), 0.0), 1.0) for value in values}) rows = [] for threshold in cleaned: cfg = ExecutionLearningConfig.from_dict(base.to_dict()) cfg.confidence_threshold = threshold result = run_execution_learning(cfg.to_dict(), workflows) rows.append(_row(f"threshold {threshold:.2f}", result) | {"threshold": threshold}) association = { "coverage_vs_ttft_r": _pearson( [float(row["prefetch_coverage"]) for row in rows], [float(row["p95_step_ttft_ms"]) for row in rows], ), "precision_vs_wrong_prefetch_r": _pearson( [float(row["prefetch_precision"]) for row in rows], [float(row["wrong_step_prefetch_gb"]) for row in rows], ), } return { "protocol": "common-shifted-agent-workflow-trace", "study": "prefetch-confidence-threshold-frontier", "rows": rows, "association": association, "note": "Correlations summarize this simulated sweep only; they are descriptive, not causal estimates.", } def execution_decay_sweep( config: dict[str, Any], decay_values: list[float] | None = None ) -> dict[str, Any]: base = ExecutionLearningConfig.from_dict(config) base.prefetch_policy = "decayed" workflows = _common_workflows(base) values = decay_values or [0.50, 0.65, 0.75, 0.85, 0.92, 0.97, 1.00] cleaned = sorted({min(max(float(value), 0.20), 1.0) for value in values}) rows = [] for decay in cleaned: cfg = ExecutionLearningConfig.from_dict(base.to_dict()) cfg.transition_decay = decay result = run_execution_learning(cfg.to_dict(), workflows) rows.append(_row(f"decay {decay:.2f}", result) | {"decay": decay}) best_accuracy = max(rows, key=lambda row: (row["post_shift_accuracy"], -row["p95_step_ttft_ms"])) if rows else None best_ttft = min(rows, key=lambda row: (row["p95_step_ttft_ms"], -row["post_shift_accuracy"])) if rows else None association = { "post_shift_accuracy_vs_ttft_r": _pearson( [float(row["post_shift_accuracy"]) for row in rows], [float(row["p95_step_ttft_ms"]) for row in rows], ), "post_shift_accuracy_vs_prefix_hit_r": _pearson( [float(row["post_shift_accuracy"]) for row in rows], [float(row["prefix_hit_rate"]) for row in rows], ), } return { "protocol": "common-shifted-agent-workflow-trace", "study": "transition-forgetting-rate-sweep", "rows": rows, "best_post_shift_accuracy_decay": best_accuracy["decay"] if best_accuracy else None, "best_ttft_decay": best_ttft["decay"] if best_ttft else None, "association": association, "note": "Better next-role prediction need not minimize serving latency because cache occupancy and transfer timing mediate the effect.", } def execution_planning_study(config: dict[str, Any]) -> dict[str, Any]: """Compare one-step, multi-step, utility-aware, and clairvoyant planning.""" base = ExecutionLearningConfig.from_dict(config) workflows = _common_workflows(base) candidates = [ ("Top-1 decayed", "decayed"), ("Multi-step top-k", "multistep"), ("Utility-aware multi-step", "utility"), ("Oracle future-set", "oracle_horizon"), ] rows: list[dict[str, Any]] = [] results: dict[str, dict[str, Any]] = {} for label, policy in candidates: cfg = ExecutionLearningConfig.from_dict(base.to_dict()) cfg.prefetch_policy = policy result = run_execution_learning(cfg.to_dict(), workflows) results[label] = result rows.append(_row(label, result)) best_ttft = min(rows, key=lambda row: row["p95_step_ttft_ms"]) if rows else None best_efficiency = max( rows, key=lambda row: ( row["prefill_tokens_saved"] / max(row["mean_hbm_gb"], 1e-9), -row["p95_step_ttft_ms"], ), ) if rows else None return { "protocol": "common-shifted-agent-workflow-trace", "study": "multi-step-prefetch-planning", "rows": rows, "results": results, "best_ttft_policy": best_ttft["label"] if best_ttft else None, "best_prefill_per_hbm_policy": best_efficiency["label"] if best_efficiency else None, "note": "Oracle future-set sees realized future roles only as an information upper bound; online multi-step policies roll the learned transition matrix forward without future trace access.", } def execution_horizon_sweep( config: dict[str, Any], horizon_values: list[int] | None = None ) -> dict[str, Any]: base = ExecutionLearningConfig.from_dict(config) base.prefetch_policy = "utility" workflows = _common_workflows(base) values = horizon_values or [1, 2, 3, 4, 5] cleaned = sorted({max(1, min(int(value), 6)) for value in values}) rows: list[dict[str, Any]] = [] for horizon in cleaned: cfg = ExecutionLearningConfig.from_dict(base.to_dict()) cfg.forecast_horizon = horizon result = run_execution_learning(cfg.to_dict(), workflows) rows.append(_row(f"horizon {horizon}", result) | {"horizon": horizon}) best_ttft = min(rows, key=lambda row: (row["p95_step_ttft_ms"], row["unused_prefetch_gb"])) if rows else None best_utilization = max(rows, key=lambda row: (row["prefetch_utilization"], -row["unused_prefetch_gb"])) if rows else None return { "protocol": "common-shifted-agent-workflow-trace", "study": "forecast-horizon-sweep", "rows": rows, "best_ttft_horizon": best_ttft["horizon"] if best_ttft else None, "best_utilization_horizon": best_utilization["horizon"] if best_utilization else None, "note": "Longer forecasts can expose future reuse but may spend bandwidth and cache capacity on prefixes that are not consumed before eviction.", } def execution_budget_sweep( config: dict[str, Any], budget_values: list[float] | None = None ) -> dict[str, Any]: base = ExecutionLearningConfig.from_dict(config) workflows = _common_workflows(base) values = budget_values or [0.30, 0.50, 0.75, 1.00] cleaned = sorted({min(max(float(value), 0.10), 1.50) for value in values}) policies = [ ("Top-1", "decayed"), ("Multi-step", "multistep"), ("Utility-aware", "utility"), ] rows: list[dict[str, Any]] = [] for budget in cleaned: for label, policy in policies: cfg = ExecutionLearningConfig.from_dict(base.to_dict()) cfg.prefix_cache_budget_fraction = budget cfg.prefetch_policy = policy result = run_execution_learning(cfg.to_dict(), workflows) rows.append(_row(label, result) | {"budget": budget, "policy_label": label}) winners = [] for budget in cleaned: group = [row for row in rows if abs(float(row["budget"]) - budget) < 1e-12] if group: winner = min(group, key=lambda row: (row["p95_step_ttft_ms"], row["unused_prefetch_gb"])) winners.append({"budget": budget, "policy": winner["policy_label"], "p95_step_ttft_ms": winner["p95_step_ttft_ms"]}) return { "protocol": "common-shifted-agent-workflow-trace", "study": "prefix-cache-budget-policy-sweep", "rows": rows, "winners": winners, "note": "All policies see the same workflows at each cache budget; the sweep exposes when broader forecasting helps versus when it increases cache pressure.", }