SIENTIAPDE-1645: code snapshot (part 1)
This commit is contained in:
vitor-aignosi
310
scripts/run_training_test.py
Normal file
310
scripts/run_training_test.py
Normal file
@@ -0,0 +1,310 @@
|
||||
# ---
|
||||
# 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)
|
||||
|
||||
# %%
|
||||
Reference in New Issue
Block a user