cloud/worker.py

"""Restore durable state before starting any real inference or publication."""
import os
from pathlib import Path
import sys

ROOT = Path(__file__).resolve().parents[1]
sys.path[:0] = [str(ROOT / "src"), str(ROOT)]
os.chdir(ROOT)

from tsfm_bench.remote_state import RemoteState
from scripts import run_online_daemon


def main():
    for key in ("HF_TOKEN", "HF_STATE_REPO", "TSFM_API_KEY", "TABPFN_TOKEN"):
        if not os.getenv(key):
            raise RuntimeError(f"Required Space secret/variable missing: {key}")
    state = RemoteState(ROOT, os.environ["HF_STATE_REPO"], os.environ["HF_TOKEN"])
    restored = state.restore()
    print(f"Restored operator checkpoint from {restored['checkpointed_at']}", flush=True)
    run_online_daemon.REMOTE_STATE = state
    run_online_daemon.main()


if __name__ == "__main__":
    main()