rpotter6298
ccd4e7a48c
The forecasting runtime gibil pilots: weather ingest (Open-Meteo forecasts and the Ecowitt station), the oracle and weather stores, the snapshot daemon, the forecast evaluator, and the TCN training scheduler. classes/datasets.py holds the shapes the stores return for display. They sit beside the stores deliberately: a store's return type belongs to the store, and keeping them with the renderers would mean this repo importing the dashboard, which is backwards. Extracted from a single repository where the forecasting engine, the database plumbing, the hardware client, the ingest daemons and the web UI all imported each other freely. Verified that every module imports with gibil resolved from its own installed checkout. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
302 lines
12 KiB
Python
302 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
from dataclasses import asdict, dataclass, field
|
|
from datetime import datetime, timedelta, timezone
|
|
from pathlib import Path
|
|
from time import monotonic
|
|
|
|
from astrape_ems.classes.env import EnvLoader
|
|
from astrape_ems.classes.store import SigenStore
|
|
from astrape_oracle.classes.weather.store import WeatherStore
|
|
from gibil.classes.usage_sequence_dataset import (
|
|
UsageSequenceDatasetBuilder,
|
|
UsageSequenceDatasetConfig,
|
|
)
|
|
from gibil.classes.usage_training import train_model
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[3]
|
|
DEFAULT_MODEL_PATH = PROJECT_ROOT / "models" / "usage_sequence_tcn_v1.pt"
|
|
DEFAULT_STATE_PATH = PROJECT_ROOT / "models" / "usage_sequence_training_state.json"
|
|
DEFAULT_LOG_PATH = PROJECT_ROOT / "logs" / "usage_sequence_training.jsonl"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TrainingRegimeConfig:
|
|
artifact_path: Path = DEFAULT_MODEL_PATH
|
|
state_path: Path = DEFAULT_STATE_PATH
|
|
log_path: Path = DEFAULT_LOG_PATH
|
|
full_epochs: int = 20
|
|
finetune_epochs: int = 1
|
|
batch_size: int = 32
|
|
full_learning_rate: float = 0.001
|
|
finetune_learning_rate: float = 0.0002
|
|
seed: int = 7
|
|
threshold_seconds: float = 10 * 60
|
|
terminal_threshold_seconds: float = 60 * 60
|
|
nightly_interval_days: float = 1.0
|
|
max_full_retrain_interval_days: float = 7.0
|
|
target_finetune_seconds: float = 5 * 60
|
|
finetune_min_samples: int = 32
|
|
|
|
@classmethod
|
|
def from_env(cls) -> "TrainingRegimeConfig":
|
|
EnvLoader().load()
|
|
return cls(
|
|
artifact_path=Path(
|
|
os.environ.get(
|
|
"ASTRAPE_USAGE_SEQUENCE_ARTIFACT_PATH",
|
|
str(DEFAULT_MODEL_PATH),
|
|
)
|
|
),
|
|
state_path=Path(
|
|
os.environ.get(
|
|
"ASTRAPE_USAGE_SEQUENCE_TRAINING_STATE_PATH",
|
|
str(DEFAULT_STATE_PATH),
|
|
)
|
|
),
|
|
log_path=Path(
|
|
os.environ.get(
|
|
"ASTRAPE_USAGE_SEQUENCE_TRAINING_LOG_PATH",
|
|
str(DEFAULT_LOG_PATH),
|
|
)
|
|
),
|
|
full_epochs=int(os.environ.get("ASTRAPE_USAGE_SEQUENCE_FULL_EPOCHS", "20")),
|
|
finetune_epochs=int(
|
|
os.environ.get("ASTRAPE_USAGE_SEQUENCE_FINETUNE_EPOCHS", "1")
|
|
),
|
|
batch_size=int(os.environ.get("ASTRAPE_USAGE_SEQUENCE_BATCH_SIZE", "32")),
|
|
full_learning_rate=float(
|
|
os.environ.get("ASTRAPE_USAGE_SEQUENCE_FULL_LR", "0.001")
|
|
),
|
|
finetune_learning_rate=float(
|
|
os.environ.get("ASTRAPE_USAGE_SEQUENCE_FINETUNE_LR", "0.0002")
|
|
),
|
|
seed=int(os.environ.get("ASTRAPE_USAGE_SEQUENCE_SEED", "7")),
|
|
threshold_seconds=float(
|
|
os.environ.get("ASTRAPE_USAGE_SEQUENCE_FULL_T_SECONDS", str(10 * 60))
|
|
),
|
|
terminal_threshold_seconds=float(
|
|
os.environ.get("ASTRAPE_USAGE_SEQUENCE_FULL_T2_SECONDS", str(60 * 60))
|
|
),
|
|
max_full_retrain_interval_days=float(
|
|
os.environ.get("ASTRAPE_USAGE_SEQUENCE_MAX_FULL_INTERVAL_DAYS", "7")
|
|
),
|
|
target_finetune_seconds=float(
|
|
os.environ.get("ASTRAPE_USAGE_SEQUENCE_TARGET_FINETUNE_SECONDS", str(5 * 60))
|
|
),
|
|
finetune_min_samples=int(
|
|
os.environ.get("ASTRAPE_USAGE_SEQUENCE_FINETUNE_MIN_SAMPLES", "32")
|
|
),
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class TrainingRegimeState:
|
|
last_full_training_seconds: float | None = None
|
|
last_full_sample_count: int | None = None
|
|
last_full_completed_at: str | None = None
|
|
last_finetune_completed_at: str | None = None
|
|
full_retrain_locked: bool = False
|
|
run_count: int = 0
|
|
history: list[dict[str, object]] = field(default_factory=list)
|
|
|
|
@classmethod
|
|
def load(cls, path: Path) -> "TrainingRegimeState":
|
|
if not path.exists():
|
|
return cls()
|
|
return cls(**json.loads(path.read_text()))
|
|
|
|
def save(self, path: Path) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(json.dumps(asdict(self), indent=2, sort_keys=True) + "\n")
|
|
|
|
|
|
class UsageSequenceTrainingScheduler:
|
|
def __init__(self, config: TrainingRegimeConfig) -> None:
|
|
self.config = config
|
|
|
|
def run(self, force_mode: str | None = None, dry_run: bool = False) -> dict[str, object]:
|
|
state = TrainingRegimeState.load(self.config.state_path)
|
|
decision = self._decision(state, force_mode=force_mode)
|
|
if dry_run:
|
|
event = {
|
|
"created_at": self._now_iso(),
|
|
"dry_run": True,
|
|
"decision": decision,
|
|
"state": asdict(state),
|
|
}
|
|
self._log_event(event)
|
|
print(json.dumps(event, sort_keys=True))
|
|
return event
|
|
|
|
dataset_config = UsageSequenceDatasetConfig.from_env()
|
|
builder = UsageSequenceDatasetBuilder(
|
|
config=dataset_config,
|
|
plant_history=SigenStore.from_env(),
|
|
weather_history=WeatherStore.from_env(),
|
|
)
|
|
examples = builder.build(limit=decision["sample_limit"])
|
|
if not examples:
|
|
raise SystemExit("No usage sequence examples available for scheduled training.")
|
|
|
|
artifact_path = self.config.artifact_path
|
|
base_artifact_path = artifact_path if decision["mode"] == "finetune" else None
|
|
if decision["mode"] == "finetune" and not artifact_path.exists():
|
|
decision = self._full_decision(reason="missing_artifact")
|
|
base_artifact_path = None
|
|
|
|
started = monotonic()
|
|
result = train_model(
|
|
examples=examples,
|
|
builder=builder,
|
|
epochs=decision["epochs"],
|
|
batch_size=self.config.batch_size,
|
|
learning_rate=decision["learning_rate"],
|
|
artifact_path=artifact_path,
|
|
seed=self.config.seed,
|
|
base_artifact_path=base_artifact_path,
|
|
mode=decision["mode"],
|
|
)
|
|
training_seconds = monotonic() - started
|
|
|
|
event = {
|
|
"created_at": self._now_iso(),
|
|
"mode": result.mode,
|
|
"reason": decision["reason"],
|
|
"training_seconds": training_seconds,
|
|
"sample_count": result.sample_count,
|
|
"training_examples": result.training_examples,
|
|
"validation_examples": result.validation_examples,
|
|
"final_train_loss": result.final_train_loss,
|
|
"final_validation_loss": result.final_validation_loss,
|
|
"artifact_path": str(result.artifact_path),
|
|
"next_full_interval_days": self._full_interval_days(training_seconds),
|
|
}
|
|
self._update_state(state, event)
|
|
state.save(self.config.state_path)
|
|
self._log_event(event)
|
|
print(json.dumps(event, sort_keys=True))
|
|
return event
|
|
|
|
def _decision(
|
|
self,
|
|
state: TrainingRegimeState,
|
|
force_mode: str | None,
|
|
) -> dict[str, object]:
|
|
if force_mode == "full":
|
|
return self._full_decision(reason="forced_full")
|
|
if force_mode == "finetune":
|
|
return self._finetune_decision(state, reason="forced_finetune")
|
|
if state.full_retrain_locked:
|
|
return self._finetune_decision(state, reason="full_retrain_locked")
|
|
if state.last_full_training_seconds is None:
|
|
return self._full_decision(reason="no_full_training_history")
|
|
if self._due_for_full_retrain(state):
|
|
return self._full_decision(reason="full_retrain_due")
|
|
return self._finetune_decision(state, reason="nightly_finetune")
|
|
|
|
def _due_for_full_retrain(self, state: TrainingRegimeState) -> bool:
|
|
if state.last_full_completed_at is None:
|
|
return True
|
|
completed_at = datetime.fromisoformat(state.last_full_completed_at)
|
|
due_after = completed_at + timedelta(
|
|
days=self._full_interval_days(state.last_full_training_seconds)
|
|
)
|
|
return datetime.now(timezone.utc) >= due_after
|
|
|
|
def _full_decision(self, reason: str) -> dict[str, object]:
|
|
return {
|
|
"mode": "full",
|
|
"reason": reason,
|
|
"sample_limit": None,
|
|
"epochs": self.config.full_epochs,
|
|
"learning_rate": self.config.full_learning_rate,
|
|
}
|
|
|
|
def _finetune_decision(
|
|
self,
|
|
state: TrainingRegimeState,
|
|
reason: str,
|
|
) -> dict[str, object]:
|
|
return {
|
|
"mode": "finetune",
|
|
"reason": reason,
|
|
"sample_limit": self._finetune_sample_limit(state),
|
|
"epochs": self.config.finetune_epochs,
|
|
"learning_rate": self.config.finetune_learning_rate,
|
|
}
|
|
|
|
def _finetune_sample_limit(self, state: TrainingRegimeState) -> int | None:
|
|
if not state.last_full_training_seconds or not state.last_full_sample_count:
|
|
return None
|
|
samples_per_second = state.last_full_sample_count / state.last_full_training_seconds
|
|
sample_limit = round(samples_per_second * self.config.target_finetune_seconds)
|
|
return max(self.config.finetune_min_samples, sample_limit)
|
|
|
|
def _full_interval_days(self, training_seconds: float | None) -> float:
|
|
if training_seconds is None or training_seconds <= self.config.threshold_seconds:
|
|
return self.config.nightly_interval_days
|
|
if training_seconds >= self.config.terminal_threshold_seconds:
|
|
return self.config.max_full_retrain_interval_days
|
|
span = self.config.terminal_threshold_seconds - self.config.threshold_seconds
|
|
progress = (training_seconds - self.config.threshold_seconds) / span
|
|
return self.config.nightly_interval_days + progress * (
|
|
self.config.max_full_retrain_interval_days - self.config.nightly_interval_days
|
|
)
|
|
|
|
def _update_state(
|
|
self,
|
|
state: TrainingRegimeState,
|
|
event: dict[str, object],
|
|
) -> None:
|
|
state.run_count += 1
|
|
if event["mode"] == "full":
|
|
state.last_full_training_seconds = float(event["training_seconds"])
|
|
state.last_full_sample_count = int(event["sample_count"])
|
|
state.last_full_completed_at = str(event["created_at"])
|
|
if state.last_full_training_seconds >= self.config.terminal_threshold_seconds:
|
|
state.full_retrain_locked = True
|
|
else:
|
|
state.last_finetune_completed_at = str(event["created_at"])
|
|
state.history.append(event)
|
|
state.history = state.history[-30:]
|
|
|
|
def _log_event(self, event: dict[str, object]) -> None:
|
|
self.config.log_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with self.config.log_path.open("a") as handle:
|
|
handle.write(json.dumps(event, sort_keys=True) + "\n")
|
|
|
|
def _now_iso(self) -> str:
|
|
return datetime.now(timezone.utc).isoformat()
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Adaptive full-retrain/fine-tune scheduler for the usage TCN."
|
|
)
|
|
parser.add_argument("--dry-run", action="store_true")
|
|
parser.add_argument(
|
|
"--force-mode",
|
|
choices=("full", "finetune"),
|
|
default=None,
|
|
help="Override the adaptive decision for this run.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
EnvLoader().load()
|
|
args = parse_args()
|
|
UsageSequenceTrainingScheduler(TrainingRegimeConfig.from_env()).run(
|
|
force_mode=args.force_mode,
|
|
dry_run=args.dry_run,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|