astrape-oracle/scripts/tcn/schedule.py

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