genesis_engine/engine/cosmos_54d_train_loop.py

#!/usr/bin/env python3
"""
Persistent continual 54D learning loop.

Watches Cosmos/data/cosmos/experience_corpus.txt and runs one bounded CPU
training burst only after enough genuinely new lines arrive. The last trained
line count survives restarts, a process lock prevents duplicate learners, and a
parent/stop-file contract lets WAKE_HER shut the learner down with Cosmos.
"""
from __future__ import annotations

import json
import os
from pathlib import Path
import subprocess
import sys
import time


PR = Path(__file__).resolve().parents[1]
CORPUS = PR / "Cosmos" / "data" / "cosmos" / "experience_corpus.txt"
TRAINER = PR / "scripts" / "train_cosmos_from_play.py"
CHECKPOINT = PR / "Cosmos" / "checkpoints" / "cosmos" / "cosmos_play.pt"
STATE = PR / "Cosmos" / "checkpoints" / "cosmos" / "cosmos_54d_loop_state.json"
LOCK = PR / "Cosmos" / "checkpoints" / "cosmos" / "cosmos_54d_loop.lock"
INTERVAL = max(30, int(os.getenv("COSMOS_54D_TRAIN_INTERVAL_S", "1800")))
MIN_NEW = max(1, int(os.getenv("COSMOS_54D_TRAIN_MIN_NEW", "12")))
START_DELAY = max(0, int(os.getenv("COSMOS_54D_TRAIN_START_DELAY_S", "0")))
PARENT_PID = int(os.getenv("COSMOS_54D_PARENT_PID", "0") or 0)
STOP_FILE = Path(os.getenv("COSMOS_54D_STOP_FILE", "")) if os.getenv("COSMOS_54D_STOP_FILE") else None


def _say(message: str) -> None:
    print(f"[54D-LOOP] {message}", flush=True)


def _lines() -> int:
    try:
        with CORPUS.open("r", encoding="utf-8", errors="ignore") as handle:
            return sum(1 for _ in handle)
    except Exception:
        return 0


def _parent_alive() -> bool:
    """Return whether the wake process still exists on every supported OS.

    ``os.kill(pid, 0)`` is not a reliable Windows existence probe: on some
    Python builds it raises ``SystemError``/WinError 87 even for a live PID.
    Use a real process handle on Windows, then retain the POSIX probe elsewhere.
    """
    if PARENT_PID <= 0:
        return True
    if os.name == "nt":
        try:
            import ctypes

            access = 0x00100000 | 0x00001000  # SYNCHRONIZE | QUERY_LIMITED_INFORMATION
            handle = ctypes.windll.kernel32.OpenProcess(access, False, PARENT_PID)
            if handle:
                ctypes.windll.kernel32.CloseHandle(handle)
                return True
            return False
        except Exception:
            # A probe failure must never crash the learner. The stop file is
            # still authoritative when WAKE_HER shuts down.
            return True
    try:
        os.kill(PARENT_PID, 0)
        return True
    except (OSError, SystemError, ValueError):
        return False


def _stopping() -> bool:
    return bool((STOP_FILE and STOP_FILE.exists()) or not _parent_alive())


def _sleep(seconds: int) -> bool:
    deadline = time.monotonic() + max(0, seconds)
    while time.monotonic() < deadline:
        if _stopping():
            return False
        time.sleep(min(2.0, max(0.0, deadline - time.monotonic())))
    return not _stopping()


def _single_instance():
    LOCK.parent.mkdir(parents=True, exist_ok=True)
    handle = LOCK.open("a+b")
    handle.seek(0, os.SEEK_END)
    if handle.tell() == 0:
        handle.write(b"0")
        handle.flush()
    handle.seek(0)
    try:
        if os.name == "nt":
            import msvcrt

            msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1)
        else:
            import fcntl

            fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
    except (OSError, IOError):
        handle.close()
        return None
    return handle


def _load_last_trained(current_lines: int) -> int:
    try:
        value = json.loads(STATE.read_text(encoding="utf-8"))
        return max(0, int(value.get("last_trained_lines", 0)))
    except Exception:
        pass
    # Migrate safely: if the checkpoint is newer than the corpus, it already
    # represents the current data. Otherwise the accumulated new corpus is due.
    try:
        if CHECKPOINT.is_file() and CHECKPOINT.stat().st_mtime_ns >= CORPUS.stat().st_mtime_ns:
            return current_lines
    except Exception:
        pass
    return 0


