Astrape/astrape-oracle
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Hand nightly training to an external_tcn service when configured

Browse Source 
Set ASTRAPE_TCN_TRAINER_URL and the scheduler ships raw series to that
service and writes back the artifact it returns, instead of training in
process. This host has no AVX2 and torch dies with SIGILL in Adam, so remote
is the working path here; local training stays for hosts that can do it.

run() now picks between _train_local and _train_remote. The remote path never
builds examples locally -- it only reads the database -- and the far side
rebuilds them with the same gibil. The artifact is written last and
atomically, because the oracle daemon loads it on every forecast. A remote
failure raises SystemExit so the unit fails visibly rather than leaving
yesterday's model in place while reporting success. Events record trained_by,
trainer_url and job_id.

Also stops deriving the model, state and log paths from __file__. Where this
package is installed says nothing about where the deployment keeps them; they
now resolve against the working directory, which every unit sets.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
rpotter6298
2026-08-30 15:47:50 +02:00
parent 9a4ee8e0cf
commit 3bc31f9f68
1 changed files with 100 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/tcn/schedule.py
+100 -24
View File
@@ -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,