355 lines
14 KiB
Python
355 lines
14 KiB
Python
"""Lightweight runtime regression tests for the SDXL Spectrum scheduler."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
from comfyui_spectrum_sdxl.config import SpectrumSDXLConfig
|
|
from comfyui_spectrum_sdxl.forecast import ChebyshevFeatureForecaster
|
|
from comfyui_spectrum_sdxl.runtime import SpectrumSDXLRuntime
|
|
from comfyui_spectrum_sdxl.sdxl import _SpectrumOuterStepController
|
|
|
|
|
|
def _make_cfg() -> SpectrumSDXLConfig:
|
|
"""Create the default test configuration."""
|
|
return SpectrumSDXLConfig(
|
|
blend_weight=0.5,
|
|
degree=4,
|
|
ridge_lambda=0.1,
|
|
window_size=2.0,
|
|
flex_window=0.75,
|
|
warmup_steps=2,
|
|
).validated()
|
|
|
|
|
|
def _step_options(
|
|
run_id,
|
|
solver_step_id: int,
|
|
time_coord: float,
|
|
actual_forward: bool | None = None,
|
|
total_steps: int = 6,
|
|
uuid: str = "stream-a",
|
|
cond: int = 0,
|
|
):
|
|
"""Create explicit outer-step context for one logical stream."""
|
|
options = {
|
|
"spectrum_run_id": run_id,
|
|
"spectrum_solver_step_id": solver_step_id,
|
|
"spectrum_time_coord": time_coord,
|
|
"spectrum_total_steps": total_steps,
|
|
"uuids": [uuid],
|
|
"cond_or_uncond": [cond],
|
|
}
|
|
if actual_forward is not None:
|
|
options["spectrum_actual_forward"] = actual_forward
|
|
return options
|
|
|
|
|
|
def test_missing_solver_step_context_fails_open() -> None:
|
|
"""Without explicit outer-step context, forecasting must stay disabled."""
|
|
runtime = SpectrumSDXLRuntime(_make_cfg())
|
|
decision = runtime.begin_step({}, torch.tensor([0.0]), (2, 8, 4, 4))
|
|
assert decision["actual_forward"] is True
|
|
assert decision["forecast_safe"] is False
|
|
assert decision["solver_step_id"] is None
|
|
assert runtime.last_info["forecast_disabled"] is True
|
|
assert runtime.last_info["forecast_disable_reason"] == "missing_solver_step_context"
|
|
|
|
|
|
def test_invalid_solver_step_context_fails_open() -> None:
|
|
"""Partially specified solver-step context must not be trusted."""
|
|
runtime = SpectrumSDXLRuntime(_make_cfg())
|
|
decision = runtime.begin_step(
|
|
{
|
|
"spectrum_run_id": "run-a",
|
|
"spectrum_solver_step_id": 0,
|
|
"spectrum_time_coord": 0.0,
|
|
"spectrum_actual_forward": 0,
|
|
},
|
|
torch.tensor([0.0]),
|
|
(2, 8, 4, 4),
|
|
)
|
|
assert decision["actual_forward"] is True
|
|
assert decision["forecast_safe"] is False
|
|
assert runtime.last_info["forecast_disable_reason"] == "invalid_solver_step_context"
|
|
|
|
|
|
def test_missing_stream_identity_fails_open() -> None:
|
|
"""Explicit step context still needs a stable logical stream identity."""
|
|
runtime = SpectrumSDXLRuntime(_make_cfg())
|
|
decision = runtime.begin_step(
|
|
{
|
|
"spectrum_run_id": "run-a",
|
|
"spectrum_solver_step_id": 0,
|
|
"spectrum_time_coord": 0.0,
|
|
"spectrum_actual_forward": False,
|
|
"spectrum_total_steps": 6,
|
|
},
|
|
torch.tensor([0.0]),
|
|
(2, 8, 4, 4),
|
|
)
|
|
assert decision["actual_forward"] is True
|
|
assert decision["forecast_safe"] is False
|
|
assert decision["stream_key"] is None
|
|
assert runtime.last_info["forecast_disable_reason"] == "missing_stream_identity"
|
|
|
|
|
|
def test_explicit_solver_step_context_allows_forecast() -> None:
|
|
"""Forecasting should work once the outer solver step is explicit."""
|
|
runtime = SpectrumSDXLRuntime(_make_cfg())
|
|
|
|
first = runtime.begin_step(_step_options("run-a", 0, 0.0, True), torch.tensor([0.0]), (2, 8, 4, 4))
|
|
runtime.observe_actual_feature(
|
|
first["stream_key"],
|
|
first["solver_step_id"],
|
|
torch.full((2, 8, 4, 4), 1.0, dtype=torch.float16),
|
|
)
|
|
|
|
second = runtime.begin_step(_step_options("run-a", 1, 1.0, True), torch.tensor([1.0]), (2, 8, 4, 4))
|
|
runtime.observe_actual_feature(
|
|
second["stream_key"],
|
|
second["solver_step_id"],
|
|
torch.full((2, 8, 4, 4), 2.0, dtype=torch.float16),
|
|
)
|
|
|
|
forecast = runtime.begin_step(_step_options("run-a", 2, 2.0, False), torch.tensor([2.0]), (2, 8, 4, 4))
|
|
assert forecast["actual_forward"] is False
|
|
assert forecast["forecast_safe"] is True
|
|
|
|
pred = runtime.predict_feature(forecast["stream_key"], forecast["solver_step_id"])
|
|
assert pred.shape == (2, 8, 4, 4)
|
|
assert torch.isfinite(pred).all()
|
|
|
|
runtime.finalize_step(forecast["stream_key"], forecast["solver_step_id"], used_forecast=True)
|
|
assert runtime.last_info["forecasted_passes"] == 1
|
|
|
|
|
|
def test_explicit_solver_step_context_without_decision_still_schedules() -> None:
|
|
"""Outer-step context alone should be enough once the controller owns step ids."""
|
|
runtime = SpectrumSDXLRuntime(_make_cfg())
|
|
|
|
first = runtime.begin_step(_step_options("run-a", 0, 0.0), torch.tensor([0.0]), (2, 8, 4, 4))
|
|
assert first["actual_forward"] is True
|
|
runtime.observe_actual_feature(
|
|
first["stream_key"],
|
|
first["solver_step_id"],
|
|
torch.full((2, 8, 4, 4), 1.0, dtype=torch.float16),
|
|
)
|
|
|
|
second = runtime.begin_step(_step_options("run-a", 1, 1.0), torch.tensor([1.0]), (2, 8, 4, 4))
|
|
assert second["actual_forward"] is True
|
|
runtime.observe_actual_feature(
|
|
second["stream_key"],
|
|
second["solver_step_id"],
|
|
torch.full((2, 8, 4, 4), 2.0, dtype=torch.float16),
|
|
)
|
|
|
|
third = runtime.begin_step(_step_options("run-a", 2, 2.0), torch.tensor([2.0]), (2, 8, 4, 4))
|
|
assert third["actual_forward"] is False
|
|
assert third["forecast_safe"] is True
|
|
|
|
|
|
def test_outer_step_controller_injects_context_and_resets_runs() -> None:
|
|
"""The sampler hook should stamp stable run and solver-step ids."""
|
|
runtime = SpectrumSDXLRuntime(_make_cfg())
|
|
controller = _SpectrumOuterStepController(
|
|
runtime=runtime,
|
|
delegate=lambda args: args["model_options"]["transformer_options"].copy(),
|
|
)
|
|
|
|
sample_sigmas = torch.linspace(1.0, 0.0, 4)
|
|
model_options = {"transformer_options": {"sample_sigmas": sample_sigmas}}
|
|
first = controller({"model_options": model_options, "sigma": torch.tensor([float(sample_sigmas[0].item())])})
|
|
second = controller({"model_options": model_options, "sigma": torch.tensor([float(sample_sigmas[1].item())])})
|
|
assert first["spectrum_run_id"] == second["spectrum_run_id"]
|
|
assert first["spectrum_solver_step_id"] == 0
|
|
assert second["spectrum_solver_step_id"] == 1
|
|
assert second["spectrum_total_steps"] == 3
|
|
|
|
next_sample_sigmas = torch.linspace(1.0, 0.0, 5)
|
|
next_model_options = {"transformer_options": {"sample_sigmas": next_sample_sigmas}}
|
|
restarted = controller(
|
|
{"model_options": next_model_options, "sigma": torch.tensor([float(next_sample_sigmas[0].item())])}
|
|
)
|
|
assert restarted["spectrum_run_id"] != first["spectrum_run_id"]
|
|
assert restarted["spectrum_solver_step_id"] == 0
|
|
assert restarted["spectrum_total_steps"] == 4
|
|
|
|
|
|
def test_outer_step_controller_same_object_restart_resets_run_state() -> None:
|
|
"""Reusing the exact same schedule tensor for a new run must reset controller state."""
|
|
runtime = SpectrumSDXLRuntime(_make_cfg())
|
|
controller = _SpectrumOuterStepController(
|
|
runtime=runtime,
|
|
delegate=lambda args: args["model_options"]["transformer_options"].copy(),
|
|
)
|
|
|
|
sample_sigmas = torch.linspace(1.0, 0.0, 6)
|
|
model_options = {"transformer_options": {"sample_sigmas": sample_sigmas}}
|
|
|
|
first = controller({"model_options": model_options, "sigma": torch.tensor([float(sample_sigmas[0].item())])})
|
|
second = controller({"model_options": model_options, "sigma": torch.tensor([float(sample_sigmas[1].item())])})
|
|
restarted = controller({"model_options": model_options, "sigma": torch.tensor([float(sample_sigmas[0].item())])})
|
|
|
|
assert first["spectrum_run_id"] == second["spectrum_run_id"]
|
|
assert second["spectrum_solver_step_id"] == 1
|
|
assert restarted["spectrum_run_id"] != first["spectrum_run_id"]
|
|
assert restarted["spectrum_solver_step_id"] == 0
|
|
assert restarted["spectrum_time_coord"] == 0.0
|
|
|
|
|
|
def test_outer_step_controller_uses_ordinal_time_coord() -> None:
|
|
"""The controller must keep time_coord in the same ordinal step space as the forecaster."""
|
|
runtime = SpectrumSDXLRuntime(_make_cfg())
|
|
controller = _SpectrumOuterStepController(
|
|
runtime=runtime,
|
|
delegate=lambda args: args["model_options"]["transformer_options"].copy(),
|
|
)
|
|
|
|
model_options = {
|
|
"transformer_options": {
|
|
"sample_sigmas": torch.tensor([14.0, 7.0, 1.5, 0.0], dtype=torch.float32),
|
|
}
|
|
}
|
|
|
|
first = controller({"model_options": model_options, "sigma": torch.tensor([14.0])})
|
|
second = controller({"model_options": model_options, "sigma": torch.tensor([7.0])})
|
|
|
|
assert first["spectrum_solver_step_id"] == 0
|
|
assert second["spectrum_solver_step_id"] == 1
|
|
assert first["spectrum_time_coord"] == 0.0
|
|
assert second["spectrum_time_coord"] == 1.0
|
|
|
|
|
|
def test_forecast_request_before_history_fails_open_per_step() -> None:
|
|
"""Outer code may request a forecast early; the runtime must still fail open."""
|
|
runtime = SpectrumSDXLRuntime(_make_cfg())
|
|
decision = runtime.begin_step(_step_options("run-a", 0, 0.0, False), torch.tensor([0.0]), (2, 8, 4, 4))
|
|
assert decision["actual_forward"] is True
|
|
assert decision["forecast_safe"] is False
|
|
|
|
|
|
def test_duplicate_actual_updates_are_deduped() -> None:
|
|
"""A repeated actual write for one stream/step must only be recorded once."""
|
|
runtime = SpectrumSDXLRuntime(_make_cfg())
|
|
decision = runtime.begin_step(_step_options("run-a", 0, 0.0, True), torch.tensor([0.0]), (2, 8, 4, 4))
|
|
assert decision["actual_forward"] is True
|
|
|
|
first_feature = torch.full((2, 8, 4, 4), 1.0, dtype=torch.float16)
|
|
second_feature = torch.full((2, 8, 4, 4), 2.0, dtype=torch.float16)
|
|
runtime.observe_actual_feature(decision["stream_key"], decision["solver_step_id"], first_feature)
|
|
runtime.observe_actual_feature(decision["stream_key"], decision["solver_step_id"], second_feature)
|
|
|
|
state = runtime.stream_states[decision["stream_key"]]
|
|
assert len(state.forecaster.history) == 1
|
|
_, stored_feature = state.forecaster.history[0]
|
|
assert torch.equal(stored_feature, first_feature)
|
|
|
|
|
|
def test_forecast_fallback_commits_actual_bookkeeping() -> None:
|
|
"""A rejected forecast should not count as a cached pass."""
|
|
runtime = SpectrumSDXLRuntime(_make_cfg())
|
|
|
|
for step_id, value in ((0, 1.0), (1, 2.0)):
|
|
decision = runtime.begin_step(
|
|
_step_options("run-a", step_id, float(step_id), True),
|
|
torch.tensor([float(step_id)]),
|
|
(2, 8, 4, 4),
|
|
)
|
|
runtime.observe_actual_feature(
|
|
decision["stream_key"],
|
|
decision["solver_step_id"],
|
|
torch.full((2, 8, 4, 4), value, dtype=torch.float16),
|
|
)
|
|
|
|
forecast = runtime.begin_step(_step_options("run-a", 2, 2.0, False), torch.tensor([2.0]), (2, 8, 4, 4))
|
|
assert forecast["actual_forward"] is False
|
|
|
|
runtime.finalize_step(forecast["stream_key"], forecast["solver_step_id"], used_forecast=False)
|
|
|
|
state = runtime.stream_states[forecast["stream_key"]]
|
|
assert runtime.last_info["forecasted_passes"] == 0
|
|
assert runtime.last_info["actual_forward_count"] == 3
|
|
assert state.num_consecutive_cached_steps == 0
|
|
assert state.decisions_by_solver_step[forecast["solver_step_id"]]["actual_forward"] is True
|
|
assert state.decisions_by_solver_step[forecast["solver_step_id"]]["finalized"] is True
|
|
|
|
|
|
def test_run_id_switch_resets_stream_state() -> None:
|
|
"""A new explicit run id must clear stale per-stream history."""
|
|
runtime = SpectrumSDXLRuntime(_make_cfg())
|
|
|
|
first = runtime.begin_step(_step_options("run-a", 0, 0.0, True), torch.tensor([0.0]), (2, 8, 4, 4))
|
|
runtime.observe_actual_feature(
|
|
first["stream_key"],
|
|
first["solver_step_id"],
|
|
torch.full((2, 8, 4, 4), 4.0, dtype=torch.float16),
|
|
)
|
|
first_state = runtime.stream_states[first["stream_key"]]
|
|
assert len(first_state.forecaster.history) == 1
|
|
|
|
restarted = runtime.begin_step(_step_options("run-b", 0, 0.0, True), torch.tensor([0.0]), (2, 8, 4, 4))
|
|
restarted_state = runtime.stream_states[restarted["stream_key"]]
|
|
assert restarted_state is not first_state
|
|
assert restarted_state.decisions_by_solver_step.keys() == {0}
|
|
assert restarted_state.observed_solver_steps == set()
|
|
assert restarted_state.forecaster.history == []
|
|
assert runtime.last_info["run_id"] == "run-b"
|
|
|
|
|
|
def test_forecaster_uses_schedule_coordinates_not_ordinal_step_index() -> None:
|
|
"""Direct forecaster predictions should accept the runtime total-steps contract."""
|
|
forecaster = ChebyshevFeatureForecaster(
|
|
degree=2,
|
|
ridge_lambda=0.1,
|
|
blend_weight=0.0,
|
|
history_size=10,
|
|
)
|
|
forecaster.update(10.0, torch.tensor([10.0]))
|
|
forecaster.update(9.0, torch.tensor([9.0]))
|
|
|
|
pred = forecaster.predict(1.0, 11)
|
|
|
|
assert torch.allclose(pred, torch.tensor([1.0]), atol=1e-5)
|
|
|
|
|
|
def test_chebyshev_prediction_varies_with_time_coord() -> None:
|
|
"""The pure Chebyshev branch must not collapse all future steps to one value."""
|
|
forecaster = ChebyshevFeatureForecaster(
|
|
degree=1,
|
|
ridge_lambda=0.0,
|
|
blend_weight=1.0,
|
|
history_size=10,
|
|
)
|
|
forecaster.update(0.0, torch.tensor([0.0]))
|
|
forecaster.update(1.0, torch.tensor([1.0]))
|
|
|
|
pred_two = forecaster.predict(2.0, 4)
|
|
pred_three = forecaster.predict(3.0, 4)
|
|
|
|
assert torch.allclose(pred_two, torch.tensor([2.0]), atol=1e-5)
|
|
assert torch.allclose(pred_three, torch.tensor([3.0]), atol=1e-5)
|
|
|
|
|
|
def main() -> None:
|
|
"""Run the lightweight regression suite without external test tooling."""
|
|
test_missing_solver_step_context_fails_open()
|
|
test_invalid_solver_step_context_fails_open()
|
|
test_missing_stream_identity_fails_open()
|
|
test_explicit_solver_step_context_allows_forecast()
|
|
test_explicit_solver_step_context_without_decision_still_schedules()
|
|
test_outer_step_controller_injects_context_and_resets_runs()
|
|
test_outer_step_controller_same_object_restart_resets_run_state()
|
|
test_outer_step_controller_uses_ordinal_time_coord()
|
|
test_forecast_request_before_history_fails_open_per_step()
|
|
test_duplicate_actual_updates_are_deduped()
|
|
test_forecast_fallback_commits_actual_bookkeeping()
|
|
test_run_id_switch_resets_stream_state()
|
|
test_forecaster_uses_schedule_coordinates_not_ordinal_step_index()
|
|
test_chebyshev_prediction_varies_with_time_coord()
|
|
print("ok")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|