311 lines
9.9 KiB
Python
311 lines
9.9 KiB
Python
# ---
|
|
# jupyter:
|
|
# jupytext:
|
|
# formats: py:percent
|
|
# text_representation:
|
|
# extension: .py
|
|
# format_name: percent
|
|
# kernelspec:
|
|
# display_name: Python 3
|
|
# language: python
|
|
# name: python3
|
|
# ---
|
|
|
|
# %% [markdown]
|
|
# # Training smoke test (`input-sample.md`)
|
|
#
|
|
# Run cells top to bottom in VS Code / Cursor (**Run Cell** on each `# %%` block).
|
|
#
|
|
# Steps mirror `input-sample.md`: optional DB delete + insert, `mc cp` to MinIO, Temporal `train_model`.
|
|
# Set `POSTGRES_*`, `TEMPORAL_*`, and configure the `mc` alias. Add/remove entries in
|
|
# `TRAINING_TEST_INPUTS` below to choose which experiments run.
|
|
|
|
# %%
|
|
from __future__ import annotations
|
|
|
|
import csv
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import uuid
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import psycopg2
|
|
from dotenv import load_dotenv
|
|
from psycopg2.extras import Json
|
|
from temporalio import client
|
|
|
|
TRAINING_TEST_INPUTS: list[str] = [
|
|
'scripts/inputs/linear_regression.json',
|
|
'scripts/inputs/xgboost.json',
|
|
]
|
|
|
|
TRAINING_TEST_INPUTS: list[str] = [
|
|
'scripts/inputs/sin-approx.json',
|
|
]
|
|
|
|
_REQUIRED_INPUT_KEYS = ('experiment', 'minio', 'temporal', 'payload', 'db_only')
|
|
|
|
|
|
def _load_input(path: Path) -> dict[str, Any]:
|
|
"""
|
|
Load a training smoke-test input JSON and validate required top-level keys.
|
|
|
|
Args:
|
|
- path: Path to the input JSON file
|
|
|
|
Return:
|
|
Parsed input dict with experiment, minio, temporal, payload, and db_only sections
|
|
"""
|
|
data: dict[str, Any] = json.loads(path.read_text(encoding='utf-8'))
|
|
missing = [key for key in _REQUIRED_INPUT_KEYS if key not in data]
|
|
if missing:
|
|
raise KeyError(f'Input JSON missing required top-level keys: {", ".join(missing)}')
|
|
return data
|
|
|
|
|
|
def _postgres_connect_kwargs() -> dict[str, str | int]:
|
|
"""
|
|
Build psycopg2.connect keyword arguments from POSTGRES_* environment variables.
|
|
|
|
Return:
|
|
host, port, user, password, and dbname suitable for psycopg2.connect
|
|
"""
|
|
host = os.getenv('POSTGRES_HOST')
|
|
user = os.getenv('POSTGRES_USER')
|
|
password = os.getenv('POSTGRES_PASSWORD')
|
|
dbname = os.getenv('POSTGRES_DBNAME')
|
|
if not host or not user or not password or not dbname:
|
|
raise RuntimeError(
|
|
'Set POSTGRES_HOST, POSTGRES_USER, POSTGRES_PASSWORD, and POSTGRES_DBNAME'
|
|
)
|
|
return {
|
|
'host': host,
|
|
'port': int(os.getenv('POSTGRES_PORT', '5432')),
|
|
'user': user,
|
|
'password': password,
|
|
'dbname': dbname,
|
|
}
|
|
|
|
|
|
def _resolve_input_paths(project_root: Path, entries: list[str]) -> list[Path]:
|
|
"""
|
|
Resolve and validate the hardcoded TRAINING_TEST_INPUTS list against the filesystem.
|
|
|
|
Args:
|
|
- project_root: Repository root used to resolve relative entries
|
|
- entries: Input file paths (absolute or relative to project_root)
|
|
|
|
Return:
|
|
Absolute paths to the input JSON files, in declaration order
|
|
"""
|
|
if not entries:
|
|
raise RuntimeError('TRAINING_TEST_INPUTS is empty; add at least one input JSON path')
|
|
resolved: list[Path] = []
|
|
for entry in entries:
|
|
candidate = Path(entry)
|
|
if not candidate.is_absolute():
|
|
candidate = project_root / candidate
|
|
if not candidate.is_file():
|
|
raise FileNotFoundError(f'Training test input file not found: {candidate}')
|
|
resolved.append(candidate.resolve())
|
|
return resolved
|
|
|
|
|
|
def _validate_csv_columns(
|
|
local_csv: Path,
|
|
payload: dict[str, Any],
|
|
input_path: Path,
|
|
) -> None:
|
|
"""
|
|
Fail fast when the local CSV header does not match payload column names.
|
|
|
|
Args:
|
|
- local_csv: Resolved path to the CSV on disk
|
|
- payload: Training payload with variable_columns, target_variable, and separators
|
|
- input_path: Input JSON path (for error messages)
|
|
|
|
Return:
|
|
None; raises ValueError when required columns are missing from the CSV header
|
|
"""
|
|
line_separator = str(payload.get('line_separator', ','))
|
|
with local_csv.open(newline='', encoding='utf-8') as handle:
|
|
header = next(csv.reader(handle, delimiter=line_separator))
|
|
required = list(payload['variable_columns']) + [str(payload['target_variable'])]
|
|
date_column = payload.get('date_column')
|
|
if date_column:
|
|
required.append(str(date_column))
|
|
missing = [column for column in required if column not in header]
|
|
if missing:
|
|
raise ValueError(
|
|
f'{input_path}: CSV {local_csv} header {header!r} is missing columns '
|
|
f'{missing!r} (check payload variable_columns, target_variable, date_column)'
|
|
)
|
|
|
|
|
|
def _resolve_local_csv(project_root: Path, local_csv: str) -> Path:
|
|
"""
|
|
Resolve a minio.local_csv entry against the project root when relative.
|
|
|
|
Args:
|
|
- project_root: Repository root
|
|
- local_csv: Path declared in the input JSON (absolute or relative)
|
|
|
|
Return:
|
|
Absolute path to the local CSV file
|
|
"""
|
|
candidate = Path(local_csv)
|
|
if not candidate.is_absolute():
|
|
candidate = project_root / candidate
|
|
return candidate
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
|
load_dotenv(PROJECT_ROOT / '.env')
|
|
|
|
experiments: list[dict[str, Any]] = []
|
|
for _input_path in _resolve_input_paths(PROJECT_ROOT, TRAINING_TEST_INPUTS):
|
|
_input = _load_input(_input_path)
|
|
_local_csv = _resolve_local_csv(PROJECT_ROOT, _input['minio']['local_csv'])
|
|
_validate_csv_columns(_local_csv, _input['payload'], _input_path)
|
|
experiments.append({'path': _input_path, **_input})
|
|
_payload = _input['payload']
|
|
_experiment = _input['experiment']
|
|
print(
|
|
f'input={_input_path} '
|
|
f'model_type={_payload["model_type"]} '
|
|
f'model_name={_payload["model_name"]} '
|
|
f'experiment_run_id={_experiment["experiment_run_id"]}'
|
|
)
|
|
|
|
# %%
|
|
# --- configuration (`.env` at repo root; per-run values from input JSON) ---
|
|
|
|
PG = _postgres_connect_kwargs()
|
|
|
|
TEMPORAL_HOST = os.getenv('TEMPORAL_HOST')
|
|
TEMPORAL_NAMESPACE = os.getenv('TEMPORAL_NAMESPACE')
|
|
TEMPORAL_TLS = os.getenv('TEMPORAL_USE_TLS', 'false').lower() in ('1', 'true', 'yes')
|
|
|
|
print(PG)
|
|
print(TEMPORAL_HOST, TEMPORAL_NAMESPACE, TEMPORAL_TLS)
|
|
for exp in experiments:
|
|
exp_minio = exp['minio']
|
|
exp_temporal = exp['temporal']
|
|
print(
|
|
exp_minio['mc_alias'],
|
|
exp_minio['bucket_name'],
|
|
exp_minio['file_name'],
|
|
_resolve_local_csv(PROJECT_ROOT, exp_minio['local_csv']),
|
|
exp_temporal['task_queue'],
|
|
)
|
|
|
|
# %%
|
|
# --- 1) database: delete previous row (same id), then insert `experiment_run` for each experiment ---
|
|
# Primary key column is `id` (see `experiment_tracking` updates).
|
|
|
|
now = datetime.utcnow()
|
|
|
|
with psycopg2.connect(
|
|
host=PG['host'],
|
|
port=PG['port'],
|
|
user=PG['user'],
|
|
password=PG['password'],
|
|
dbname=PG['dbname'],
|
|
) as conn:
|
|
with conn.cursor() as cur:
|
|
for exp in experiments:
|
|
exp_experiment = exp['experiment']
|
|
exp_minio = exp['minio']
|
|
exp_payload = exp['payload']
|
|
exp_db_only = exp['db_only']
|
|
request_data = {
|
|
**exp_payload,
|
|
'bucket_name': exp_minio['bucket_name'],
|
|
'file_name': exp_minio['file_name'],
|
|
'model_metadata': exp_db_only['model_metadata'],
|
|
}
|
|
cur.execute(
|
|
'DELETE FROM public.experiment_run WHERE id = %s',
|
|
(exp_experiment['experiment_run_id'],),
|
|
)
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO public.experiment_run (
|
|
id, experiment_name, run_name, username, status, error_message,
|
|
created_at, updated_at, bucket_name, file_name, request_data, orchestrator_response_data
|
|
)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
""",
|
|
(
|
|
exp_experiment['experiment_run_id'],
|
|
exp_experiment['experiment_name'],
|
|
exp_experiment['run_name'],
|
|
exp_experiment['username'],
|
|
exp_experiment['status'],
|
|
None,
|
|
now,
|
|
now,
|
|
exp_minio['bucket_name'],
|
|
exp_minio['file_name'],
|
|
Json(request_data),
|
|
None,
|
|
),
|
|
)
|
|
|
|
# %%
|
|
# --- 2) MinIO: upload local CSV for each experiment (requires `mc` CLI and alias configured) ---
|
|
for exp in experiments:
|
|
exp_minio = exp['minio']
|
|
local_csv = _resolve_local_csv(PROJECT_ROOT, exp_minio['local_csv'])
|
|
subprocess.run(
|
|
[
|
|
'mc',
|
|
'cp',
|
|
'--insecure',
|
|
str(local_csv),
|
|
f'{exp_minio["mc_alias"]}/{exp_minio["bucket_name"]}/{exp_minio["file_name"]}',
|
|
],
|
|
check=True,
|
|
)
|
|
|
|
# %%
|
|
# --- 3) Temporal: connect once, then start `train_model` per experiment ---
|
|
|
|
if not TEMPORAL_HOST or not TEMPORAL_NAMESPACE:
|
|
raise RuntimeError('Set TEMPORAL_HOST and TEMPORAL_NAMESPACE in the environment')
|
|
|
|
c = await client.Client.connect( # type: ignore[top-level-await]
|
|
target_host=TEMPORAL_HOST,
|
|
namespace=TEMPORAL_NAMESPACE,
|
|
tls=TEMPORAL_TLS,
|
|
)
|
|
|
|
# %%
|
|
|
|
for exp in experiments:
|
|
exp_minio = exp['minio']
|
|
exp_temporal = exp['temporal']
|
|
exp_payload = exp['payload']
|
|
workflow_input = {
|
|
**exp_payload,
|
|
'bucket_name': exp_minio['bucket_name'],
|
|
'file_name': exp_minio['file_name'],
|
|
}
|
|
wid = f'train-model-test-{uuid.uuid4()}'
|
|
result = await c.execute_workflow( # type: ignore[top-level-await, call-overload]
|
|
exp_temporal['workflow_name'],
|
|
workflow_input,
|
|
id=wid,
|
|
task_queue=exp_temporal['task_queue'],
|
|
execution_timeout=timedelta(minutes=exp_temporal['execution_timeout_minutes']),
|
|
run_timeout=timedelta(minutes=exp_temporal['run_timeout_minutes']),
|
|
task_timeout=timedelta(minutes=exp_temporal['task_timeout_minutes']),
|
|
)
|
|
print(wid)
|
|
print(result)
|
|
|
|
# %%
|