diff --git a/scripts/tcn/schedule.py b/scripts/tcn/schedule.py index 5ad65ab..bc6a530 100644 --- a/scripts/tcn/schedule.py +++ b/scripts/tcn/schedule.py @@ -11,6 +11,10 @@ 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, @@ -18,10 +22,12 @@ from gibil.classes.usage_sequence_dataset import ( 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" +# 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) @@ -41,11 +47,20 @@ class TrainingRegimeConfig: 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", @@ -134,6 +149,41 @@ class UsageSequenceTrainingScheduler: 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(), @@ -143,13 +193,6 @@ class UsageSequenceTrainingScheduler: 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, @@ -161,26 +204,59 @@ class UsageSequenceTrainingScheduler: base_artifact_path=base_artifact_path, mode=decision["mode"], ) - training_seconds = monotonic() - started - - event = { - "created_at": self._now_iso(), + return { "mode": result.mode, - "reason": decision["reason"], - "training_seconds": training_seconds, + "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, - "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 _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,