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()