Qwen3-VL/sim_scripts/sim_worker.py

"""Persistent simulation worker process for GRPO reward server.

Each worker binds to one GPU, initializes warp + GarmentCode once,
then loops on a task queue.  Avoids per-request subprocess cold-start.
"""

from __future__ import annotations

import copy
import io
import multiprocessing as mp
import os
import sys
import time
import traceback
import uuid
from contextlib import redirect_stderr, redirect_stdout
from pathlib import Path
from typing import Any, Optional

_SCRIPT_DIR = Path(__file__).resolve().parent
_PROJECT_ROOT = _SCRIPT_DIR.parent


def _worker_loop(
    gpu_id: int,
    garmentcode_root: str,
    sim_config_path: str,
    task_queue: mp.Queue,
    result_map: dict,
    result_lock: mp.Lock,
    result_event_map: dict,
):
    """Main loop for one persistent sim worker.

    Runs in a child process with CUDA_VISIBLE_DEVICES pinned.
    """
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)

    gc_root = Path(garmentcode_root).resolve()
    os.chdir(str(gc_root))
    if str(gc_root) not in sys.path:
        sys.path.insert(0, str(gc_root))

    import yaml
    from pygarment.meshgen.boxmeshgen import BoxMesh
    from pygarment.meshgen.simulation import run_sim
    import pygarment.data_config as data_config
    from pygarment.meshgen.sim_config import PathCofig

    sim_props_template = data_config.Properties(str(gc_root / sim_config_path))
    sim_props_template_dict = copy.deepcopy(sim_props_template.properties)
    sim_res_scale = sim_props_template_dict["sim"]["config"]["resolution_scale"]
    sim_uv_config = copy.deepcopy(sim_props_template_dict["render"]["config"]["uv_texture"])

    print(f"[SimWorker GPU={gpu_id}] Initialized, waiting for tasks...", flush=True)

    while True:
        try:
            item = task_queue.get()
            if item is None:
                break

            task_id = item["task_id"]
            spec_json_path = item["spec_json_path"]
            out_dir = Path(item["out_dir"])
            sim_timeout_s = int(item.get("sim_timeout_s", 120))
            t0 = time.perf_counter()

            spec_path = Path(spec_json_path)
            garment_name, _, _ = spec_path.stem.rpartition("_")
            if not garment_name:
                garment_name = spec_path.stem

            # Some GarmentCode visualization paths expect this file to exist.
            # Create a tiny placeholder to avoid repeated warning logs.
            design_params_path = out_dir / "design_params.yaml"
            if not design_params_path.exists():
                design_params_path.write_text("{}\n", encoding="utf-8")

            props = data_config.Properties()
            props.properties = copy.deepcopy(sim_props_template_dict)
            props.properties_on_load = copy.deepcopy(sim_props_template_dict)
            props.set_section_stats(
                "sim", fails={}, sim_time={}, spf={},
                fin_frame={}, body_collisions={}, self_collisions={},
            )
            props.set_section_stats("render", render_time={})
            # Fail fast for long-tail simulations.
            if sim_timeout_s > 0:
                cur_timeout = int(props["sim"]["config"].get("max_sim_time", sim_timeout_s))
                props["sim"]["config"]["max_sim_time"] = min(cur_timeout, sim_timeout_s)

            out_name = "sim_result"
            paths = PathCofig(
                in_element_path=spec_path.parent,
                out_path=str(out_dir),
                in_name=garment_name,
                out_name=out_name,
                body_name="mean_all",
                smpl_body=False,
                add_timestamp=False,
            )

            box_mesh = BoxMesh(str(spec_path), sim_res_scale)
            box_mesh.load()
            box_mesh.serialize(
                paths, store_panels=False,
                uv_config=sim_uv_config,
            )
            props.serialize(paths.element_sim_props)

            # GarmentCode still prints per-frame progress even with verbose=False.
            # Swallow worker stdout/stderr to avoid severe multi-process log I/O slowdown.
            with redirect_stdout(io.StringIO()), redirect_stderr(io.StringIO()):
                run_sim(
                    box_mesh.name,
                    props,
                    paths,
                    save_v_norms=False,
                    store_usd=False,
                    optimize_storage=False,
                    verbose=False,
                )
            props.serialize(paths.element_sim_props)

            rd = out_dir / out_name
            front = rd / f"{out_name}_render_front.png"
            back = rd / f"{out_name}_render_back.png"
            sim_props_file = rd / "sim_props.yaml"
            sim_mesh = rd / f"{out_name}_sim.obj"

            render_ok = front.exists() and back.exists()
            physics_ok = False
            if sim_props_file.exists():
                try:
                    sp = yaml.safe_load(sim_props_file.read_text())
                    fails = sp["sim"]["stats"]["fails"]
                    fatal = (
                        "crashes", "frame_timeout", "simulation_timeout",
                        "cloth_body_intersection", "cloth_self_intersection",
                        "static_equilibrium",
                    )
                    physics_ok = not any(fails.get(k, []) for k in fatal)
                except Exception:
                    physics_ok = False

            result = {
                "success": render_ok and physics_ok,
                "sim_mesh": str(sim_mesh) if sim_mesh.exists() else None,
                "sim_time": round(time.perf_counter() - t0, 4),
                "gpu_id": gpu_id,
                "error": None,
            }
        except Exception as e:
            result = {
                "success": False,
                "sim_mesh": None,
                "sim_time": 0,
                "gpu_id": gpu_id,
                "error": f"{type(e).__name__}: {e}\n{traceback.format_exc()}",
            }

        with result_lock:
            # Task may already be timed out on server side; drop late results.
            if task_id in result_event_map:
                result_map[task_id] = result
                result_event_map[task_id].set()


