feat(generator): добавлен артефакт стартовой истории
- Зачем: - чистый стенд должен восстанавливать стартовую историю без повторной генерации. - Что: - добавлены export/import артефакта через Kafka и compact-топики. - добавлена manifest-aware сверка ClickHouse и защита от смешения state. - добавлен runbook использования стартовой истории. - Проверка: - uv run --with-requirements generator/requirements.txt pytest generator/tests -q. - bash -n scripts/check_startup_history_manifest.sh scripts/export_startup_history_artifact.sh scripts/import_startup_history_artifact.sh. - git diff --check.
This commit is contained in:
@@ -92,6 +92,11 @@ class Config:
|
||||
run_mode: str = field(
|
||||
default_factory=lambda: os.getenv("GEN_RUN_MODE", "live")
|
||||
)
|
||||
startup_history_artifact: Path | None = field(
|
||||
default_factory=lambda: Path(os.getenv("GEN_STARTUP_HISTORY_ARTIFACT"))
|
||||
if os.getenv("GEN_STARTUP_HISTORY_ARTIFACT")
|
||||
else None
|
||||
)
|
||||
|
||||
def __post_init__(self):
|
||||
if self.tick_seconds < 1:
|
||||
|
||||
@@ -62,6 +62,16 @@ def _kafka_importer():
|
||||
return _facade_attr("_import_kafka", _import_kafka)
|
||||
|
||||
|
||||
def kafka_event_value_json(event: dict) -> str:
|
||||
"""Возвращает JSON value ровно в формате Kafka producer генератора."""
|
||||
return json.dumps(event)
|
||||
|
||||
|
||||
def kafka_event_value_bytes(event: dict) -> bytes:
|
||||
"""Возвращает Kafka value bytes для события генератора."""
|
||||
return kafka_event_value_json(event).encode("utf-8")
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchRecord:
|
||||
"""Запись об отправленном батче."""
|
||||
@@ -380,7 +390,7 @@ class KafkaPublisher:
|
||||
logger.info(f"Connecting to Kafka at {self.bootstrap_servers}")
|
||||
self.producer = KafkaProducerCls(
|
||||
bootstrap_servers=self.bootstrap_servers,
|
||||
value_serializer=lambda v: json.dumps(v).encode("utf-8"),
|
||||
value_serializer=kafka_event_value_bytes,
|
||||
key_serializer=lambda k: k.encode("utf-8") if k else None,
|
||||
batch_size=16384,
|
||||
linger_ms=100,
|
||||
|
||||
@@ -1,7 +1,5 @@
|
||||
"""Основной сервисный цикл генератора."""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
@@ -27,108 +25,19 @@ from clickstream_generator.metrics import (
|
||||
METRICS_TICK_DURATION,
|
||||
)
|
||||
from clickstream_generator.runtime import TickStreamGenerator
|
||||
from clickstream_generator.startup_history_artifact import (
|
||||
StartupHistoryArtifactBuilder,
|
||||
build_manifest,
|
||||
generation_settings_from_config,
|
||||
write_startup_history_artifact,
|
||||
)
|
||||
|
||||
|
||||
logger = logging.getLogger("generator")
|
||||
|
||||
|
||||
class _TopicManifestStats:
|
||||
"""Накопительные счётчики одного Kafka-топика для manifest."""
|
||||
|
||||
def __init__(self):
|
||||
self.rows = 0
|
||||
self.min_event_timestamp: str | None = None
|
||||
self.max_event_timestamp: str | None = None
|
||||
self._checksum = hashlib.sha256()
|
||||
|
||||
@property
|
||||
def checksum(self) -> str:
|
||||
return self._checksum.hexdigest()
|
||||
|
||||
def add(self, event: dict, event_timestamp: str | None = None) -> None:
|
||||
self.rows += 1
|
||||
self._checksum.update(
|
||||
json.dumps(event, sort_keys=True, ensure_ascii=True).encode("utf-8")
|
||||
)
|
||||
timestamp = event_timestamp if event_timestamp is not None else event.get("event_timestamp")
|
||||
self.add_timestamp(timestamp)
|
||||
|
||||
def add_timestamp(self, timestamp: str | None) -> None:
|
||||
if timestamp is None:
|
||||
return
|
||||
if self.min_event_timestamp is None or timestamp < self.min_event_timestamp:
|
||||
self.min_event_timestamp = timestamp
|
||||
if self.max_event_timestamp is None or timestamp > self.max_event_timestamp:
|
||||
self.max_event_timestamp = timestamp
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"rows": self.rows,
|
||||
"min_event_timestamp": self.min_event_timestamp,
|
||||
"max_event_timestamp": self.max_event_timestamp,
|
||||
"checksum_sha256": self.checksum,
|
||||
}
|
||||
|
||||
|
||||
class _ManifestCounters:
|
||||
"""Счётчики стартовой истории для manifest."""
|
||||
|
||||
def __init__(self):
|
||||
self.topic_stats = {
|
||||
topic: _TopicManifestStats()
|
||||
for topic in ("browser_events", "location_events", "device_events", "geo_events")
|
||||
}
|
||||
self.click_ids: set[str] = set()
|
||||
self.user_ids: set[str] = set()
|
||||
|
||||
def add_batch(self, batch: dict[str, list[dict]]) -> None:
|
||||
browser_events = batch.get("browser_events", [])
|
||||
event_timestamps = {
|
||||
event["event_id"]: event.get("event_timestamp")
|
||||
for event in browser_events
|
||||
if event.get("event_id") and event.get("event_timestamp")
|
||||
}
|
||||
click_timestamps: dict[str, list[str]] = {}
|
||||
for event in browser_events:
|
||||
click_id = event.get("click_id")
|
||||
timestamp = event.get("event_timestamp")
|
||||
if click_id and timestamp:
|
||||
click_timestamps.setdefault(click_id, []).append(timestamp)
|
||||
|
||||
for topic, events in batch.items():
|
||||
stats = self.topic_stats[topic]
|
||||
for event in events:
|
||||
event_timestamp = event.get("event_timestamp")
|
||||
if event_timestamp is None and topic == "location_events":
|
||||
event_timestamp = event_timestamps.get(event.get("event_id"))
|
||||
if event_timestamp is None and topic in ("device_events", "geo_events"):
|
||||
timestamps = click_timestamps.get(event.get("click_id"))
|
||||
if timestamps:
|
||||
event_timestamp = min(timestamps)
|
||||
stats.add_timestamp(max(timestamps))
|
||||
stats.add(event, event_timestamp=event_timestamp)
|
||||
click_id = event.get("click_id")
|
||||
if topic == "browser_events" and click_id:
|
||||
self.click_ids.add(click_id)
|
||||
user_id = event.get("user_domain_id")
|
||||
if topic == "device_events" and user_id:
|
||||
self.user_ids.add(user_id)
|
||||
|
||||
def to_manifest_topics(self) -> dict:
|
||||
return {
|
||||
topic: stats.to_dict()
|
||||
for topic, stats in self.topic_stats.items()
|
||||
}
|
||||
|
||||
def to_manifest_totals(self) -> dict:
|
||||
browser_stats = self.topic_stats["browser_events"]
|
||||
return {
|
||||
"events": browser_stats.rows,
|
||||
"visits": len(self.click_ids),
|
||||
"users": len(self.user_ids),
|
||||
"min_event_timestamp": browser_stats.min_event_timestamp,
|
||||
"max_event_timestamp": browser_stats.max_event_timestamp,
|
||||
}
|
||||
class IncompatibleStateError(ValueError):
|
||||
"""Читаемый state относится к другому миру генерации."""
|
||||
|
||||
|
||||
class GeneratorService:
|
||||
@@ -201,8 +110,15 @@ class GeneratorService:
|
||||
model_t_end=model_t_end,
|
||||
)
|
||||
elif self._is_startup_history_marker(restored_state):
|
||||
raise ValueError(
|
||||
"startup-history state without matching manifest"
|
||||
mismatches = self._startup_history_mismatch_fields(
|
||||
restored_state,
|
||||
manifest,
|
||||
)
|
||||
raise IncompatibleStateError(
|
||||
"startup-history state mismatch: "
|
||||
+ ", ".join(mismatches)
|
||||
+ "; set GEN_STATE_RESET=true to start a new "
|
||||
"world intentionally"
|
||||
)
|
||||
else:
|
||||
self._restore_live_state(
|
||||
@@ -214,6 +130,13 @@ class GeneratorService:
|
||||
f"model_time={self._model_time.isoformat()}, "
|
||||
f"last_batch_id={restored_state.last_batch_id}"
|
||||
)
|
||||
except IncompatibleStateError:
|
||||
logger.error(
|
||||
"Readable generator state is incompatible with "
|
||||
"current settings; set GEN_STATE_RESET=true to "
|
||||
"start a new world intentionally."
|
||||
)
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"State data was invalid, starting fresh: {e}"
|
||||
@@ -290,27 +213,36 @@ class GeneratorService:
|
||||
"""Проверяет, что state относится к текущей конфигурации live-запуска."""
|
||||
mismatches = []
|
||||
if state.gen_seed != self.config.seed:
|
||||
mismatches.append("gen_seed")
|
||||
mismatches.append("GEN_SEED")
|
||||
if self._as_aware_utc(state.model_t0) != self.config.model_t0:
|
||||
mismatches.append("model_t0")
|
||||
mismatches.append("GEN_MODEL_T0")
|
||||
if state.model_timezone != self.config.model_timezone:
|
||||
mismatches.append("model_timezone")
|
||||
mismatches.append("GEN_MODEL_TIMEZONE")
|
||||
if abs(state.model_time_speed - self.config.model_time_speed) > 1e-9:
|
||||
mismatches.append("model_time_speed")
|
||||
mismatches.append("GEN_MODEL_TIME_SPEED")
|
||||
|
||||
if mismatches:
|
||||
raise ValueError(
|
||||
"state config mismatch: " + ", ".join(mismatches)
|
||||
raise IncompatibleStateError(
|
||||
"state config mismatch: "
|
||||
+ ", ".join(mismatches)
|
||||
+ "; set GEN_STATE_RESET=true to start a new world intentionally"
|
||||
)
|
||||
|
||||
def _is_startup_history_state(self, state, manifest: dict | None) -> bool:
|
||||
"""Проверяет, что state совпадает со слепком стартовой истории."""
|
||||
return not self._startup_history_mismatch_fields(state, manifest)
|
||||
|
||||
def _startup_history_mismatch_fields(self, state, manifest: dict | None) -> list[str]:
|
||||
"""Возвращает поля, по которым startup-history state не совпал."""
|
||||
mismatches = []
|
||||
if not manifest or manifest.get("run_mode") != "backfill":
|
||||
return False
|
||||
return ["generator_startup_history_manifest"]
|
||||
|
||||
expected_state = manifest.get("state") or {}
|
||||
if manifest.get("state_version") != state.version:
|
||||
mismatches.append("state_version")
|
||||
if expected_state.get("last_batch_id") != state.last_batch_id:
|
||||
return False
|
||||
mismatches.append("state.last_batch_id")
|
||||
|
||||
try:
|
||||
model_t_end = self._as_aware_utc(
|
||||
@@ -320,32 +252,34 @@ class GeneratorService:
|
||||
datetime.fromisoformat(manifest["model_t0"])
|
||||
)
|
||||
except (KeyError, TypeError, ValueError):
|
||||
return False
|
||||
return ["manifest.model_t0", "manifest.model_t_end"]
|
||||
|
||||
if self._as_aware_utc(state.model_timestamp) != model_t_end:
|
||||
return False
|
||||
mismatches.append("GEN_MODEL_T_END")
|
||||
if self._as_aware_utc(state.model_t0) != model_t0:
|
||||
return False
|
||||
mismatches.append("GEN_MODEL_T0")
|
||||
if state.gen_seed != manifest.get("gen_seed"):
|
||||
return False
|
||||
mismatches.append("GEN_SEED")
|
||||
if state.model_timezone != manifest.get("model_timezone"):
|
||||
return False
|
||||
mismatches.append("GEN_MODEL_TIMEZONE")
|
||||
if (
|
||||
self.config.model_t_end is not None
|
||||
and self.config.model_t_end != model_t_end
|
||||
):
|
||||
return False
|
||||
mismatches.append("GEN_MODEL_T_END")
|
||||
if model_t0 != self.config.model_t0:
|
||||
return False
|
||||
mismatches.append("GEN_MODEL_T0")
|
||||
if manifest.get("gen_seed") != self.config.seed:
|
||||
return False
|
||||
mismatches.append("GEN_SEED")
|
||||
if manifest.get("model_timezone") != self.config.model_timezone:
|
||||
return False
|
||||
mismatches.append("GEN_MODEL_TIMEZONE")
|
||||
|
||||
settings = manifest.get("generation_settings") or {}
|
||||
if abs(state.model_time_speed - settings.get("model_time_speed", -1)) > 1e-9:
|
||||
return False
|
||||
return self._generation_settings() == settings
|
||||
mismatches.append("GEN_MODEL_TIME_SPEED")
|
||||
if self._generation_settings() != settings:
|
||||
mismatches.append("generation_settings")
|
||||
return list(dict.fromkeys(mismatches))
|
||||
|
||||
def _is_startup_history_marker(self, state) -> bool:
|
||||
"""Отличает state стартовой истории от обычного live-state."""
|
||||
@@ -400,7 +334,7 @@ class GeneratorService:
|
||||
self.config.model_t0.isoformat(),
|
||||
self.config.model_t_end.isoformat(),
|
||||
)
|
||||
counters = _ManifestCounters()
|
||||
artifact_builder = StartupHistoryArtifactBuilder()
|
||||
|
||||
while self._model_time < self.config.model_t_end:
|
||||
self._tick += 1
|
||||
@@ -415,7 +349,7 @@ class GeneratorService:
|
||||
)
|
||||
total_sent, sent_counts, status = self._publish_batch(batch)
|
||||
self._raise_on_backfill_publish_error(batch_id, status, sent_counts)
|
||||
counters.add_batch(batch)
|
||||
artifact_builder.add_batch(batch)
|
||||
self._write_batch_history(
|
||||
batch_id=batch_id,
|
||||
started_at=started_at,
|
||||
@@ -441,7 +375,7 @@ class GeneratorService:
|
||||
final_status,
|
||||
sent_counts,
|
||||
)
|
||||
counters.add_batch(final_batch)
|
||||
artifact_builder.add_batch(final_batch)
|
||||
self._write_batch_history(
|
||||
batch_id=batch_id,
|
||||
started_at=started_at,
|
||||
@@ -455,7 +389,7 @@ class GeneratorService:
|
||||
self.publisher.flush()
|
||||
|
||||
self._model_time = self.config.model_t_end
|
||||
state_batch_id = self._startup_state_batch_id(counters)
|
||||
state_batch_id = self._startup_state_batch_id(artifact_builder.counters)
|
||||
state = self.stream.to_state(
|
||||
tick=self._tick,
|
||||
rng_state=self.generator.rng.getstate(),
|
||||
@@ -469,8 +403,9 @@ class GeneratorService:
|
||||
gen_seed=self.config.seed,
|
||||
)
|
||||
|
||||
manifest = self._build_startup_history_manifest(
|
||||
counters=counters,
|
||||
manifest = build_manifest(
|
||||
config=self.config,
|
||||
counters=artifact_builder.counters,
|
||||
state=state,
|
||||
)
|
||||
if self.manifest_manager:
|
||||
@@ -481,6 +416,17 @@ class GeneratorService:
|
||||
self.state_manager.save(state)
|
||||
self.state_manager.flush()
|
||||
|
||||
if self.config.startup_history_artifact is not None:
|
||||
artifact = artifact_builder.to_artifact(manifest=manifest, state=state)
|
||||
write_startup_history_artifact(
|
||||
self.config.startup_history_artifact,
|
||||
artifact,
|
||||
)
|
||||
logger.info(
|
||||
"Startup history artifact written: %s",
|
||||
self.config.startup_history_artifact,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Backfill completed: events=%s, visits=%s, users=%s, final_sent=%s",
|
||||
manifest["totals"]["events"],
|
||||
@@ -565,40 +511,10 @@ class GeneratorService:
|
||||
return f"startup-history-{digest}"
|
||||
|
||||
def _build_startup_history_manifest(self, counters, state) -> dict:
|
||||
return {
|
||||
"manifest_version": "1.0",
|
||||
"generated_at": datetime.now(timezone.utc).isoformat(),
|
||||
"gen_seed": self.config.seed,
|
||||
"model_t0": self.config.model_t0.isoformat(),
|
||||
"model_t_end": self.config.model_t_end.isoformat(),
|
||||
"model_timezone": self.config.model_timezone,
|
||||
"run_mode": "backfill",
|
||||
"generation_settings": self._generation_settings(),
|
||||
"state_version": state.version,
|
||||
"state": {
|
||||
"topic": KafkaStateManager.STATE_TOPIC,
|
||||
"key": KafkaStateManager.STATE_KEY,
|
||||
"last_batch_id": state.last_batch_id,
|
||||
"model_timestamp": state.model_timestamp.isoformat(),
|
||||
},
|
||||
"topics": counters.to_manifest_topics(),
|
||||
"totals": counters.to_manifest_totals(),
|
||||
}
|
||||
return build_manifest(config=self.config, counters=counters, state=state)
|
||||
|
||||
def _generation_settings(self) -> dict:
|
||||
return {
|
||||
"tick_seconds": self.config.tick_seconds,
|
||||
"lambda_base_per_min": self.config.lambda_base_per_min,
|
||||
"jitter_pct": self.config.jitter_pct,
|
||||
"min_events_per_tick": self.config.min_events_per_tick,
|
||||
"max_events_per_tick": self.config.max_events_per_tick,
|
||||
"max_session_events": self.config.max_session_events,
|
||||
"max_active_sessions": self.config.max_active_sessions,
|
||||
"population_max": self.config.population_max,
|
||||
"p_new_user": self.config.p_new_user,
|
||||
"min_return_minutes": self.config.min_return_minutes,
|
||||
"model_time_speed": self.config.model_time_speed,
|
||||
}
|
||||
return generation_settings_from_config(self.config)
|
||||
|
||||
def _main_loop(self):
|
||||
"""Основной цикл тиков."""
|
||||
|
||||
@@ -0,0 +1,642 @@
|
||||
"""Портативный артефакт стартовой истории."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
from copy import deepcopy
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from clickstream_generator.config import Config
|
||||
from clickstream_generator.kafka_io import (
|
||||
KafkaStateManager,
|
||||
KafkaStartupHistoryManifest,
|
||||
ensure_topics,
|
||||
kafka_event_value_json,
|
||||
_kafka_importer,
|
||||
)
|
||||
from clickstream_generator.state import GeneratorState
|
||||
|
||||
|
||||
logger = logging.getLogger("generator")
|
||||
|
||||
ARTIFACT_KIND = "clickstream_generator_startup_history"
|
||||
ARTIFACT_VERSION = "1.0"
|
||||
TOPICS = ("browser_events", "location_events", "device_events", "geo_events")
|
||||
IMPORT_TOPICS = (
|
||||
"browser_events",
|
||||
"location_events",
|
||||
"device_events",
|
||||
"geo_events",
|
||||
KafkaStateManager.STATE_TOPIC,
|
||||
KafkaStartupHistoryManifest.MANIFEST_TOPIC,
|
||||
)
|
||||
|
||||
|
||||
class _TopicManifestStats:
|
||||
"""Накопительные счётчики одного Kafka-топика для manifest."""
|
||||
|
||||
def __init__(self):
|
||||
self.rows = 0
|
||||
self.min_event_timestamp: str | None = None
|
||||
self.max_event_timestamp: str | None = None
|
||||
self._checksum = hashlib.sha256()
|
||||
|
||||
@property
|
||||
def checksum(self) -> str:
|
||||
return self._checksum.hexdigest()
|
||||
|
||||
def add(self, event: dict, event_timestamp: str | None = None) -> None:
|
||||
self.rows += 1
|
||||
self._checksum.update(
|
||||
json.dumps(event, sort_keys=True, ensure_ascii=True).encode("utf-8")
|
||||
)
|
||||
timestamp = (
|
||||
event_timestamp
|
||||
if event_timestamp is not None
|
||||
else event.get("event_timestamp")
|
||||
)
|
||||
self.add_timestamp(timestamp)
|
||||
|
||||
def add_timestamp(self, timestamp: str | None) -> None:
|
||||
if timestamp is None:
|
||||
return
|
||||
if self.min_event_timestamp is None or timestamp < self.min_event_timestamp:
|
||||
self.min_event_timestamp = timestamp
|
||||
if self.max_event_timestamp is None or timestamp > self.max_event_timestamp:
|
||||
self.max_event_timestamp = timestamp
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"rows": self.rows,
|
||||
"min_event_timestamp": self.min_event_timestamp,
|
||||
"max_event_timestamp": self.max_event_timestamp,
|
||||
"checksum_sha256": self.checksum,
|
||||
}
|
||||
|
||||
|
||||
class ManifestCounters:
|
||||
"""Счётчики стартовой истории для manifest."""
|
||||
|
||||
def __init__(self):
|
||||
self.topic_stats = {topic: _TopicManifestStats() for topic in TOPICS}
|
||||
self.click_ids: set[str] = set()
|
||||
self.user_ids: set[str] = set()
|
||||
|
||||
def add_batch(self, batch: dict[str, list[dict]]) -> None:
|
||||
browser_events = batch.get("browser_events", [])
|
||||
event_timestamps = {
|
||||
event["event_id"]: event.get("event_timestamp")
|
||||
for event in browser_events
|
||||
if event.get("event_id") and event.get("event_timestamp")
|
||||
}
|
||||
click_timestamps: dict[str, list[str]] = {}
|
||||
for event in browser_events:
|
||||
click_id = event.get("click_id")
|
||||
timestamp = event.get("event_timestamp")
|
||||
if click_id and timestamp:
|
||||
click_timestamps.setdefault(click_id, []).append(timestamp)
|
||||
|
||||
for topic, events in batch.items():
|
||||
if topic not in self.topic_stats:
|
||||
raise ValueError(f"unknown startup history topic: {topic}")
|
||||
stats = self.topic_stats[topic]
|
||||
for event in events:
|
||||
event_timestamp = event.get("event_timestamp")
|
||||
if event_timestamp is None and topic == "location_events":
|
||||
event_timestamp = event_timestamps.get(event.get("event_id"))
|
||||
if event_timestamp is None and topic in ("device_events", "geo_events"):
|
||||
timestamps = click_timestamps.get(event.get("click_id"))
|
||||
if timestamps:
|
||||
event_timestamp = min(timestamps)
|
||||
stats.add_timestamp(max(timestamps))
|
||||
stats.add(event, event_timestamp=event_timestamp)
|
||||
click_id = event.get("click_id")
|
||||
if topic == "browser_events" and click_id:
|
||||
self.click_ids.add(click_id)
|
||||
user_id = event.get("user_domain_id")
|
||||
if topic == "device_events" and user_id:
|
||||
self.user_ids.add(user_id)
|
||||
|
||||
def to_manifest_topics(self) -> dict:
|
||||
return {topic: stats.to_dict() for topic, stats in self.topic_stats.items()}
|
||||
|
||||
def to_manifest_totals(self) -> dict:
|
||||
browser_stats = self.topic_stats["browser_events"]
|
||||
return {
|
||||
"events": browser_stats.rows,
|
||||
"visits": len(self.click_ids),
|
||||
"users": len(self.user_ids),
|
||||
"min_event_timestamp": browser_stats.min_event_timestamp,
|
||||
"max_event_timestamp": browser_stats.max_event_timestamp,
|
||||
}
|
||||
|
||||
|
||||
class StartupHistoryArtifactBuilder:
|
||||
"""Собирает сообщения топиков для портативного файла."""
|
||||
|
||||
def __init__(self):
|
||||
self.topics = {topic: [] for topic in TOPICS}
|
||||
self.raw_topics = {topic: [] for topic in TOPICS}
|
||||
self.counters = ManifestCounters()
|
||||
|
||||
def add_batch(self, batch: dict[str, list[dict]]) -> None:
|
||||
self.counters.add_batch(batch)
|
||||
for topic in TOPICS:
|
||||
for event in batch.get(topic, []):
|
||||
self.topics[topic].append(deepcopy(event))
|
||||
self.raw_topics[topic].append(
|
||||
{
|
||||
"key": _event_key(event),
|
||||
"value_json": _event_value_json(event),
|
||||
}
|
||||
)
|
||||
|
||||
def to_artifact(self, manifest: dict, state: GeneratorState) -> dict:
|
||||
return {
|
||||
"artifact_version": ARTIFACT_VERSION,
|
||||
"kind": ARTIFACT_KIND,
|
||||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||||
"manifest": deepcopy(manifest),
|
||||
"state": state.to_dict(),
|
||||
"topics": deepcopy(self.topics),
|
||||
"raw_topics": deepcopy(self.raw_topics),
|
||||
}
|
||||
|
||||
|
||||
class KafkaRawPublisher:
|
||||
"""Публикует сохранённые Kafka value bytes без повторной JSON-сериализации."""
|
||||
|
||||
def __init__(self, bootstrap_servers: str):
|
||||
self.bootstrap_servers = bootstrap_servers
|
||||
KafkaProducerCls, _ = _kafka_importer()()
|
||||
self.producer = KafkaProducerCls(
|
||||
bootstrap_servers=self.bootstrap_servers,
|
||||
key_serializer=lambda k: k.encode("utf-8") if k else None,
|
||||
retries=3,
|
||||
retry_backoff_ms=1000,
|
||||
)
|
||||
|
||||
def publish_records(self, topic: str, records: list[dict]) -> tuple[int, int]:
|
||||
sent = 0
|
||||
errors = 0
|
||||
futures = []
|
||||
for record in records:
|
||||
future = self.producer.send(
|
||||
topic,
|
||||
key=record.get("key"),
|
||||
value=record["value_json"].encode("utf-8"),
|
||||
)
|
||||
futures.append(future)
|
||||
for future in futures:
|
||||
try:
|
||||
future.get(timeout=10)
|
||||
sent += 1
|
||||
except Exception:
|
||||
errors += 1
|
||||
return sent, errors
|
||||
|
||||
def flush(self) -> None:
|
||||
self.producer.flush()
|
||||
|
||||
def close(self) -> None:
|
||||
self.producer.close()
|
||||
|
||||
|
||||
class KafkaTopicInspector:
|
||||
"""Проверяет Kafka-топики перед clean-stand импортом."""
|
||||
|
||||
def __init__(self, bootstrap_servers: str):
|
||||
self.bootstrap_servers = bootstrap_servers
|
||||
|
||||
def assert_data_topics_empty(self) -> None:
|
||||
"""Падает, если в data-топиках уже есть сообщения."""
|
||||
from kafka import KafkaConsumer, TopicPartition
|
||||
|
||||
consumer = KafkaConsumer(
|
||||
bootstrap_servers=self.bootstrap_servers,
|
||||
enable_auto_commit=False,
|
||||
consumer_timeout_ms=1000,
|
||||
)
|
||||
try:
|
||||
dirty_topics = []
|
||||
for topic in TOPICS:
|
||||
partitions = consumer.partitions_for_topic(topic)
|
||||
if not partitions:
|
||||
continue
|
||||
topic_partitions = [
|
||||
TopicPartition(topic, partition)
|
||||
for partition in partitions
|
||||
]
|
||||
end_offsets = consumer.end_offsets(topic_partitions)
|
||||
if any(offset > 0 for offset in end_offsets.values()):
|
||||
dirty_topics.append(topic)
|
||||
if dirty_topics:
|
||||
raise RuntimeError(
|
||||
"Kafka data topics are not empty; clean the stand before "
|
||||
"startup history import: " + ", ".join(dirty_topics)
|
||||
)
|
||||
finally:
|
||||
consumer.close()
|
||||
|
||||
def snapshot_import_topics(self) -> dict:
|
||||
"""Снимок для отката clean-stand импорта."""
|
||||
return {"topics": list(IMPORT_TOPICS)}
|
||||
|
||||
def rollback_import_topics(self, snapshot: dict) -> None:
|
||||
"""Удаляет топики, которые могли получить частичный импорт."""
|
||||
from kafka import KafkaAdminClient
|
||||
from kafka.errors import UnknownTopicOrPartitionError
|
||||
|
||||
admin = KafkaAdminClient(bootstrap_servers=self.bootstrap_servers)
|
||||
try:
|
||||
existing_topics = set(admin.list_topics())
|
||||
topics = [
|
||||
topic
|
||||
for topic in snapshot.get("topics", IMPORT_TOPICS)
|
||||
if topic in existing_topics
|
||||
]
|
||||
if not topics:
|
||||
return
|
||||
try:
|
||||
admin.delete_topics(topics)
|
||||
except UnknownTopicOrPartitionError:
|
||||
return
|
||||
finally:
|
||||
admin.close()
|
||||
|
||||
|
||||
def generation_settings_from_config(config: Config) -> dict:
|
||||
"""Возвращает настройки, влияющие на поток генерации."""
|
||||
return {
|
||||
"tick_seconds": config.tick_seconds,
|
||||
"lambda_base_per_min": config.lambda_base_per_min,
|
||||
"jitter_pct": config.jitter_pct,
|
||||
"min_events_per_tick": config.min_events_per_tick,
|
||||
"max_events_per_tick": config.max_events_per_tick,
|
||||
"max_session_events": config.max_session_events,
|
||||
"max_active_sessions": config.max_active_sessions,
|
||||
"population_max": config.population_max,
|
||||
"p_new_user": config.p_new_user,
|
||||
"min_return_minutes": config.min_return_minutes,
|
||||
"model_time_speed": config.model_time_speed,
|
||||
}
|
||||
|
||||
|
||||
def build_manifest(config: Config, counters: ManifestCounters, state: GeneratorState) -> dict:
|
||||
"""Строит manifest стартовой истории."""
|
||||
if config.model_t_end is None:
|
||||
raise ValueError("GEN_MODEL_T_END is required for startup history manifest")
|
||||
return {
|
||||
"manifest_version": "1.0",
|
||||
"generated_at": datetime.now(timezone.utc).isoformat(),
|
||||
"gen_seed": config.seed,
|
||||
"model_t0": config.model_t0.isoformat(),
|
||||
"model_t_end": config.model_t_end.isoformat(),
|
||||
"model_timezone": config.model_timezone,
|
||||
"run_mode": "backfill",
|
||||
"generation_settings": generation_settings_from_config(config),
|
||||
"state_version": state.version,
|
||||
"state": {
|
||||
"topic": KafkaStateManager.STATE_TOPIC,
|
||||
"key": KafkaStateManager.STATE_KEY,
|
||||
"last_batch_id": state.last_batch_id,
|
||||
"model_timestamp": state.model_timestamp.isoformat(),
|
||||
},
|
||||
"topics": counters.to_manifest_topics(),
|
||||
"totals": counters.to_manifest_totals(),
|
||||
}
|
||||
|
||||
|
||||
def write_startup_history_artifact(path: str | Path, artifact: dict) -> None:
|
||||
"""Пишет артефакт в JSON-файл."""
|
||||
target = Path(path)
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_text(
|
||||
json.dumps(artifact, ensure_ascii=True, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
def load_startup_history_artifact(path: str | Path) -> dict:
|
||||
"""Читает артефакт из JSON-файла."""
|
||||
return json.loads(Path(path).read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
def validate_startup_history_artifact(
|
||||
artifact: dict,
|
||||
expected_config: Config | None = None,
|
||||
) -> tuple[GeneratorState, dict, dict[str, list[dict]], dict[str, list[dict]]]:
|
||||
"""Проверяет связку manifest, state и сообщений топиков."""
|
||||
if not isinstance(artifact, dict):
|
||||
raise ValueError("artifact must be an object")
|
||||
if artifact.get("kind") != ARTIFACT_KIND:
|
||||
raise ValueError("artifact kind mismatch")
|
||||
if artifact.get("artifact_version") != ARTIFACT_VERSION:
|
||||
raise ValueError("artifact version mismatch")
|
||||
|
||||
manifest = artifact.get("manifest")
|
||||
state_payload = artifact.get("state")
|
||||
topics = artifact.get("topics")
|
||||
raw_topics = artifact.get("raw_topics")
|
||||
if not isinstance(manifest, dict):
|
||||
raise ValueError("artifact manifest must be an object")
|
||||
if not isinstance(state_payload, dict):
|
||||
raise ValueError("artifact state must be an object")
|
||||
if not isinstance(topics, dict):
|
||||
raise ValueError("artifact topics must be an object")
|
||||
if not isinstance(raw_topics, dict):
|
||||
raise ValueError("artifact raw_topics must be an object")
|
||||
for topic in TOPICS:
|
||||
if not isinstance(topics.get(topic), list):
|
||||
raise ValueError(f"artifact topic {topic} must be a list")
|
||||
if not isinstance(raw_topics.get(topic), list):
|
||||
raise ValueError(f"artifact raw topic {topic} must be a list")
|
||||
|
||||
state = GeneratorState.from_dict(state_payload)
|
||||
_validate_manifest_state(manifest, state)
|
||||
_validate_manifest_topics(manifest, topics)
|
||||
_validate_raw_topics(topics, raw_topics)
|
||||
if expected_config is not None:
|
||||
_validate_manifest_config(manifest, expected_config)
|
||||
return (
|
||||
state,
|
||||
manifest,
|
||||
{topic: topics[topic] for topic in TOPICS},
|
||||
{topic: raw_topics[topic] for topic in TOPICS},
|
||||
)
|
||||
|
||||
|
||||
def import_startup_history_artifact(
|
||||
artifact: dict,
|
||||
*,
|
||||
publisher,
|
||||
state_manager,
|
||||
manifest_manager,
|
||||
expected_config: Config | None = None,
|
||||
topic_inspector=None,
|
||||
) -> dict:
|
||||
"""Воспроизводит артефакт в Kafka и compact-топиках генератора."""
|
||||
state, manifest, topics, raw_topics = validate_startup_history_artifact(
|
||||
artifact,
|
||||
expected_config=expected_config,
|
||||
)
|
||||
if topic_inspector is not None:
|
||||
topic_inspector.assert_data_topics_empty()
|
||||
snapshot = (
|
||||
topic_inspector.snapshot_import_topics()
|
||||
if hasattr(topic_inspector, "snapshot_import_topics")
|
||||
else None
|
||||
)
|
||||
else:
|
||||
snapshot = None
|
||||
|
||||
if not hasattr(publisher, "publish_records"):
|
||||
raise TypeError(
|
||||
"startup history import requires publisher.publish_records "
|
||||
"for raw Kafka values"
|
||||
)
|
||||
|
||||
try:
|
||||
sent_total = 0
|
||||
sent_by_topic = {}
|
||||
|
||||
for topic in TOPICS:
|
||||
raw_records = raw_topics[topic]
|
||||
if not raw_records:
|
||||
sent_by_topic[topic] = 0
|
||||
continue
|
||||
sent, errors = publisher.publish_records(topic, raw_records)
|
||||
if errors:
|
||||
raise RuntimeError(
|
||||
f"Startup history import failed for {topic}: sent={sent}, errors={errors}"
|
||||
)
|
||||
sent_total += sent
|
||||
sent_by_topic[topic] = sent
|
||||
|
||||
publisher.flush()
|
||||
manifest_manager.save(manifest)
|
||||
manifest_manager.flush()
|
||||
state_manager.save(state)
|
||||
state_manager.flush()
|
||||
except Exception:
|
||||
if (
|
||||
topic_inspector is not None
|
||||
and snapshot is not None
|
||||
and hasattr(topic_inspector, "rollback_import_topics")
|
||||
):
|
||||
topic_inspector.rollback_import_topics(snapshot)
|
||||
raise
|
||||
|
||||
return {"events": sent_total, "topics": sent_by_topic}
|
||||
|
||||
|
||||
def compare_clickhouse_stats_to_manifest(
|
||||
manifest: dict,
|
||||
stats: dict[str, str],
|
||||
) -> list[str]:
|
||||
"""Возвращает поля ClickHouse, которые не совпали с manifest."""
|
||||
totals = manifest.get("totals") or {}
|
||||
expected = {
|
||||
"events": str(totals.get("events")),
|
||||
"visits": str(totals.get("visits")),
|
||||
"users": str(totals.get("users")),
|
||||
"min_event_timestamp": str(totals.get("min_event_timestamp")),
|
||||
"max_event_timestamp": str(totals.get("max_event_timestamp")),
|
||||
}
|
||||
return [
|
||||
field
|
||||
for field, expected_value in expected.items()
|
||||
if str(stats.get(field)) != expected_value
|
||||
]
|
||||
|
||||
|
||||
def _validate_manifest_state(manifest: dict, state: GeneratorState) -> None:
|
||||
mismatches = []
|
||||
expected_state = manifest.get("state") or {}
|
||||
if manifest.get("run_mode") != "backfill":
|
||||
mismatches.append("run_mode")
|
||||
if manifest.get("state_version") != state.version:
|
||||
mismatches.append("state_version")
|
||||
if expected_state.get("last_batch_id") != state.last_batch_id:
|
||||
mismatches.append("state.last_batch_id")
|
||||
if expected_state.get("model_timestamp") != state.model_timestamp.isoformat():
|
||||
mismatches.append("state.model_timestamp")
|
||||
try:
|
||||
model_t0 = _parse_manifest_timestamp(manifest["model_t0"])
|
||||
model_t_end = _parse_manifest_timestamp(manifest["model_t_end"])
|
||||
except (KeyError, TypeError, ValueError) as exc:
|
||||
raise ValueError(f"manifest time fields are invalid: {exc}") from exc
|
||||
if state.model_t0 != model_t0:
|
||||
mismatches.append("GEN_MODEL_T0")
|
||||
if state.model_timestamp != model_t_end:
|
||||
mismatches.append("GEN_MODEL_T_END")
|
||||
if state.gen_seed != manifest.get("gen_seed"):
|
||||
mismatches.append("GEN_SEED")
|
||||
if state.model_timezone != manifest.get("model_timezone"):
|
||||
mismatches.append("GEN_MODEL_TIMEZONE")
|
||||
settings = manifest.get("generation_settings") or {}
|
||||
if abs(state.model_time_speed - settings.get("model_time_speed", -1)) > 1e-9:
|
||||
mismatches.append("GEN_MODEL_TIME_SPEED")
|
||||
if mismatches:
|
||||
raise ValueError("startup history artifact mismatch: " + ", ".join(mismatches))
|
||||
|
||||
|
||||
def _validate_manifest_topics(manifest: dict, topics: dict[str, list[dict]]) -> None:
|
||||
counters = ManifestCounters()
|
||||
counters.add_batch({topic: topics[topic] for topic in TOPICS})
|
||||
expected_topics = manifest.get("topics") or {}
|
||||
expected_totals = manifest.get("totals") or {}
|
||||
if counters.to_manifest_topics() != expected_topics:
|
||||
raise ValueError("startup history artifact mismatch: topic checksums")
|
||||
if counters.to_manifest_totals() != expected_totals:
|
||||
raise ValueError("startup history artifact mismatch: totals")
|
||||
|
||||
|
||||
def _validate_raw_topics(
|
||||
topics: dict[str, list[dict]],
|
||||
raw_topics: dict[str, list[dict]],
|
||||
) -> None:
|
||||
for topic in TOPICS:
|
||||
if len(raw_topics[topic]) != len(topics[topic]):
|
||||
raise ValueError(f"artifact raw topic {topic} row count mismatch")
|
||||
for event, record in zip(topics[topic], raw_topics[topic]):
|
||||
if record.get("key") != _event_key(event):
|
||||
raise ValueError(f"artifact raw topic {topic} key mismatch")
|
||||
if _raw_value_to_event(record.get("value_json"), topic) != event:
|
||||
raise ValueError(f"artifact raw topic {topic} value mismatch")
|
||||
|
||||
|
||||
def _event_key(event: dict) -> str | None:
|
||||
return event.get("event_id") or event.get("click_id")
|
||||
|
||||
|
||||
def _event_value_json(event: dict) -> str:
|
||||
return kafka_event_value_json(event)
|
||||
|
||||
|
||||
def _raw_value_to_event(value: Any, topic: str) -> dict:
|
||||
if not isinstance(value, str):
|
||||
raise ValueError(f"artifact raw topic {topic} value must be a string")
|
||||
try:
|
||||
event = json.loads(value)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError(f"artifact raw topic {topic} value must be valid JSON") from exc
|
||||
if not isinstance(event, dict):
|
||||
raise ValueError(f"artifact raw topic {topic} value must be a JSON object")
|
||||
return event
|
||||
|
||||
|
||||
def _validate_manifest_config(manifest: dict, config: Config) -> None:
|
||||
mismatches = []
|
||||
if manifest.get("gen_seed") != config.seed:
|
||||
mismatches.append("GEN_SEED")
|
||||
if _parse_manifest_timestamp(manifest.get("model_t0")) != config.model_t0:
|
||||
mismatches.append("GEN_MODEL_T0")
|
||||
if config.model_t_end is None:
|
||||
mismatches.append("GEN_MODEL_T_END")
|
||||
elif _parse_manifest_timestamp(manifest.get("model_t_end")) != config.model_t_end:
|
||||
mismatches.append("GEN_MODEL_T_END")
|
||||
if manifest.get("model_timezone") != config.model_timezone:
|
||||
mismatches.append("GEN_MODEL_TIMEZONE")
|
||||
if manifest.get("generation_settings") != generation_settings_from_config(config):
|
||||
mismatches.append("generation_settings")
|
||||
if mismatches:
|
||||
raise ValueError(
|
||||
"startup history artifact is incompatible with current settings: "
|
||||
+ ", ".join(mismatches)
|
||||
)
|
||||
|
||||
|
||||
def _parse_manifest_timestamp(value: Any) -> datetime:
|
||||
if not isinstance(value, str):
|
||||
raise ValueError("timestamp must be a string")
|
||||
timestamp = datetime.fromisoformat(value.replace("Z", "+00:00"))
|
||||
if timestamp.tzinfo is None:
|
||||
raise ValueError("timestamp must include timezone")
|
||||
return timestamp.astimezone(timezone.utc)
|
||||
|
||||
|
||||
def _run_import(args: argparse.Namespace) -> int:
|
||||
config = Config()
|
||||
artifact = load_startup_history_artifact(args.artifact)
|
||||
ensure_topics(config.kafka_bootstrap_servers)
|
||||
publisher = KafkaRawPublisher(config.kafka_bootstrap_servers)
|
||||
state_manager = KafkaStateManager(config.kafka_bootstrap_servers)
|
||||
manifest_manager = KafkaStartupHistoryManifest(config.kafka_bootstrap_servers)
|
||||
topic_inspector = KafkaTopicInspector(config.kafka_bootstrap_servers)
|
||||
try:
|
||||
result = import_startup_history_artifact(
|
||||
artifact,
|
||||
publisher=publisher,
|
||||
state_manager=state_manager,
|
||||
manifest_manager=manifest_manager,
|
||||
expected_config=config,
|
||||
topic_inspector=topic_inspector,
|
||||
)
|
||||
finally:
|
||||
publisher.close()
|
||||
state_manager.close()
|
||||
manifest_manager.close()
|
||||
logger.info(
|
||||
"Startup history artifact imported: events=%s, topics=%s",
|
||||
result["events"],
|
||||
result["topics"],
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
def _run_validate(args: argparse.Namespace) -> int:
|
||||
config = Config()
|
||||
artifact = load_startup_history_artifact(args.artifact)
|
||||
validate_startup_history_artifact(artifact, expected_config=config)
|
||||
logger.info("Startup history artifact is valid: %s", args.artifact)
|
||||
return 0
|
||||
|
||||
|
||||
def _run_manifest_summary(args: argparse.Namespace) -> int:
|
||||
artifact = load_startup_history_artifact(args.artifact)
|
||||
_, manifest, _, _ = validate_startup_history_artifact(artifact)
|
||||
totals = manifest["totals"]
|
||||
print(
|
||||
"\t".join(
|
||||
[
|
||||
str(totals["events"]),
|
||||
str(totals["visits"]),
|
||||
str(totals["users"]),
|
||||
str(totals["min_event_timestamp"]),
|
||||
str(totals["max_event_timestamp"]),
|
||||
str(manifest["model_t0"]),
|
||||
str(manifest["model_t_end"]),
|
||||
]
|
||||
)
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(description="Startup history artifact tools")
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
validate_parser = subparsers.add_parser("validate")
|
||||
validate_parser.add_argument("--artifact", required=True)
|
||||
validate_parser.set_defaults(func=_run_validate)
|
||||
|
||||
import_parser = subparsers.add_parser("import")
|
||||
import_parser.add_argument("--artifact", required=True)
|
||||
import_parser.set_defaults(func=_run_import)
|
||||
|
||||
summary_parser = subparsers.add_parser("manifest-summary")
|
||||
summary_parser.add_argument("--artifact", required=True)
|
||||
summary_parser.set_defaults(func=_run_manifest_summary)
|
||||
|
||||
args = parser.parse_args(argv)
|
||||
return args.func(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user