"""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)