321 lines
12 KiB
Python
321 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import math
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict, Optional
|
|
|
|
import torch
|
|
|
|
from .config import SpectrumConfig
|
|
from .forecast import ChebyshevSpectrumForecaster
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class RuntimeStats:
|
|
actual_forward_count: int = 0
|
|
forecasted_count: int = 0
|
|
total_steps: int = 0
|
|
current_window: float = 0.0
|
|
run_id: int = 0
|
|
forecast_disabled: bool = False
|
|
disable_reason: Optional[str] = None
|
|
sampler_name: Optional[str] = None
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class _ActiveRun:
|
|
run_id: int
|
|
sampler_name: str
|
|
total_steps: int
|
|
schedule_values: tuple[float, ...]
|
|
schedule_coords: tuple[float, ...]
|
|
supports_solver_steps: bool
|
|
next_solver_step_id: int = 0
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class _ActiveStep:
|
|
solver_step_id: int
|
|
time_coord: float
|
|
decision: Dict[str, Any]
|
|
expected_shape: Optional[tuple[int, ...]] = None
|
|
branch_signature: Optional[tuple[Any, ...]] = None
|
|
hook_call_count: int = 0
|
|
observed_actual: bool = False
|
|
used_forecast: bool = False
|
|
predicted_feature: Optional[torch.Tensor] = None
|
|
|
|
|
|
class SpectrumRuntime:
|
|
def __init__(self, cfg: SpectrumConfig):
|
|
self.cfg = cfg.validate()
|
|
self.forecaster = ChebyshevSpectrumForecaster(
|
|
degree=self.cfg.degree,
|
|
ridge_lambda=self.cfg.ridge_lambda,
|
|
max_history=self.cfg.max_history,
|
|
)
|
|
self.run_id = 0
|
|
self.stats = RuntimeStats(current_window=float(self.cfg.window_size))
|
|
self._active_run: Optional[_ActiveRun] = None
|
|
self._active_steps: Dict[int, _ActiveStep] = {}
|
|
self._reset_scheduler_state()
|
|
|
|
@property
|
|
def min_fit_points(self) -> int:
|
|
return max(2, self.cfg.degree + 1)
|
|
|
|
@property
|
|
def active_run_id(self) -> Optional[int]:
|
|
if self._active_run is None:
|
|
return None
|
|
return self._active_run.run_id
|
|
|
|
@property
|
|
def next_solver_step_id(self) -> int:
|
|
if self._active_run is None:
|
|
return 0
|
|
return self._active_run.next_solver_step_id
|
|
|
|
@property
|
|
def active_run_supports_solver_steps(self) -> bool:
|
|
if self._active_run is None:
|
|
return False
|
|
return bool(self._active_run.supports_solver_steps)
|
|
|
|
def _reset_scheduler_state(self) -> None:
|
|
self.curr_ws = float(self.cfg.window_size)
|
|
self.num_consecutive_cached_steps = 0
|
|
self.forecast_disabled = False
|
|
self.forecast_disable_reason: Optional[str] = None
|
|
self.forecaster.reset()
|
|
self._active_steps = {}
|
|
self.stats.current_window = float(self.cfg.window_size)
|
|
self.stats.forecast_disabled = False
|
|
self.stats.disable_reason = None
|
|
|
|
def reset_all(self) -> None:
|
|
self.run_id += 1
|
|
self.stats = RuntimeStats(current_window=float(self.cfg.window_size), run_id=self.run_id)
|
|
self._active_run = None
|
|
self._reset_scheduler_state()
|
|
|
|
def _disable_forecasting(self, reason: str) -> None:
|
|
if self.forecast_disabled:
|
|
return
|
|
self.forecast_disabled = True
|
|
self.forecast_disable_reason = reason
|
|
self.num_consecutive_cached_steps = 0
|
|
self.curr_ws = float(self.cfg.window_size)
|
|
self.forecaster.reset()
|
|
for step in self._active_steps.values():
|
|
step.predicted_feature = None
|
|
self.stats.current_window = self.curr_ws
|
|
self.stats.forecast_disabled = True
|
|
self.stats.disable_reason = reason
|
|
|
|
def num_steps(self) -> int:
|
|
if self.stats.total_steps > 0:
|
|
return self.stats.total_steps
|
|
return 50
|
|
|
|
@staticmethod
|
|
def _build_schedule_coords(sample_sigmas: torch.Tensor) -> tuple[tuple[float, ...], tuple[float, ...]]:
|
|
values = tuple(float(v) for v in sample_sigmas.detach().flatten().tolist()[:-1])
|
|
if not values:
|
|
return (), ()
|
|
|
|
start = values[0]
|
|
end = values[-1]
|
|
denom = end - start
|
|
if abs(denom) < 1e-12:
|
|
coords = tuple(0.0 for _ in values)
|
|
else:
|
|
coords = tuple(((v - start) / denom) * 2.0 - 1.0 for v in values)
|
|
return values, coords
|
|
|
|
def time_coord_for_step(self, solver_step_id: int) -> float:
|
|
if self._active_run is None:
|
|
raise RuntimeError("Spectrum runtime is not inside an active run.")
|
|
idx = int(solver_step_id)
|
|
if idx < 0 or idx >= len(self._active_run.schedule_coords):
|
|
raise RuntimeError(f"Spectrum solver step {solver_step_id} is outside the active schedule.")
|
|
return float(self._active_run.schedule_coords[idx])
|
|
|
|
def start_run(self, sample_sigmas: torch.Tensor, sampler_name: str, *, supports_solver_steps: bool) -> int:
|
|
self.run_id += 1
|
|
schedule_values, schedule_coords = self._build_schedule_coords(sample_sigmas)
|
|
total_steps = max(len(schedule_coords), 1)
|
|
self.stats = RuntimeStats(
|
|
current_window=float(self.cfg.window_size),
|
|
total_steps=total_steps,
|
|
run_id=self.run_id,
|
|
sampler_name=sampler_name,
|
|
)
|
|
self._active_run = _ActiveRun(
|
|
run_id=self.run_id,
|
|
sampler_name=sampler_name,
|
|
total_steps=total_steps,
|
|
schedule_values=schedule_values,
|
|
schedule_coords=schedule_coords,
|
|
supports_solver_steps=bool(supports_solver_steps),
|
|
)
|
|
self._reset_scheduler_state()
|
|
self.stats.run_id = self.run_id
|
|
self.stats.total_steps = total_steps
|
|
self.stats.sampler_name = sampler_name
|
|
if not supports_solver_steps:
|
|
self._disable_forecasting(
|
|
f"sampler {sampler_name!r} does not expose one predict_noise call per solver step"
|
|
)
|
|
return self.run_id
|
|
|
|
def end_run(self, run_id: int) -> None:
|
|
if self._active_run is None or self._active_run.run_id != int(run_id):
|
|
return
|
|
self._active_run = None
|
|
self._active_steps = {}
|
|
|
|
def begin_solver_step(self, run_id: int, solver_step_id: int, time_coord: float, total_steps: int) -> Dict[str, Any]:
|
|
if self._active_run is None or self._active_run.run_id != int(run_id):
|
|
raise RuntimeError("Spectrum runtime is not inside the requested sampling run.")
|
|
|
|
if int(total_steps) != self._active_run.total_steps:
|
|
self._disable_forecasting("solver-step total_steps changed inside one sampling run")
|
|
|
|
expected_time_coord = self.time_coord_for_step(int(solver_step_id))
|
|
if not math.isclose(float(time_coord), expected_time_coord, rel_tol=0.0, abs_tol=1e-8):
|
|
self._disable_forecasting("solver-step time_coord did not match the active schedule")
|
|
|
|
existing = self._active_steps.get(int(solver_step_id))
|
|
if existing is not None:
|
|
return existing.decision
|
|
|
|
if int(solver_step_id) != self._active_run.next_solver_step_id:
|
|
self._disable_forecasting("solver-step ids are not sequential within the sampling run")
|
|
|
|
actual_forward = True
|
|
if (
|
|
not self.forecast_disabled
|
|
and int(solver_step_id) >= self.cfg.warmup_steps
|
|
and self.forecaster.ready(self.min_fit_points)
|
|
):
|
|
ws_floor = max(1, int(math.floor(self.curr_ws)))
|
|
actual_forward = ((self.num_consecutive_cached_steps + 1) % ws_floor) == 0
|
|
|
|
if self.forecast_disabled or not self.forecaster.ready(self.min_fit_points):
|
|
actual_forward = True
|
|
|
|
decision = {
|
|
"run_id": int(run_id),
|
|
"solver_step_id": int(solver_step_id),
|
|
"time_coord": float(time_coord),
|
|
"total_steps": int(total_steps),
|
|
"actual_forward": bool(actual_forward),
|
|
"forecast_disabled": self.forecast_disabled,
|
|
}
|
|
self._active_steps[int(solver_step_id)] = _ActiveStep(
|
|
solver_step_id=int(solver_step_id),
|
|
time_coord=float(time_coord),
|
|
decision=decision,
|
|
)
|
|
self._active_run.next_solver_step_id = int(solver_step_id) + 1
|
|
self.stats.current_window = self.curr_ws
|
|
return decision
|
|
|
|
def get_step_decision(self, run_id: int, solver_step_id: int) -> Optional[Dict[str, Any]]:
|
|
step = self._active_steps.get(int(solver_step_id))
|
|
if step is None or step.decision["run_id"] != int(run_id):
|
|
return None
|
|
return step.decision
|
|
|
|
def step_used_forecast(self, run_id: int, solver_step_id: int) -> bool:
|
|
return self._require_active_step(run_id, solver_step_id).used_forecast
|
|
|
|
def register_model_hook_call(
|
|
self,
|
|
run_id: int,
|
|
solver_step_id: int,
|
|
*,
|
|
expected_shape: tuple[int, ...],
|
|
branch_signature: Optional[tuple[Any, ...]] = None,
|
|
) -> None:
|
|
step = self._require_active_step(run_id, solver_step_id)
|
|
step.hook_call_count += 1
|
|
if step.hook_call_count > 1:
|
|
self._disable_forecasting("multiple model-hook calls observed within one solver step")
|
|
if step.expected_shape is None:
|
|
step.expected_shape = tuple(expected_shape)
|
|
elif tuple(expected_shape) != step.expected_shape:
|
|
self._disable_forecasting("model-hook feature shape changed within one solver step")
|
|
|
|
if branch_signature is None:
|
|
return
|
|
if step.branch_signature is None:
|
|
step.branch_signature = branch_signature
|
|
elif branch_signature != step.branch_signature:
|
|
self._disable_forecasting("model-hook branch signature changed within one solver step")
|
|
|
|
def observe_actual_feature(self, run_id: int, solver_step_id: int, feature: torch.Tensor) -> None:
|
|
step = self._require_active_step(run_id, solver_step_id)
|
|
step.observed_actual = True
|
|
step.used_forecast = False
|
|
step.predicted_feature = None
|
|
if self.forecast_disabled:
|
|
return
|
|
self.forecaster.update(step.time_coord, feature)
|
|
|
|
def predict_feature(
|
|
self,
|
|
run_id: int,
|
|
solver_step_id: int,
|
|
*,
|
|
expected_shape: Optional[tuple[int, ...]] = None,
|
|
) -> Optional[torch.Tensor]:
|
|
step = self._require_active_step(run_id, solver_step_id)
|
|
if self.forecast_disabled or not self.forecaster.ready(self.min_fit_points):
|
|
return None
|
|
|
|
if step.predicted_feature is None:
|
|
step.predicted_feature = self.forecaster.predict(
|
|
time_coord=step.time_coord,
|
|
blend_weight=self.cfg.blend_weight,
|
|
)
|
|
|
|
if expected_shape is not None and tuple(step.predicted_feature.shape) != tuple(expected_shape):
|
|
self._disable_forecasting("predicted feature shape did not match the current solver-step input")
|
|
return None
|
|
|
|
step.used_forecast = True
|
|
return step.predicted_feature
|
|
|
|
def finalize_solver_step(self, run_id: int, solver_step_id: int, *, used_forecast: bool) -> None:
|
|
step = self._require_active_step(run_id, solver_step_id)
|
|
if bool(used_forecast):
|
|
step.used_forecast = True
|
|
if step.decision["actual_forward"] and not step.observed_actual:
|
|
self._disable_forecasting("solver step requested an actual forward but no actual feature was observed")
|
|
if not step.observed_actual and not step.used_forecast:
|
|
self._disable_forecasting("solver step finished without an actual feature or a forecasted feature")
|
|
|
|
if step.used_forecast:
|
|
self.num_consecutive_cached_steps += 1
|
|
self.stats.forecasted_count += 1
|
|
step.decision["actual_forward"] = False
|
|
else:
|
|
if not self.forecast_disabled and step.solver_step_id >= self.cfg.warmup_steps:
|
|
self.curr_ws = round(self.curr_ws + float(self.cfg.flex_window), 6)
|
|
self.num_consecutive_cached_steps = 0
|
|
self.stats.actual_forward_count += 1
|
|
step.decision["actual_forward"] = True
|
|
|
|
self.stats.current_window = self.curr_ws
|
|
self._active_steps.pop(int(solver_step_id), None)
|
|
|
|
def _require_active_step(self, run_id: int, solver_step_id: int) -> _ActiveStep:
|
|
if self._active_run is None or self._active_run.run_id != int(run_id):
|
|
raise RuntimeError("Spectrum runtime is not inside the requested sampling run.")
|
|
step = self._active_steps.get(int(solver_step_id))
|
|
if step is None:
|
|
raise RuntimeError(f"Spectrum solver step {solver_step_id} is not active.")
|
|
return step
|