src/server/spec.py

from dataclasses import dataclass
from torch import Tensor
from typing import Dict, Optional, List, Tuple

import atexit
import io
import os
import signal
import subprocess
import sys
import time
import torch
import uuid

from ..rig_package.info.asset import Asset
from ..model.tokenrig import TokenRig

# ---------------------------------------------------------------------------
# File-based IPC directory (replaces TCP port listening)
# ---------------------------------------------------------------------------
BPY_IPC_DIR = os.path.join(os.environ.get("TMPDIR", "/tmp"), "bpy_ipc")
BPY_IPC_POLL_INTERVAL = 0.05   # 50 ms
BPY_IPC_TIMEOUT = 180          # 3 min max wait for a single response
BPY_IPC_STARTUP_TIMEOUT = 30   # 0.5 min for bpy import (cold start)

@dataclass
class TensorPacket:
    """make sure stays on cpu"""
    validate: bool=False
    know_skeleton: bool=False
    learned_mesh_cond: Optional[Tensor]=None
    cond_latents: Optional[Tensor]=None
    mesh_cond: Optional[Tensor]=None
    vertices: Optional[Tensor]=None
    assets: Optional[List[Asset]]=None
    output_ids: Optional[Tensor]=None
    start_embed_list: Optional[List[Tensor]]=None
    start_tokens_list: Optional[List[List[int]]]=None

    def to_device(self, device):
        if self.learned_mesh_cond is not None:
            self.learned_mesh_cond = self.learned_mesh_cond.to(device)
        if self.cond_latents is not None:
            self.cond_latents = self.cond_latents.to(device)
        if self.mesh_cond is not None:
            self.mesh_cond = self.mesh_cond.to(device)
        if self.vertices is not None:
            self.vertices = self.vertices.to(device)
        if self.output_ids is not None:
            self.output_ids = self.output_ids.to(device)
        if self.start_embed_list is not None:
            self.start_embed_list = [x.to(device) for x in self.start_embed_list]

    @property
    def B(self):
        assert self.learned_mesh_cond is not None
        return self.learned_mesh_cond.shape[0]

    def to_bytes(self):
        return object_to_bytes(self)

    @classmethod
    def from_bytes(cls, bytes) -> 'TensorPacket':
        return bytes_to_object(bytes)


def object_to_bytes(t):
    buffer = io.BytesIO()
    torch.save(t, buffer)
    return buffer.getvalue()

def bytes_to_object(b, map_location=None):
    return torch.load(io.BytesIO(b), weights_only=False, map_location=map_location)

def get_model(
    ckpt_path: str,
    hf_path: Optional[str]=None,
    device='cuda',
) -> TokenRig:
    model = TokenRig.load_from_system_checkpoint(checkpoint_path=ckpt_path)
    if hf_path is not None:
        from transformers import AutoModel
        a = AutoModel.from_pretrained(
            hf_path,
            local_files_only=True,
            _attn_implementation="flash_attention_2",
            torch_dtype=torch.bfloat16,
        )
        model.transformer.model.load_state_dict(a.state_dict())

    model = model.to(device)
    return model


# ---------------------------------------------------------------------------
# BpyClient — file-based IPC with the bpy worker process
# ---------------------------------------------------------------------------
# Communication happens entirely via the filesystem under ``BPY_IPC_DIR``:
#
#   Request:   {BPY_IPC_DIR}/req_{uuid}.torch   (torch-serialised (op, data))
#              {BPY_IPC_DIR}/req_{uuid}.ready    (empty flag — signals ready)
#   Response:  {BPY_IPC_DIR}/resp_{uuid}.torch   (torch-serialised result)
#              {BPY_IPC_DIR}/resp_{uuid}.done    (empty flag — signals done)
#
# The bpy worker polls for ``*.ready`` files, processes them, writes the
# matching ``resp_*.torch``+``.done``, then removes the request pair.
#
# ``bpy_server_request`` is the low-level stateless function — safe to call
# from any process (main or ZeroGPU worker).  ``BpyClient`` wraps it with
# subprocess lifecycle management (launch / shutdown) for the main process.
# ---------------------------------------------------------------------------

