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.remote_training import ( RemoteTrainingError, RemoteUsageSequenceTrainer, ) from gibil.classes.usage_sequence_dataset import ( UsageSequenceDatasetBuilder, UsageSequenceDatasetConfig, ) from gibil.classes.usage_training import train_model # Resolved against the working directory, which every unit sets to # /srv/astrape. Not derived from __file__: where this package is installed says # nothing about where the deployment keeps its models and logs. DEFAULT_MODEL_PATH = Path("models") / "usage_sequence_tcn_v1.pt" DEFAULT_STATE_PATH = Path("models") / "usage_sequence_training_state.json" DEFAULT_LOG_PATH = Path("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 # When set, training is handed to an external_tcn service instead of run # locally. This host has no AVX2 and torch dies in Adam, so in production # this is the working path; local training stays for hosts that can do it. trainer_url: str | None = None trainer_timeout_seconds: float = 5400.0 @classmethod def from_env(cls) -> "TrainingRegimeConfig": EnvLoader().load() return cls( trainer_url=os.environ.get("ASTRAPE_TCN_TRAINER_URL") or None, trainer_timeout_seconds=float( os.environ.get("ASTRAPE_TCN_TRAINER_TIMEOUT_SECONDS", "5400") ), 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() 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() if self.config.trainer_url: metrics = self._train_remote(dataset_config, decision, artifact_path, base_artifact_path) else: metrics = self._train_local(dataset_config, decision, artifact_path, base_artifact_path) training_seconds = monotonic() - started event = { "created_at": self._now_iso(), "reason": decision["reason"], "training_seconds": training_seconds, "artifact_path": str(artifact_path), "next_full_interval_days": self._full_interval_days(training_seconds), **metrics, } 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 _train_local( self, dataset_config: UsageSequenceDatasetConfig, decision: dict[str, object], artifact_path: Path, base_artifact_path: Path | None, ) -> dict[str, object]: """Build examples and train in this process. Needs a torch that runs here.""" 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.") 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"], ) return { "mode": result.mode, "trained_by": "local", "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, } def _train_remote( self, dataset_config: UsageSequenceDatasetConfig, decision: dict[str, object], artifact_path: Path, base_artifact_path: Path | None, ) -> dict[str, object]: """Ship the raw series to an external_tcn service and fetch the artifact. Examples are rebuilt on the far side by the same gibil, so nothing is built here -- this host only reads the database. A failure raises SystemExit so the unit fails visibly rather than leaving yesterday's model in place while reporting success. """ trainer = RemoteUsageSequenceTrainer( self.config.trainer_url, job_timeout_seconds=self.config.trainer_timeout_seconds, ) try: result = trainer.train( plant_history=SigenStore.from_env(), weather_history=WeatherStore.from_env(), dataset_config=dataset_config, epochs=decision["epochs"], batch_size=self.config.batch_size, learning_rate=decision["learning_rate"], seed=self.config.seed, mode=decision["mode"], limit=decision["sample_limit"], base_artifact_path=base_artifact_path, ) except RemoteTrainingError as error: raise SystemExit(f"remote training failed via {self.config.trainer_url}: {error}") # Written last, and atomically: the oracle daemon loads this file on # every forecast and must never see a partial model. result.write_artifact(artifact_path) metrics = dict(result.metrics) metrics["mode"] = metrics.get("mode", decision["mode"]) metrics["trained_by"] = "remote" metrics["trainer_url"] = self.config.trainer_url metrics["job_id"] = result.job_id return metrics 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()