perf(model-reference): skip reparsing an unchanged fleet snapshot

- The snapshot loader separates the raw Redis read from parsing and revision verification
- `refresh_from_redis` keeps the last parsed document and returns early when the bytes match and the served state still carries that document's revision and degraded flag; every process refreshes every 5 s while the document changes about once a minute, so the pydantic parse and the sha256 run only on a change
- A corrupt document is not recorded, so it is parsed and logged again on the next refresh; a state replaced locally is applied again
This commit is contained in:
tazlin committed 2026-10-23 13:12:25 -04:00
1 parent 89f6a1587e
commit 8fb8c4ccfd
3 files changed
+106 -7

No files matched your search

+48 -3
View File
@@ -96,6 +96,20 @@ class _ImageReferenceState:
degraded: bool = False
@dataclass(frozen=True)
class _LoadedSnapshotDocument:
"""Represents the last fleet snapshot document parsed from Redis, so an unchanged one is not parsed again."""
encoded: str | bytes
"""The document as read from Redis, compared byte for byte with the next read."""
revision: str
"""The revision the document produced, compared with the served state's."""
degraded: bool
"""Whether the document was degraded, compared with the served state's flag."""
@dataclass(frozen=True)
class _TextReferenceState:
models: dict[str, dict[str, object]]
@@ -204,6 +218,7 @@ class ModelReference:
def __init__(self) -> None:
"""Start with HMR's packaged policy until production initialization loads Redis."""
self._last_stale_warning: float | None = None
self._last_loaded_document: _LoadedSnapshotDocument | None = None
self._image_state = _ImageReferenceState(
models={},
baselines=dict(_packaged_baseline_catalog().baselines),
@@ -324,20 +339,50 @@ class ModelReference:
)
def refresh_from_redis(self, *, log_missing: bool = True) -> bool:
"""Load the central fleet snapshot, retaining local state on any Redis failure."""
"""Load the central fleet snapshot, retaining local state on any Redis failure.
Every process refreshes every few seconds while the document rarely changes, so a document
byte-identical to the last one parsed is not parsed or hash-verified again when the served state
still carries the revision it produced.
"""
try:
snapshot = redis_image_reference_snapshots.load()
encoded = redis_image_reference_snapshots.load_encoded()
except Exception as err:
logger.error(f"Failed to load the fleet image-reference snapshot from Redis: {err}")
return False
if snapshot is None:
if encoded is None:
if log_missing:
logger.warning("The fleet image-reference snapshot is absent from central Redis.")
return False
if self._serves_loaded_document(encoded):
self._report_snapshot_health()
return False
try:
snapshot = redis_image_reference_snapshots.parse(encoded)
except Exception as err:
logger.error(f"Failed to load the fleet image-reference snapshot from Redis: {err}")
return False
changed = self._apply_snapshot(snapshot)
self._last_loaded_document = _LoadedSnapshotDocument(
encoded=encoded,
revision=snapshot.revision,
degraded=snapshot.degraded,
)
self._report_snapshot_health()
return changed
def _serves_loaded_document(self, encoded: str | bytes) -> bool:
"""Return whether applying this document would leave the served state unchanged.
The served state is compared as well as the bytes because other paths, such as a local
publication, replace the state without going through Redis.
"""
loaded_document = self._last_loaded_document
if loaded_document is None or loaded_document.encoded != encoded:
return False
state = self._image_state
return state.revision == loaded_document.revision and state.degraded == loaded_document.degraded
@staticmethod
def _records_from_payload(payload: dict[str, object]) -> dict[str, ImageGenerationModelRecord]:
"""Validate a PRIMARY category response into its typed HMR records."""
+19 -4
View File
@@ -133,13 +133,28 @@ class RedisImageReferenceSnapshots:
def load(self) -> ImageReferenceSnapshot | None:
"""Read directly from central Redis, deliberately bypassing the per-host cache."""
encoded = self.load_encoded()
if encoded is None:
return None
return self.parse(encoded)
@staticmethod
def load_encoded() -> str | bytes | None:
"""Return the stored snapshot document unparsed, read directly from central Redis."""
redis_backend = hr.horde_r
if redis_backend is None:
return None
raw = redis_backend.get(SNAPSHOT_KEY)
if raw is None:
return None
snapshot = ImageReferenceSnapshot.model_validate_json(raw)
return redis_backend.get(SNAPSHOT_KEY)
@staticmethod
def parse(encoded: str | bytes) -> ImageReferenceSnapshot:
"""Parse a stored snapshot document and verify its content against its revision hash.
Raises:
pydantic.ValidationError: If the document is not a valid snapshot.
ValueError: If the content does not match the revision hash.
"""
snapshot = ImageReferenceSnapshot.model_validate_json(encoded)
snapshot.verify_revision()
return snapshot
@@ -309,6 +309,45 @@ def test_missing_or_malformed_redis_value_does_not_clear_local_state(fake_redis,
assert process.reference == previous.models
def test_unchanged_redis_document_is_not_parsed_again(fake_redis, monkeypatch) -> None:
process = ModelReference()
served = _snapshot("unchanged-model")
fake_redis.horde_r.set(SNAPSHOT_KEY, served.model_dump_json())
assert process.refresh_from_redis() is True
def fail_parse(_encoded: str | bytes) -> ImageReferenceSnapshot:
raise AssertionError("an unchanged snapshot document was parsed again")
monkeypatch.setattr(model_reference_module.redis_image_reference_snapshots, "parse", fail_parse)
assert process.refresh_from_redis() is False
assert process._snapshot_revision == served.revision
def test_changed_redis_document_is_parsed_and_applied(fake_redis) -> None:
process = ModelReference()
fake_redis.horde_r.set(SNAPSHOT_KEY, _snapshot("first-model").model_dump_json())
assert process.refresh_from_redis() is True
replacement = _snapshot("second-model")
fake_redis.horde_r.set(SNAPSHOT_KEY, replacement.model_dump_json())
assert process.refresh_from_redis() is True
assert process._snapshot_revision == replacement.revision
def test_unchanged_redis_document_is_reapplied_once_local_state_diverges(fake_redis) -> None:
process = ModelReference()
served = _snapshot("redis-model")
fake_redis.horde_r.set(SNAPSHOT_KEY, served.model_dump_json())
assert process.refresh_from_redis() is True
process._apply_snapshot(_snapshot("locally-applied-model").model_copy(update={"degraded": True}))
assert process.refresh_from_redis() is True
assert process._snapshot_revision == served.revision
assert process._image_state.degraded is False
assert set(process.reference or {}) == {"redis-model"}
class _PrimaryBackend:
def __init__(self, *, fallback: bool = False) -> None:
self.fallback = fallback