def bpy_server_request(op: str, data, proc: "Optional[subprocess.Popen]" = None,
                       timeout: "Optional[float]" = None):
    """Send *op* with *data* to the bpy worker via the filesystem IPC dir.

    This is a **stateless** function that only does file I/O — it can be
    called from the main process or any ZeroGPU worker process.

    Parameters
    ----------
    op : str
        One of ``"ping"``, ``"load"``, ``"export"``, ``"transfer"``.
    data : object
        Arbitrary pickle-able / torch-save-able payload.
    proc : subprocess.Popen or None
        If a handle is provided the function will abort early when the
        worker process dies unexpectedly.
    timeout : float or None
        Maximum seconds to wait.  Defaults to ``BPY_IPC_TIMEOUT``.

    Returns
    -------
    object
        Deserialised result.

    Raises
    ------
    RuntimeError
        On timeout, worker crash, or worker-reported error.
    """
    if timeout is None:
        timeout = BPY_IPC_TIMEOUT

    # Ensure IPC directory exists (may be called from worker before main init)
    os.makedirs(BPY_IPC_DIR, exist_ok=True)

    req_id = f"req_{uuid.uuid4().hex}"
    req_torch = os.path.join(BPY_IPC_DIR, f"{req_id}.torch")
    req_ready = os.path.join(BPY_IPC_DIR, f"{req_id}.ready")
    resp_torch = os.path.join(BPY_IPC_DIR, f"resp_{req_id[4:]}.torch")
    resp_done = os.path.join(BPY_IPC_DIR, f"resp_{req_id[4:]}.done")

    # Write request data to temp file
    with open(req_torch, "wb") as f:
        f.write(object_to_bytes((op, data)))

    # Signal readiness (atomic — after file is fully written)
    with open(req_ready, "w") as f:
        f.write("")

    # Wait for response
    t0 = time.time()
    _pid_check_interval = 2.0   # check PID every 2 s
    _last_pid_check = 0.0
    while not os.path.exists(resp_done):
        elapsed = time.time() - t0
        if elapsed > timeout:
            _cleanup_files(req_torch, req_ready)
            raise RuntimeError(
                f"bpy_server_request timeout after {timeout:.0f}s "
                f"waiting for op='{op}' (id={req_id})"
            )
        if proc is not None and proc.poll() is not None:
            _cleanup_files(req_torch, req_ready)
            raise RuntimeError(
                f"bpy worker process exited unexpectedly (code={proc.returncode})"
            )
        # Periodically verify the bpy worker PID is still alive
        # (catches crashes even when we don't have a proc handle).
        if elapsed - _last_pid_check > _pid_check_interval:
            _last_pid_check = elapsed
            if _bpy_server_pid() is None:
                _cleanup_files(req_torch, req_ready)
                raise RuntimeError(
                    f"bpy worker process died while waiting for op='{op}' (id={req_id})"
                )
        time.sleep(BPY_IPC_POLL_INTERVAL)

    # Read response
    try:
        with open(resp_torch, "rb") as f:
            result = bytes_to_object(f.read())
    except Exception as exc:
        _cleanup_files(req_torch, req_ready, resp_torch, resp_done)
        raise RuntimeError(f"Failed to read response: {exc}")

    # Cleanup
    _cleanup_files(req_torch, req_ready, resp_torch, resp_done)

    # Surface errors from the worker
    if isinstance(result, str):
        if result.startswith("error:"):
            raise RuntimeError(f"bpy worker error: {result}")
        if result.startswith("unsupported op"):
            raise RuntimeError(f"bpy worker: {result}")

    return result


def bpy_server_ping(timeout: float = 2.0) -> bool:
    """Return ``True`` if the bpy worker is alive and responsive.

    Uses a short *timeout* (default 2 s) so callers are not blocked when
    no bpy worker is running yet.
    """
    try:
        return bpy_server_request("ping", None, timeout=timeout) == "pong"
    except Exception:
        return False


def _bpy_server_pid() -> "Optional[int]":
    """Return the PID of a running bpy worker, or ``None``.

    Reads ``{BPY_IPC_DIR}/.pid`` and verifies the process still exists.
    This is a fast check (no I/O wait).
    """
    pid_file = os.path.join(BPY_IPC_DIR, ".pid")
    try:
        with open(pid_file, "r") as f:
            pid = int(f.read().strip())
        os.kill(pid, 0)  # signal 0 = existence check only
        return pid
    except (FileNotFoundError, ValueError, OSError):
        return None


def is_bpy_server_alive() -> bool:
    """Fast check: is a bpy worker process running?

    First checks the PID file (instant).  If the PID is alive, optionally
    verifies responsiveness with a short ping.
    """
    return _bpy_server_pid() is not None


def _cleanup_files(*paths: str):
    """Remove each file if it exists; never raise."""
    for p in paths:
        try:
            os.remove(p)
        except OSError:
            pass


