meshscale_worker_space/app.py

from __future__ import annotations

import os
import threading
import time
import traceback
from typing import Any

from fastapi import FastAPI

from key_farm.contracts import FarmPolicy
from key_farm.store import FarmStore
from key_farm.worker import worker_loop


app = FastAPI(title="MeshScale CPU Worker")

_worker_thread: threading.Thread | None = None
_supervisor_thread: threading.Thread | None = None
_worker_error: str | None = None
_worker_lock = threading.Lock()


def _int_env(name: str, default: int) -> int:
    try:
        return int(os.environ.get(name, str(default)))
    except ValueError:
        return default


def _float_env(name: str, default: float) -> float:
    try:
        return float(os.environ.get(name, str(default)))
    except ValueError:
        return default


def _bool_env(name: str, default: bool) -> bool:
    raw = os.environ.get(name)
    if raw is None:
        return default
    return raw.strip().lower() not in {"0", "false", "no", "off"}


def _root() -> str:
    return os.environ.get("KEY_FARM_ROOT", "/data/farm")


def _formation() -> str:
    return os.environ.get("KEY_FARM_FORMATION", "cpu-worker")


def _worker_id() -> str:
    base = os.environ.get("KEY_FARM_WORKER_ID") or os.environ.get("SPACE_ID") or "meshscale-worker"
    runtime = (
        os.environ.get("SPACE_REPLICA_ID")
        or os.environ.get("REPLICA_ID")
        or os.environ.get("HOSTNAME")
        or str(os.getpid())
    )
    return f"{base}:{runtime}"


def _policy() -> FarmPolicy:
    formation = _formation()
    return FarmPolicy(
        target_batch_size=_int_env("KEY_FARM_TARGET_BATCH_SIZE", 24),
        min_batch_size=_int_env("KEY_FARM_MIN_BATCH_SIZE", 4),
        max_depth=_int_env("KEY_FARM_MAX_DEPTH", 4),
        max_workers=_int_env("KEY_FARM_MAX_WORKERS", 24),
        claim_ttl_seconds=_int_env("KEY_FARM_CLAIM_TTL_SECONDS", 900),
        pulse_ttl_seconds=_int_env("KEY_FARM_PULSE_TTL_SECONDS", 120),
        controller_pulse_ttl_seconds=_int_env("KEY_FARM_CONTROLLER_PULSE_TTL_SECONDS", 180),
        budget_hourly_usd=_float_env("KEY_FARM_BUDGET_HOURLY_USD", 0.72),
        cpu_space_hourly_usd=_float_env("KEY_FARM_CPU_SPACE_HOURLY_USD", 0.03),
        exit_when_pulse_lost=_bool_env("KEY_FARM_EXIT_WHEN_PULSE_LOST", True),
        pulse_rate=_float_env("KEY_FARM_PULSE_RATE", 1.0),
        allowed_formations=(formation,),
    )


def _run_worker() -> None:
    global _worker_error
    try:
        worker_loop(
            FarmStore(_root()),
            _worker_id(),
            _policy(),
            poll_seconds=_float_env("KEY_FARM_POLL_SECONDS", 2.0),
            once=False,
        )
    except Exception:
        _worker_error = traceback.format_exc()


def _autostart_enabled() -> bool:
    return _bool_env("MESHSCALE_WORKER_AUTOSTART", True)


def _controller_ready(store: FarmStore | None = None, policy: FarmPolicy | None = None) -> bool:
    if not _bool_env("MESHSCALE_WORKER_WAIT_FOR_CONTROLLER", True):
        return True
    store = store or FarmStore(_root())
    policy = policy or _policy()
    return store.controller_pulse_alive(policy.controller_pulse_ttl_seconds)


def _start_worker_if_ready() -> bool:
    global _worker_thread, _worker_error
    if not _autostart_enabled():
        return False
    if _worker_error and not _bool_env("MESHSCALE_WORKER_RESTART_ON_ERROR", False):
        return False
    store = FarmStore(_root())
    policy = _policy()
    if not _controller_ready(store, policy):
        return False
    with _worker_lock:
        if _worker_thread and _worker_thread.is_alive():
            return True
        _worker_error = None
        _worker_thread = threading.Thread(target=_run_worker, name="meshscale-worker", daemon=True)
        _worker_thread.start()
        return True


def _supervise_worker() -> None:
    while _autostart_enabled():
        _start_worker_if_ready()
        time.sleep(_float_env("MESHSCALE_WORKER_SUPERVISOR_SECONDS", 5.0))


def _ensure_supervisor_started() -> None:
    global _supervisor_thread
    if not _autostart_enabled():
        return
    if _supervisor_thread and _supervisor_thread.is_alive():
        return
    _supervisor_thread = threading.Thread(
        target=_supervise_worker,
        name="meshscale-worker-supervisor",
        daemon=True,
    )
    _supervisor_thread.start()


def _ensure_worker_started() -> None:
    if not _bool_env("MESHSCALE_WORKER_AUTOSTART", True):
        return
    _ensure_supervisor_started()
    _start_worker_if_ready()


@app.on_event("startup")
def startup() -> None:
    _ensure_worker_started()


@app.get("/health")
def health() -> dict[str, Any]:
    store = FarmStore(_root())
    policy = _policy()
    _ensure_worker_started()
    return {
        "ok": _worker_error is None,
        "worker_id": _worker_id(),
        "formation": _formation(),
        "thread_alive": bool(_worker_thread and _worker_thread.is_alive()),
        "supervisor_alive": bool(_supervisor_thread and _supervisor_thread.is_alive()),
        "root": str(store.root),
        "controller_alive": store.controller_pulse_alive(policy.controller_pulse_ttl_seconds),
        "regulation": store.read_regulation().to_dict(),
        "error": _worker_error,
    }


@app.get("/state")
def state() -> dict[str, Any]:
    store = FarmStore(_root())
    policy = _policy()
    _ensure_worker_started()
    return {
        "worker_id": _worker_id(),
        "formation": _formation(),
        "pending_jobs": len(store.pending_jobs(_formation())),
        "active_workers": store.active_worker_count(policy.pulse_ttl_seconds, _formation()),
        "controller_alive": store.controller_pulse_alive(policy.controller_pulse_ttl_seconds),
        "policy": policy.to_dict(),
    }