Flatten to a single package directory
The checkout is now astrape/astrape-oracle, not astrape/astrape-oracle/<package>. The package contents move to the repository root and pyproject maps the package name to '.' via package-dir, with the subpackages listed explicitly. The systemd units move out to /srv/astrape/deploy/systemd, outside every repository, because they describe this deployment rather than any one component. Verified both install modes: editable, which is how the deployment runs, and a plain wheel install, which is what the container build does. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
rpotter6298
@@ -0,0 +1,301 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user