feat(airflow): добавлен пульт управления генератором
- Зачем: - нужен основной ручной интерфейс стенда для backfill/import/check без консольной матрицы переменных. - Что: - добавлен DAG generator_control с параметрами Airflow, ветвлением операций и ожиданием ETL. - вынесена общая логика запуска и предпроверок генератора для Airflow. - обновлены compose-настройки, зависимости, тесты и документация по пульту. - Проверка: - uv run --with pytest --with-requirements generator/requirements.txt pytest generator/tests -q. - docker compose config --quiet.
This commit is contained in:
@@ -0,0 +1,257 @@
|
||||
"""
|
||||
DAG-пульт генератора стартовой истории.
|
||||
|
||||
Пульт выполняет только Python-код генератора внутри Airflow worker. Жизненный
|
||||
цикл контейнеров остаётся в Makefile.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
from airflow import DAG
|
||||
from airflow.models.param import Param
|
||||
from airflow.operators.empty import EmptyOperator
|
||||
from airflow.operators.python import BranchPythonOperator, PythonOperator
|
||||
from airflow.operators.trigger_dagrun import TriggerDagRunOperator
|
||||
from airflow.utils.trigger_rule import TriggerRule
|
||||
from airflow_clickhouse_plugin.hooks.clickhouse import ClickHouseHook
|
||||
|
||||
from clickstream_generator.airflow_control import (
|
||||
KAFKA_BOOTSTRAP_SERVERS,
|
||||
assert_clickhouse_matches_manifest,
|
||||
assert_live_generator_not_running,
|
||||
assert_stand_clean,
|
||||
build_control_env,
|
||||
default_artifact_path,
|
||||
load_manifest_from_kafka,
|
||||
run_backfill,
|
||||
run_import,
|
||||
validate_import_artifact,
|
||||
)
|
||||
from clickstream_generator.launch import PROFILES
|
||||
|
||||
|
||||
default_args = {
|
||||
"owner": "airflow",
|
||||
"depends_on_past": False,
|
||||
"email_on_failure": False,
|
||||
"email_on_retry": False,
|
||||
"retries": 0,
|
||||
"retry_delay": timedelta(minutes=1),
|
||||
}
|
||||
|
||||
|
||||
def _conf(context) -> dict:
|
||||
dag_run = context.get("dag_run")
|
||||
return dag_run.conf if dag_run and dag_run.conf else {}
|
||||
|
||||
|
||||
def _param(context, name: str):
|
||||
return _conf(context).get(name, context["params"][name])
|
||||
|
||||
|
||||
def _operation(context) -> str:
|
||||
return str(_param(context, "operation"))
|
||||
|
||||
|
||||
def _artifact_path(context) -> str:
|
||||
value = str(_param(context, "artifact_path") or "").strip()
|
||||
if value:
|
||||
return value
|
||||
return default_artifact_path(_operation(context))
|
||||
|
||||
|
||||
def _overrides(context) -> dict[str, str]:
|
||||
return {
|
||||
key: str(_param(context, param_name) or "").strip()
|
||||
for param_name, key in {
|
||||
"seed": "GEN_SEED",
|
||||
"model_time_speed": "GEN_MODEL_TIME_SPEED",
|
||||
}.items()
|
||||
}
|
||||
|
||||
|
||||
def _build_env(context, operation: str) -> dict[str, str]:
|
||||
artifact_path = _artifact_path(context)
|
||||
if (
|
||||
operation == "backfill"
|
||||
and not str(_param(context, "artifact_path") or "").strip()
|
||||
):
|
||||
artifact_path = None
|
||||
return build_control_env(
|
||||
operation,
|
||||
profile_name=str(_param(context, "profile")),
|
||||
duration=str(_param(context, "duration") or "").strip(),
|
||||
artifact_path=artifact_path,
|
||||
overrides=_overrides(context),
|
||||
)
|
||||
|
||||
|
||||
def choose_operation(**context) -> str:
|
||||
"""Выбирает ветку пульта по параметру operation."""
|
||||
operation = _operation(context)
|
||||
if operation == "backfill":
|
||||
return "precheck_backfill"
|
||||
if operation == "import":
|
||||
return "precheck_import"
|
||||
if operation == "check":
|
||||
return "check_only"
|
||||
raise ValueError(f"Неизвестная операция: {operation}")
|
||||
|
||||
|
||||
def precheck_backfill(**context) -> None:
|
||||
"""Проверяет чистоту стенда перед backfill."""
|
||||
assert_live_generator_not_running()
|
||||
hook = ClickHouseHook(clickhouse_conn_id="clickhouse_default", database="default")
|
||||
assert_stand_clean(KAFKA_BOOTSTRAP_SERVERS, hook)
|
||||
context["ti"].xcom_push(
|
||||
key="generator_env",
|
||||
value=_build_env(context, "backfill"),
|
||||
)
|
||||
|
||||
|
||||
def run_backfill_task(**context) -> None:
|
||||
"""Выполняет backfill через код генератора."""
|
||||
env = context["ti"].xcom_pull(task_ids="precheck_backfill", key="generator_env")
|
||||
run_backfill(env)
|
||||
|
||||
|
||||
def precheck_import(**context) -> None:
|
||||
"""Проверяет чистоту стенда и совместимость артефакта перед import."""
|
||||
assert_live_generator_not_running()
|
||||
hook = ClickHouseHook(clickhouse_conn_id="clickhouse_default", database="default")
|
||||
assert_stand_clean(KAFKA_BOOTSTRAP_SERVERS, hook)
|
||||
env = _build_env(context, "import")
|
||||
artifact_path = _artifact_path(context)
|
||||
validate_import_artifact(env, artifact_path)
|
||||
context["ti"].xcom_push(key="generator_env", value=env)
|
||||
context["ti"].xcom_push(key="artifact_path", value=artifact_path)
|
||||
|
||||
|
||||
def run_import_task(**context) -> None:
|
||||
"""Выполняет import портативного артефакта."""
|
||||
ti = context["ti"]
|
||||
env = ti.xcom_pull(task_ids="precheck_import", key="generator_env")
|
||||
artifact_path = ti.xcom_pull(task_ids="precheck_import", key="artifact_path")
|
||||
run_import(env, artifact_path)
|
||||
|
||||
|
||||
def check_manifest_task(**context) -> None:
|
||||
"""Сверяет витрину ClickHouse с manifest стартовой истории."""
|
||||
manifest = load_manifest_from_kafka()
|
||||
hook = ClickHouseHook(clickhouse_conn_id="clickhouse_default", database="default")
|
||||
assert_clickhouse_matches_manifest(manifest, hook)
|
||||
|
||||
|
||||
with DAG(
|
||||
dag_id="generator_control",
|
||||
description="Пульт стартовой истории генератора",
|
||||
default_args=default_args,
|
||||
schedule=None,
|
||||
start_date=datetime(2024, 1, 1),
|
||||
catchup=False,
|
||||
max_active_runs=1,
|
||||
is_paused_upon_creation=True,
|
||||
tags=["generator", "startup-history"],
|
||||
params={
|
||||
"operation": Param(
|
||||
"backfill",
|
||||
type="string",
|
||||
enum=["backfill", "import", "check"],
|
||||
title="Операция",
|
||||
description=(
|
||||
"Что сделать: создать историю, импортировать артефакт "
|
||||
"или проверить витрины."
|
||||
),
|
||||
),
|
||||
"profile": Param(
|
||||
sorted(PROFILES)[0],
|
||||
type="string",
|
||||
enum=sorted(PROFILES),
|
||||
title="Профиль",
|
||||
description="Именованный набор настроек генератора.",
|
||||
),
|
||||
"duration": Param(
|
||||
"",
|
||||
type="string",
|
||||
title="Длительность",
|
||||
description="Например 6h или 2d. Пусто — взять длительность из профиля.",
|
||||
),
|
||||
"seed": Param(
|
||||
"",
|
||||
type="string",
|
||||
title="GEN_SEED",
|
||||
description="Пусто — взять seed из профиля.",
|
||||
),
|
||||
"model_time_speed": Param(
|
||||
"",
|
||||
type="string",
|
||||
title="GEN_MODEL_TIME_SPEED",
|
||||
description="Пусто — взять скорость модельного времени из профиля.",
|
||||
),
|
||||
"artifact_path": Param(
|
||||
"",
|
||||
type="string",
|
||||
title="Артефакт",
|
||||
description=(
|
||||
"Backfill: куда сохранить файл; пусто — не сохранять. "
|
||||
"Import: что читать; пусто — путь по умолчанию в data."
|
||||
),
|
||||
),
|
||||
},
|
||||
) as dag:
|
||||
route = BranchPythonOperator(
|
||||
task_id="choose_operation",
|
||||
python_callable=choose_operation,
|
||||
)
|
||||
|
||||
precheck_backfill_task = PythonOperator(
|
||||
task_id="precheck_backfill",
|
||||
python_callable=precheck_backfill,
|
||||
)
|
||||
backfill_task = PythonOperator(
|
||||
task_id="run_backfill",
|
||||
python_callable=run_backfill_task,
|
||||
)
|
||||
|
||||
precheck_import_task = PythonOperator(
|
||||
task_id="precheck_import",
|
||||
python_callable=precheck_import,
|
||||
)
|
||||
import_task = PythonOperator(
|
||||
task_id="run_import",
|
||||
python_callable=run_import_task,
|
||||
)
|
||||
|
||||
trigger_etl = TriggerDagRunOperator(
|
||||
task_id="trigger_etl",
|
||||
trigger_dag_id="etl_pipeline",
|
||||
conf={"full_refresh": True},
|
||||
wait_for_completion=True,
|
||||
allowed_states=["success"],
|
||||
failed_states=["failed"],
|
||||
poke_interval=30,
|
||||
trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS,
|
||||
)
|
||||
|
||||
check_after_etl = PythonOperator(
|
||||
task_id="check_after_etl",
|
||||
python_callable=check_manifest_task,
|
||||
)
|
||||
|
||||
check_only = PythonOperator(
|
||||
task_id="check_only",
|
||||
python_callable=check_manifest_task,
|
||||
)
|
||||
|
||||
done = EmptyOperator(
|
||||
task_id="done",
|
||||
trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS,
|
||||
)
|
||||
|
||||
route >> [precheck_backfill_task, precheck_import_task, check_only]
|
||||
precheck_backfill_task >> backfill_task >> trigger_etl
|
||||
precheck_import_task >> import_task >> trigger_etl
|
||||
trigger_etl >> check_after_etl >> done
|
||||
check_only >> done
|
||||
Reference in New Issue
Block a user