dfbdd4b605
The probes only check app_up, which was set to 1 when the metrics server started, before install_runtime. A pod that would fail the guard could turn Ready first, so a rollout would remove the old pod. app_up now goes to 1 right after the guard. Also: copy sys.modules keys before iterating, raise the sientia_model floor to the runtime's >=0.14.0, pin the same runtime versions in requirements-local.txt, test that the snapshot is taken before the install and the real importlib.metadata path, and document the boot failure in README and CLAUDE.md.
79 lines
8.2 KiB
Markdown
79 lines
8.2 KiB
Markdown
# CLAUDE.md
|
|
|
|
This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
|
|
|
|
## What this is
|
|
|
|
Enterprise ML model training orchestration platform built on **Temporal**. Orchestrates training workflows (validation → train → save to MLFlow → cleanup), backed by PostgreSQL (experiment tracking), MinIO (data/artifacts), MongoDB (notifications), and a Git-based "Plugin Store" that supplies model runtimes/wrappers dynamically. Python 3.11+.
|
|
|
|
## Commands
|
|
|
|
```bash
|
|
# Local run (loads .env, starts worker)
|
|
./run_local.sh
|
|
# or manually:
|
|
python -m model_manager.worker.worker
|
|
|
|
# Full validation (format check, lint, mypy, bandit, tests w/ coverage) — run before committing
|
|
./validate.sh
|
|
./validate.sh --fix # auto-fix format/lint, no tests
|
|
./validate.sh --no-tests # skip tests
|
|
./validate.sh --only-tests # tests only
|
|
|
|
# Individual tools
|
|
ruff format --check model_manager/ tests/
|
|
ruff format model_manager/ tests/ # auto-fix
|
|
ruff check model_manager/ tests/
|
|
ruff check --fix model_manager/ tests/
|
|
mypy model_manager/
|
|
bandit -r model_manager/ -ll
|
|
|
|
# Tests
|
|
pytest # runs tests/ and e2e/ (see pyproject testpaths)
|
|
pytest tests/ # unit tests only
|
|
pytest tests/activities/test_training.py -v # single file
|
|
pytest tests/activities/test_training.py::test_name -v # single test
|
|
pytest tests/ --cov=model_manager --cov-report=term-missing --cov-report=html
|
|
|
|
# Manual integration/E2E scripts (require live Temporal/Postgres/MinIO stack; run via IDE "Run Cell" or directly)
|
|
python scripts/run_training_test.py # supports --all for batch scenario execution
|
|
python scripts/run_cleanup_test.py
|
|
```
|
|
|
|
Coverage gate in CI is 80% minimum; this repo targets 100% in practice — don't add uncovered code paths without tests.
|
|
|
|
## Architecture
|
|
|
|
### Temporal layering
|
|
- **Worker** (`model_manager/worker/worker.py`): entrypoint. Builds config from env, connects Postgres/MLFlow/MinIO/Plugin Store, installs the `RUNTIME` model runtime via `PluginStore.install_runtime` (boot fails if that changes the version of an already-imported package, so the `requirements.txt` pins must match the factory runtimes; `app_up` goes to 1 only after this check), starts two Temporal workers (one per task queue) plus Prometheus metrics server and the cleanup schedule. All heavy imports are wrapped in `workflow.unsafe.imports_passed_through()` — required because Temporal workflow sandboxing forbids non-deterministic imports at module scope outside that block.
|
|
- **Workflows** (`model_manager/workflows/`): pure orchestration, no I/O beyond `execute_activity_method` calls.
|
|
- `train_model.py` (`TrainModel`, workflow name `train_model`): validate `experiment_run_id` → `load_model_metadata` → `validate_train_params` → update DB status `ORCHESTRATOR_WAITING_PROC` → `train_model` activity → update DB status `TRAINING_SUCCESS`/`TRAINING_ERROR` → `cleanup_resources` (always runs in `finally` if a `run_dir` exists, even on failure). Uses distinct `RetryPolicy`s per operation class: `network_retry_policy` (transient I/O), `no_retry_policy` (training/validation — data errors are permanent, and retrying would duplicate MLflow writes), `database_retry_policy` (Postgres). Timeouts are env-driven (`TIMEOUT_*`), extend `TIMEOUT_TRAIN_MODEL` instead of adding retries for slow training.
|
|
- `cleanup_files.py` (`CleanupFiles`): scheduled (`cleanup_schedule.py`, default daily at midnight UTC via `CLEANUP_CRON`) deletion of stale local temp dirs matching `name_YYYYMMDD_HHMMSS_microseconds`, older than `CLEANUP_RETENTION_HOURS`.
|
|
- Task queues are derived from `RUNTIME` env var (default `single`) via `prepare_worker.build_queue_name`: `train_model-<runtime>-queue`, `cleanup_files-<runtime>-queue`. A Temporal client starting a workflow must target the matching queue.
|
|
- **Activities** (`model_manager/activities/`): `Activities` (`activities.py`) combines `ExperimentTracking` + `Training` + `Cleanup` via multiple inheritance into one class registered with the worker — mind MRO and `__init__`/`__del__` ordering when touching this class (see the defensive `hasattr(self, 'engine')` check in `__del__`).
|
|
- `experiment_tracking.py`: single `update_experiment_run()` entrypoint handling three update kinds via `UpdateType`: `STATUS`, `STATUS_WITH_ERROR`, `MODEL_SAVED`. Error messages truncated to 1024 chars before persisting.
|
|
- `training.py`: single `train_model()` entrypoint. Receives pre-downloaded data as `BytesIO` (not re-downloaded per-activity) to avoid memory leaks; never raises on training failure — returns a result/error so the workflow can update DB status without an unhandled exception mid-activity.
|
|
- `cleanup.py`: local-filesystem-only cleanup (`cleanup_temp_directories`), no MinIO involved — MinIO artifacts are managed externally.
|
|
|
|
### Data/config layer
|
|
- `model_manager/utils/connectors_config.py`: single source of truth for env-var → config-dict builders (`build_postgres_config`, `build_mlflow_config`, `build_minio_config`, `build_mongodb_config`, `build_plugin_store_config`). Read this file (and `worker.py`) before trusting any env var default — the README table mirrors it but this file is authoritative.
|
|
- `model_manager/utils/models/train_model_params.py`: `TrainModelParams` — the workflow's single input contract. `from_dict` does type coercion; `validate_business_rules()` runs after, enforcing 7 rules (train_size range, non-empty variable_columns, required model_metadata, JSON-Schema validation of `data_model_kwargs`/`model_kwargs`/`opt_params` against schemas in `model_metadata` via `Draft202012Validator`, required non-blank strings, allowed `date_format` values, numeric `experiment_run_id`). Extending training params means updating this file's validation, not just the workflow.
|
|
- `model_manager/utils/repository/data_manager_repository.py`: data loading, feature prep, metrics, report generation — the actual ML/data logic activities delegate into.
|
|
- `model_manager/sientia/`: RCE drift metrics (`metrics.py`), Evidently-based HTML report generation (`reports.py`), custom exceptions.
|
|
- `model_manager/metrics.py`: Prometheus metric definitions (`app_up`, `workflow_execution_total`, `activity_execution_total`, training/cleanup/RCE-drift gauges). New metrics go here, not inline in activities.
|
|
|
|
### External dependencies
|
|
Two private Aignosi packages carry non-trivial logic: `sientia_do` (observability/logger, notification handler, metrics controller, Postgres/MinIO repositories) and `sientia_model` (Plugin Store, MLflow repository, model wrappers). Both are pinned in `requirements.txt`/`requirements-local.txt`; when behavior looks like it lives outside this repo, it's likely in one of these.
|
|
|
|
### Test bootstrap quirk
|
|
`tests/conftest.py` stubs `evidently` and parts of `sientia_do` at `pytest_configure` time so unit tests can import `model_manager.sientia.reports` without the full (heavy/optional) Evidently install. If you add new imports from these packages in production code paths that unit tests exercise, you likely need to extend the stubs here too.
|
|
|
|
### Test layout
|
|
- `tests/` — unit tests, mirrors `model_manager/` package structure (`tests/activities/`, `tests/schedules/`, `tests/sientia/`, etc.).
|
|
- `e2e/` — integration tests against a real Temporal/Postgres/MinIO/MongoDB stack (via `testcontainers`), driven by JSON scenario files in `docs/test-scenarios/` (see `e2e/scenarios.md` for the scenario→test mapping). Scenario payloads are snake_case, shaped like `TrainModelParams`/`input-sample.json`.
|
|
- New scenario: copy an existing `docs/test-scenarios/*.json`, add/extend a test in `e2e/test_train_model_workflow.py`, update `e2e/scenarios.md`.
|
|
|
|
## Conventions
|
|
- Ruff: line-length 100, single quotes, `select = [E,W,F,I,B,C4,UP,N,YTT,S,BLE,A,C90]`, max complexity 15. `N802`/`N806` ignored (Temporal decorators use non-lowercase names). `S101`/`S105`/`S106` ignored in tests.
|
|
- mypy: `disallow_untyped_defs = false` but `check_untyped_defs = true` — existing untyped defs are tolerated, but code inside them is still checked. Several external packages (`temporalio`, `mlflow`, `sientia_model`, etc.) have `ignore_missing_imports`; `model_repository.py`/`data_manager_repository.py` have `ignore_errors = true` (legacy — don't extend that list without reason).
|