def _cleanup_ipc_dir():
    """Remove all files from the IPC directory (not the dir itself)."""
    if not os.path.isdir(BPY_IPC_DIR):
        return
    for fn in os.listdir(BPY_IPC_DIR):
        try:
            os.remove(os.path.join(BPY_IPC_DIR, fn))
        except OSError:
            pass


def _kill_stale_bpy_worker():
    """If a bpy worker PID file exists and the process is alive, kill it.

    This prevents orphaned/hung workers from interfering with a new one.
    """
    pid = _bpy_server_pid()
    if pid is None:
        return
    print(f"[BpyClient] Killing stale bpy worker (pid={pid})", flush=True)
    try:
        os.kill(pid, signal.SIGTERM)
        # Wait briefly for graceful shutdown
        for _ in range(50):  # 5 s max
            try:
                os.kill(pid, 0)
            except OSError:
                break
            time.sleep(0.1)
        else:
            # Force kill
            os.kill(pid, signal.SIGKILL)
    except OSError:
        pass


class BpyClient:
    """Manages the bpy subprocess lifecycle and provides a ``request()`` API.

    Typical usage (singleton in the main process)::

        client = BpyClient.launch()
        asset  = client.request("load", "/path/to/model.fbx")
        client.request("export", {"asset": asset, "filepath": "/out.glb"})

    Worker processes that don't own the subprocess should use the stateless
    ``bpy_server_request()`` function directly instead.
    """

    def __init__(self, proc: subprocess.Popen):
        self._proc = proc

    # -- public API ---------------------------------------------------------

    def request(self, op: str, data):
        """Thin wrapper around :func:`bpy_server_request` that also monitors
        the subprocess handle for unexpected exits."""
        return bpy_server_request(op, data, proc=self._proc)

    def ping(self) -> bool:
        return bpy_server_ping()

    def is_alive(self) -> bool:
        """Check whether the subprocess is still running."""
        return self._proc is not None and self._proc.poll() is None

    def shutdown(self):
        """Terminate the bpy worker process and clean up the IPC directory."""
        if self._proc is None:
            return
        print(f"[BpyClient] Terminating bpy worker (pid={self._proc.pid})", flush=True)
        try:
            os.killpg(os.getpgid(self._proc.pid), signal.SIGTERM)
        except ProcessLookupError:
            pass
        try:
            self._proc.wait(timeout=10)
        except subprocess.TimeoutExpired:
            os.killpg(os.getpgid(self._proc.pid), signal.SIGKILL)
            self._proc.wait()
        self._proc = None
        # Remove leftover IPC files
        _cleanup_ipc_dir()

    # -- launch -------------------------------------------------------------

    @classmethod
    def launch(cls) -> "BpyClient":
        """Start the bpy worker subprocess and wait until it is ready.

        Returns a ``BpyClient`` instance ready to accept requests.
        """
        # Kill any stale bpy worker still running (hung/crashed).
        _kill_stale_bpy_worker()

        # Clean up any stale files from a previous run
        _cleanup_ipc_dir()
        os.makedirs(BPY_IPC_DIR, exist_ok=True)

        # Launch the worker
        here = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
        proc = subprocess.Popen(
            [sys.executable, os.path.join(here, "bpy_server.py")],
            stdout=None,
            stderr=None,
            preexec_fn=os.setsid,
        )
        print(f"[BpyClient] bpy worker started (pid={proc.pid})", flush=True)

        client = cls(proc)

        # Register cleanup on normal exit
        atexit.register(client.shutdown)

        # Wait for the worker to signal readiness
        ready_flag = os.path.join(BPY_IPC_DIR, ".ready")
        t0 = time.time()
        last_log = 0.0
        while not os.path.exists(ready_flag):
            if time.time() - t0 > BPY_IPC_STARTUP_TIMEOUT:
                client.shutdown()
                raise RuntimeError(
                    f"bpy worker failed to start within {BPY_IPC_STARTUP_TIMEOUT}s"
                )
            if proc.poll() is not None:
                raise RuntimeError(
                    f"bpy worker exited during startup (code={proc.returncode})"
                )
            now = time.time()
            if now - last_log > 10:
                print(f"[BpyClient] still waiting for bpy worker ({now - t0:.0f}s elapsed)", flush=True)
                last_log = now
            time.sleep(0.5)

        print(f"[BpyClient] bpy worker is ready (after {time.time() - t0:.1f}s)", flush=True)
        return client