def _save_state(lines: int) -> None:
    STATE.parent.mkdir(parents=True, exist_ok=True)
    value = {
        "schema": "cosmos.54d_continual_state.v1",
        "last_trained_lines": int(lines),
        "checkpoint": str(CHECKPOINT),
        "checkpoint_mtime_ns": CHECKPOINT.stat().st_mtime_ns if CHECKPOINT.is_file() else None,
        "updated_at": time.time(),
    }
    temp = STATE.with_name(f"{STATE.name}.tmp.{os.getpid()}")
    temp.write_text(json.dumps(value, indent=2) + "\n", encoding="utf-8")
    os.replace(temp, STATE)


def _mem_too_high() -> bool:
    try:
        import psutil

        vm = psutil.virtual_memory()
        return vm.percent >= 90.0 or vm.available < 1_500_000_000
    except Exception:
        return False


def _run_burst() -> tuple[int | None, str]:
    env = dict(os.environ)
    env.setdefault("COSMOS_PLAY_TRAIN_STEPS", "120")
    env.setdefault("COSMOS_PLAY_TRAIN_SEC", "360")
    env.setdefault("COSMOS_PLAY_TRAIN_THREADS", "2")
    env.setdefault("PYTHONIOENCODING", "utf-8")
    proc = subprocess.Popen(
        [sys.executable, str(TRAINER)],
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        encoding="utf-8",
        errors="replace",
    )
    deadline = time.monotonic() + 900
    while proc.poll() is None:
        if _stopping():
            proc.terminate()
            try:
                proc.wait(timeout=10)
            except Exception:
                proc.kill()
            output, _ = proc.communicate()
            return None, output or ""
        if time.monotonic() >= deadline:
            proc.kill()
            output, _ = proc.communicate()
            return -1, (output or "") + "\ntrainer timeout"
        time.sleep(1)
    output, _ = proc.communicate()
    return proc.returncode, output or ""


def main() -> int:
    lock_handle = _single_instance()
    if lock_handle is None:
        _say("another continual learner already owns the lock; this copy exits")
        return 0
    try:
        _say(
            f"armed: {MIN_NEW}+ new lines, check every {INTERVAL}s, "
            f"start delay {START_DELAY}s"
        )
        if START_DELAY and not _sleep(START_DELAY):
            return 0
        current = _lines()
        last = _load_last_trained(current)
        if last > current:
            last = current
            _save_state(last)
        elif not STATE.exists() and last == current:
            _save_state(last)
        _say(f"persistent baseline: corpus={current}, last_trained={last}")

        while not _stopping():
            current = _lines()
            new_lines = current - last
            if new_lines >= MIN_NEW:
                if _mem_too_high():
                    _say(
                        f"corpus {current} (+{new_lines}) - RAM guard active; "
                        "training deferred"
                    )
                else:
                    before = CHECKPOINT.stat().st_mtime_ns if CHECKPOINT.is_file() else 0
                    _say(f"corpus {current} (+{new_lines}) - bounded training burst")
                    code, output = _run_burst()
                    tail = "\n".join(output.strip().splitlines()[-8:])
                    if tail:
                        _say("trainer tail:\n" + tail)
                    after = CHECKPOINT.stat().st_mtime_ns if CHECKPOINT.is_file() else 0
                    if code == 0 and after != before:
                        last = current
                        _save_state(last)
                        _say(f"checkpoint advanced; durable baseline={last}")
                    elif code is None:
                        _say("stop requested; active trainer terminated cleanly")
                        break
                    else:
                        _say(
                            f"burst did not advance checkpoint (return={code}); "
                            "new lines remain due for retry"
                        )
            else:
                _say(f"corpus {current} (+{new_lines}) - waiting for {MIN_NEW - new_lines} more")
            if not _sleep(INTERVAL):
                break
        _say("stopped with Cosmos; checkpoint and corpus remain persistent")
        return 0
    finally:
        lock_handle.close()


if __name__ == "__main__":
    raise SystemExit(main())