class SimWorkerPool:
    """Pool of persistent simulation workers with dynamic scheduling."""

    def __init__(
        self,
        gpus: list[int],
        garmentcode_root: str,
        sim_config: str,
        sim_timeout_s: int = 120,
    ):
        self.gpus = gpus
        self.sim_timeout_s = sim_timeout_s
        self.manager = mp.Manager()
        self.result_map = self.manager.dict()
        self.result_lock = mp.Lock()
        self.result_event_map = self.manager.dict()
        self.task_queue: mp.Queue = mp.Queue()
        self.workers: dict[int, mp.Process] = {}

        for gpu_id in gpus:
            p = mp.Process(
                target=_worker_loop,
                args=(
                    gpu_id, garmentcode_root, sim_config,
                    self.task_queue, self.result_map, self.result_lock,
                    self.result_event_map,
                ),
                daemon=True,
            )
            p.start()
            self.workers[gpu_id] = p

    def submit(
        self,
        spec_json_path: str,
        out_dir: str,
    ) -> str:
        """Submit a simulation task. Returns task_id."""
        task_id = uuid.uuid4().hex[:12]
        evt = self.manager.Event()
        with self.result_lock:
            self.result_event_map[task_id] = evt
        self.task_queue.put({
            "task_id": task_id,
            "spec_json_path": spec_json_path,
            "out_dir": out_dir,
            "sim_timeout_s": self.sim_timeout_s,
        })
        return task_id

    def wait(self, task_id: str, timeout: float = 600) -> dict:
        """Block until task completes. Returns result dict."""
        evt = self.result_event_map.get(task_id)
        if evt is not None:
            evt.wait(timeout=timeout)
        with self.result_lock:
            result = dict(self.result_map.pop(task_id, {
                "success": False, "sim_mesh": None, "sim_time": 0,
                "gpu_id": None,
                "error": "timeout",
            }))
            self.result_event_map.pop(task_id, None)
        return result

    def submit_and_wait(
        self,
        spec_json_path: str,
        out_dir: str,
        timeout: float = 600,
    ) -> dict:
        """Submit and block for result."""
        tid = self.submit(spec_json_path, out_dir)
        return self.wait(tid, timeout)

    def shutdown(self):
        for _ in self.workers.values():
            self.task_queue.put(None)
        for p in self.workers.values():
            p.join(timeout=10)