Update Wall-X to 1.1.0 (#104)
This commit is contained in:
Starrick Liu
@@ -0,0 +1,197 @@
|
||||
"""In-process driver: DirectModelHandle + N x SubprocEnvHandle.
|
||||
|
||||
The model lives in the driver process while simulator envs live in subprocesses.
|
||||
Evaluation proceeds in chunk-level lockstep:
|
||||
1. Seed the driver process and construct the model handle.
|
||||
2. Start one subprocess env handle per worker.
|
||||
3. Claim a task-local frame from JobState.
|
||||
4. Reset envs, batch active observations, run model.predict_batch, then
|
||||
fan out action chunks to env subprocesses.
|
||||
5. Complete all episodes in the frame and move to the next frame.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
|
||||
from wall_x._vendor.harrix.eval_config import EvalConfig
|
||||
from wall_x._vendor.harrix.drivers.inproc.env_handle import SubprocEnvHandle
|
||||
from wall_x._vendor.harrix.drivers.inproc.model_handle import DirectModelHandle
|
||||
from wall_x._vendor.harrix.drivers.job_state import JobState
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def run(cfg: EvalConfig) -> None:
|
||||
import wall_x._vendor.harrix.envs # noqa: F401 trigger env register
|
||||
import wall_x._vendor.harrix.adapters # noqa: F401 trigger adapter register
|
||||
from wall_x._vendor.harrix.envs.registry import enumerate_episodes_for
|
||||
from wall_x._vendor.harrix.utils.seed import set_seed_everywhere
|
||||
|
||||
# 1) Seed the driver process.
|
||||
set_seed_everywhere(cfg.env.seed)
|
||||
|
||||
# 2) Build DirectModelHandle; the model is loaded in the driver process.
|
||||
logger.info("Constructing DirectModelHandle in the driver process")
|
||||
t_model = time.time()
|
||||
model_handle = DirectModelHandle(cfg)
|
||||
logger.info("Model loaded in %.1fs", time.time() - t_model)
|
||||
|
||||
# 3) Start env subprocesses.
|
||||
logger.info("Starting %s SubprocEnvHandle(s)", cfg.runtime.num_workers)
|
||||
env_handles = [
|
||||
SubprocEnvHandle(cfg, worker_id=i) for i in range(cfg.runtime.num_workers)
|
||||
]
|
||||
|
||||
# 4) JobState runs in frame-sync mode for lockstep evaluation.
|
||||
os.makedirs(cfg.runtime.log_dir, exist_ok=True)
|
||||
log_path = os.path.join(cfg.runtime.log_dir, "state.jsonl")
|
||||
report_path = os.path.join(cfg.runtime.log_dir, "report.json")
|
||||
episodes = enumerate_episodes_for(cfg)
|
||||
logger.info("env.type=%r; scheduled %s episodes", cfg.env.type, len(episodes))
|
||||
state = JobState(
|
||||
episodes,
|
||||
log_path,
|
||||
batch_sync_mode=True,
|
||||
batch_size=cfg.runtime.num_workers,
|
||||
)
|
||||
|
||||
# 5) Main loop, one frame at a time.
|
||||
t0 = time.time()
|
||||
frame_idx = 0
|
||||
while not state.is_drained():
|
||||
frame_eps = state.claim_frame()
|
||||
if not frame_eps:
|
||||
# Previous frame still has in-flight episodes.
|
||||
time.sleep(0.05)
|
||||
continue
|
||||
|
||||
results = _run_frame(frame_eps, model_handle, env_handles, cfg, frame_idx)
|
||||
for ep, res in zip(frame_eps, results):
|
||||
if "_error" in res:
|
||||
state.fail(ep, str(res["_error"]))
|
||||
else:
|
||||
state.complete(ep, res)
|
||||
frame_idx += 1
|
||||
|
||||
elapsed = time.time() - t0
|
||||
logger.info("All episodes finished in %.1fs (%.1f min)", elapsed, elapsed / 60)
|
||||
|
||||
state.dump_final(report_path)
|
||||
with open(report_path) as f:
|
||||
report = json.load(f)
|
||||
overall = report["overall"]
|
||||
logger.info(
|
||||
"attempted=%s, successes=%s, success_rate=%.2f%%, failed=%s",
|
||||
overall["attempted"],
|
||||
overall["successes"],
|
||||
overall["success_rate"] * 100,
|
||||
overall["failed"],
|
||||
)
|
||||
logger.info("Report: %s", report_path)
|
||||
logger.info("State log: %s", log_path)
|
||||
|
||||
for h in env_handles:
|
||||
h.shutdown()
|
||||
model_handle.shutdown()
|
||||
|
||||
|
||||
def _run_frame(
|
||||
frame_eps: list,
|
||||
model_handle: DirectModelHandle,
|
||||
env_handles: list,
|
||||
cfg: EvalConfig,
|
||||
frame_idx: int,
|
||||
) -> list[dict]:
|
||||
"""Run one task-local frame with chunk-level lockstep.
|
||||
|
||||
Active workers are batched together at each chunk boundary. Workers that
|
||||
already finished no longer participate in later forwards.
|
||||
"""
|
||||
from wall_x._vendor.harrix.envs.libero_common import encode_raw_obs
|
||||
|
||||
n = len(frame_eps)
|
||||
t_frame_start = time.time()
|
||||
|
||||
# ---- a) reset: fan out, then gather ----
|
||||
for i in range(n):
|
||||
env_handles[i].submit_reset(tuple(frame_eps[i]))
|
||||
initials = [env_handles[i].wait_reset() for i in range(n)]
|
||||
|
||||
obs_list = [r["obs"] for r in initials]
|
||||
instr_list = [r["instruction"] for r in initials]
|
||||
status = [
|
||||
{
|
||||
"done": False,
|
||||
"success": False,
|
||||
"steps": 0,
|
||||
"task_desc": initials[i].get("task_desc", ""),
|
||||
}
|
||||
for i in range(n)
|
||||
]
|
||||
|
||||
max_rounds = cfg.env.libero.max_infer_times
|
||||
|
||||
# ---- b) chunk lockstep ----
|
||||
for round_idx in range(max_rounds):
|
||||
active = [i for i in range(n) if not status[i]["done"]]
|
||||
if not active:
|
||||
break
|
||||
|
||||
payloads = [
|
||||
{
|
||||
"observation": encode_raw_obs(obs_list[i]),
|
||||
"instruction": instr_list[i],
|
||||
"noise": None,
|
||||
}
|
||||
for i in active
|
||||
]
|
||||
chunks = model_handle.predict_batch(payloads)
|
||||
|
||||
# fan-out submit
|
||||
for k, i in enumerate(active):
|
||||
env_handles[i].submit_execute_chunk(chunks[k])
|
||||
# gather
|
||||
for k, i in enumerate(active):
|
||||
try:
|
||||
r = env_handles[i].wait_execute_chunk()
|
||||
except Exception as e:
|
||||
status[i]["_error"] = str(e)
|
||||
status[i]["done"] = True
|
||||
continue
|
||||
obs_list[i] = r["obs"]
|
||||
status[i]["steps"] += r["steps"]
|
||||
if r["done"]:
|
||||
status[i]["done"] = True
|
||||
status[i]["success"] = True
|
||||
|
||||
for i in range(n):
|
||||
env_handles[i].submit_finalize_episode(status[i]["success"])
|
||||
for i in range(n):
|
||||
env_handles[i].wait_finalize_episode()
|
||||
|
||||
elapsed_frame = time.time() - t_frame_start
|
||||
logger.info(
|
||||
"frame=%s n=%s succ=%s/%s elapsed=%.1fs",
|
||||
frame_idx,
|
||||
n,
|
||||
sum(s["success"] for s in status),
|
||||
n,
|
||||
elapsed_frame,
|
||||
)
|
||||
|
||||
return [
|
||||
{
|
||||
"success": s["success"],
|
||||
"steps": s["steps"],
|
||||
"elapsed_sec": round(elapsed_frame / max(1, n), 3),
|
||||
"task_desc": s["task_desc"],
|
||||
**({"_error": s["_error"]} if "_error" in s else {}),
|
||||
}
|
||||
for s in status
|
||||
]
|
||||
@@ -0,0 +1,152 @@
|
||||
"""Subprocess env handle used by the in-process driver.
|
||||
|
||||
The model remains in the driver process. Each env subprocess receives reset and
|
||||
execute-chunk commands through a pipe.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def _subproc_main(cfg, worker_id, child_conn):
|
||||
"""Env subprocess entry point."""
|
||||
# Spawned children inherit env vars, but set these explicitly for launchers
|
||||
# that did not configure them. Robosuite validates the EGL id against the
|
||||
# CUDA_VISIBLE_DEVICES environment string, so keep the same visible id here.
|
||||
cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") or "0"
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices
|
||||
# EGL device index is always 0 after CUDA_VISIBLE_DEVICES remapping.
|
||||
os.environ["MUJOCO_EGL_DEVICE_ID"] = os.environ.get("MUJOCO_EGL_DEVICE_ID") or "0"
|
||||
|
||||
# Env workers run robosuite only. Do not import torch here; otherwise many
|
||||
# subprocesses may initialize CUDA contexts and compete with the driver model.
|
||||
seed = cfg.env.seed + worker_id
|
||||
np.random.seed(seed)
|
||||
random.seed(seed)
|
||||
os.environ["PYTHONHASHSEED"] = str(seed)
|
||||
|
||||
import wall_x._vendor.harrix.envs # noqa: F401 trigger register
|
||||
from wall_x._vendor.harrix.envs.registry import build_env
|
||||
|
||||
env = build_env(cfg, worker_id)
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
cmd, args = child_conn.recv()
|
||||
except EOFError:
|
||||
break
|
||||
try:
|
||||
if cmd == "reset_episode":
|
||||
result = env.reset_episode(tuple(args))
|
||||
elif cmd == "execute_chunk":
|
||||
result = env.execute_chunk(args)
|
||||
elif cmd == "finalize_episode":
|
||||
env.finalize_episode(bool(args))
|
||||
result = None
|
||||
elif cmd == "shutdown":
|
||||
env.shutdown()
|
||||
child_conn.send(("ok", None))
|
||||
break
|
||||
else:
|
||||
raise ValueError(f"unknown cmd {cmd!r}")
|
||||
child_conn.send(("ok", result))
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
child_conn.send(("err", f"{e}\n{traceback.format_exc()}"))
|
||||
finally:
|
||||
try:
|
||||
child_conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
class SubprocEnvHandle:
|
||||
"""Synchronous pipe wrapper around one env subprocess.
|
||||
|
||||
Usage:
|
||||
h.submit_reset(ep_id)
|
||||
...
|
||||
h.wait_reset() # -> {"obs", "instruction", "task_desc"}
|
||||
h.submit_execute_chunk(actions)
|
||||
...
|
||||
h.wait_execute_chunk() # -> {"obs", "done", "steps"}
|
||||
"""
|
||||
|
||||
def __init__(self, cfg, worker_id: int):
|
||||
# Use spawn instead of fork because the driver may already hold a CUDA
|
||||
# context for the model.
|
||||
ctx = mp.get_context("spawn")
|
||||
self._parent_conn, child_conn = ctx.Pipe()
|
||||
self._proc = ctx.Process(
|
||||
target=_subproc_main,
|
||||
args=(cfg, worker_id, child_conn),
|
||||
daemon=False,
|
||||
name=f"infer-subproc-w{worker_id}",
|
||||
)
|
||||
self._proc.start()
|
||||
child_conn.close()
|
||||
self._has_pending = False
|
||||
|
||||
def _send(self, cmd: str, args):
|
||||
self._parent_conn.send((cmd, args))
|
||||
self._has_pending = True
|
||||
|
||||
def _recv(self):
|
||||
if not self._has_pending:
|
||||
raise RuntimeError("no pending request to wait for")
|
||||
status, payload = self._parent_conn.recv()
|
||||
self._has_pending = False
|
||||
if status == "err":
|
||||
raise RuntimeError(f"subproc env error (w={self._proc.name}): {payload}")
|
||||
return payload
|
||||
|
||||
def submit_reset(self, ep_id):
|
||||
self._send("reset_episode", ep_id)
|
||||
|
||||
def wait_reset(self):
|
||||
return self._recv()
|
||||
|
||||
def submit_execute_chunk(self, actions):
|
||||
self._send("execute_chunk", np.asarray(actions, dtype=np.float32))
|
||||
|
||||
def wait_execute_chunk(self):
|
||||
return self._recv()
|
||||
|
||||
def submit_finalize_episode(self, success: bool):
|
||||
self._send("finalize_episode", success)
|
||||
|
||||
def wait_finalize_episode(self):
|
||||
return self._recv()
|
||||
|
||||
def finalize_episode(self, success: bool):
|
||||
"""Blocking convenience wrapper."""
|
||||
self.submit_finalize_episode(success)
|
||||
return self.wait_finalize_episode()
|
||||
|
||||
def reset_episode(self, ep_id):
|
||||
"""Blocking convenience wrapper."""
|
||||
self.submit_reset(ep_id)
|
||||
return self.wait_reset()
|
||||
|
||||
def execute_chunk(self, actions):
|
||||
"""Blocking convenience wrapper."""
|
||||
self.submit_execute_chunk(actions)
|
||||
return self.wait_execute_chunk()
|
||||
|
||||
def shutdown(self) -> None:
|
||||
try:
|
||||
self._send("shutdown", None)
|
||||
self._recv()
|
||||
except Exception:
|
||||
pass
|
||||
if self._proc.is_alive():
|
||||
self._proc.join(timeout=5)
|
||||
if self._proc.is_alive():
|
||||
self._proc.terminate()
|
||||
self._proc.join(timeout=2)
|
||||
@@ -0,0 +1,25 @@
|
||||
"""In-process model handle.
|
||||
|
||||
The model is constructed in the driver process and calls the adapter directly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
class DirectModelHandle:
|
||||
def __init__(self, cfg):
|
||||
# Trigger adapter registration.
|
||||
import wall_x._vendor.harrix.adapters # noqa: F401
|
||||
from wall_x._vendor.harrix.adapters.registry import build_adapter
|
||||
|
||||
self._adapter = build_adapter(cfg)
|
||||
|
||||
@property
|
||||
def chunk_horizon(self) -> int:
|
||||
return self._adapter.chunk_horizon
|
||||
|
||||
def predict_batch(self, payloads):
|
||||
return self._adapter.predict_batch(payloads)
|
||||
|
||||
def shutdown(self) -> None:
|
||||
pass
|
||||
@@ -0,0 +1,343 @@
|
||||
"""Episode work queue with JSONL persistence.
|
||||
|
||||
Drivers enumerate episode ids into ``pending``; env handles claim work and
|
||||
return results through ``complete`` or ``fail``. Every state transition is
|
||||
appended to ``state.jsonl``, so completed episodes can be skipped on restart.
|
||||
|
||||
Two scheduling modes are supported:
|
||||
- FIFO: workers claim the next pending episode.
|
||||
- batch_sync: episodes are grouped into task-local frames. A new frame is not
|
||||
released until the previous frame is complete, which gives deterministic
|
||||
lockstep batches at the cost of possible idle workers.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class JobState:
|
||||
def __init__(
|
||||
self,
|
||||
all_episodes: list[tuple],
|
||||
log_path: str,
|
||||
batch_sync_mode: bool = False,
|
||||
batch_size: int = 1,
|
||||
):
|
||||
"""
|
||||
all_episodes: list of (suite_name, task_idx, ep_idx) tuples
|
||||
log_path: state.jsonl path; existing completed episodes are skipped
|
||||
batch_sync_mode: enable task-local frame barriers
|
||||
batch_size: number of episodes per frame
|
||||
"""
|
||||
self._lock = threading.Lock()
|
||||
self._log_path = log_path
|
||||
self._log_fh = None
|
||||
self._batch_sync_mode = bool(batch_sync_mode)
|
||||
self._batch_size = int(batch_size)
|
||||
|
||||
completed_set = (
|
||||
self._load_completed(log_path) if os.path.exists(log_path) else set()
|
||||
)
|
||||
remaining: list[tuple] = [
|
||||
tuple(ep) for ep in all_episodes if tuple(ep) not in completed_set
|
||||
]
|
||||
|
||||
self._in_progress: dict[tuple, int] = {}
|
||||
self._completed: dict[tuple, dict] = {}
|
||||
self._failed: dict[tuple, str] = {}
|
||||
|
||||
if self._batch_sync_mode:
|
||||
self._frames: list[dict] = self._build_frames(remaining, self._batch_size)
|
||||
self._cur_frame_idx: int = 0
|
||||
self._cur_frame_inflight: int = 0
|
||||
self._pending = None
|
||||
else:
|
||||
self._frames = []
|
||||
self._cur_frame_idx = 0
|
||||
self._cur_frame_inflight = 0
|
||||
self._pending: list[tuple] = remaining
|
||||
|
||||
os.makedirs(os.path.dirname(log_path) or ".", exist_ok=True)
|
||||
self._log_fh = open(log_path, "a")
|
||||
|
||||
log_entry = {
|
||||
"event": "session_start",
|
||||
"pending": (
|
||||
len(self._pending)
|
||||
if not self._batch_sync_mode
|
||||
else sum(len(f["eps"]) for f in self._frames)
|
||||
),
|
||||
"skipped_completed": len(completed_set),
|
||||
"batch_sync_mode": self._batch_sync_mode,
|
||||
}
|
||||
if self._batch_sync_mode:
|
||||
log_entry["num_frames"] = len(self._frames)
|
||||
log_entry["batch_size"] = self._batch_size
|
||||
self._append_log(log_entry)
|
||||
|
||||
@staticmethod
|
||||
def _build_frames(remaining: list[tuple], batch_size: int) -> list[dict]:
|
||||
"""Group episodes by ``(suite, task_idx)`` into fixed-size frames.
|
||||
|
||||
Partial tail frames are kept.
|
||||
"""
|
||||
frames: list[dict] = []
|
||||
cur_key: Optional[tuple] = None
|
||||
cur_bucket: list[tuple] = []
|
||||
for ep in remaining:
|
||||
key = (ep[0], ep[1])
|
||||
if key != cur_key and cur_bucket:
|
||||
for i in range(0, len(cur_bucket), batch_size):
|
||||
frames.append(
|
||||
{"eps": cur_bucket[i : i + batch_size], "next_slot": 0}
|
||||
)
|
||||
cur_bucket = []
|
||||
cur_key = key
|
||||
cur_bucket.append(ep)
|
||||
if cur_bucket:
|
||||
for i in range(0, len(cur_bucket), batch_size):
|
||||
frames.append({"eps": cur_bucket[i : i + batch_size], "next_slot": 0})
|
||||
return frames
|
||||
|
||||
@staticmethod
|
||||
def _load_completed(log_path: str) -> set[tuple]:
|
||||
completed = set()
|
||||
with open(log_path, "r") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
rec = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if rec.get("event") == "completed" and "ep_id" in rec:
|
||||
completed.add(tuple(rec["ep_id"]))
|
||||
return completed
|
||||
|
||||
def _append_log(self, extra: dict):
|
||||
rec = {"ts": time.strftime("%Y-%m-%dT%H:%M:%S"), **extra}
|
||||
if "ep_id" in rec and isinstance(rec["ep_id"], tuple):
|
||||
rec["ep_id"] = list(rec["ep_id"])
|
||||
self._log_fh.write(json.dumps(rec, ensure_ascii=False) + "\n")
|
||||
self._log_fh.flush()
|
||||
|
||||
# ---- claim API ----
|
||||
|
||||
def claim(self, worker_id: int) -> Optional[list]:
|
||||
"""Claim one episode id.
|
||||
|
||||
FIFO returns ``None`` when the pending queue is empty. In batch-sync
|
||||
mode, ``None`` can also mean the current frame has been fully issued but
|
||||
still has in-flight episodes.
|
||||
"""
|
||||
with self._lock:
|
||||
if self._batch_sync_mode:
|
||||
return self._claim_sync(worker_id)
|
||||
return self._claim_fifo(worker_id)
|
||||
|
||||
def claim_frame(self) -> list[list]:
|
||||
"""Claim one complete frame for lockstep in-process evaluation.
|
||||
|
||||
In FIFO mode this returns up to ``batch_size`` pending episodes.
|
||||
"""
|
||||
with self._lock:
|
||||
if self._batch_sync_mode:
|
||||
return self._claim_frame_sync()
|
||||
return self._claim_frame_fifo()
|
||||
|
||||
def _claim_fifo(self, worker_id: int) -> Optional[list]:
|
||||
if not self._pending:
|
||||
return None
|
||||
ep = self._pending.pop(0)
|
||||
self._in_progress[ep] = worker_id
|
||||
self._append_log({"event": "claimed", "ep_id": ep, "worker": worker_id})
|
||||
return list(ep)
|
||||
|
||||
def _claim_sync(self, worker_id: int) -> Optional[list]:
|
||||
while self._cur_frame_idx < len(self._frames):
|
||||
frame = self._frames[self._cur_frame_idx]
|
||||
if frame["next_slot"] < len(frame["eps"]):
|
||||
ep = frame["eps"][frame["next_slot"]]
|
||||
frame["next_slot"] += 1
|
||||
self._cur_frame_inflight += 1
|
||||
self._in_progress[ep] = worker_id
|
||||
self._append_log(
|
||||
{
|
||||
"event": "claimed",
|
||||
"ep_id": ep,
|
||||
"worker": worker_id,
|
||||
"frame": self._cur_frame_idx,
|
||||
}
|
||||
)
|
||||
return list(ep)
|
||||
if self._cur_frame_inflight > 0:
|
||||
return None
|
||||
self._append_log(
|
||||
{
|
||||
"event": "frame_done",
|
||||
"frame": self._cur_frame_idx,
|
||||
"size": len(frame["eps"]),
|
||||
}
|
||||
)
|
||||
self._cur_frame_idx += 1
|
||||
return None
|
||||
|
||||
def _claim_frame_sync(self) -> list[list]:
|
||||
# Advance completed frames.
|
||||
while self._cur_frame_idx < len(self._frames):
|
||||
frame = self._frames[self._cur_frame_idx]
|
||||
if frame["next_slot"] < len(frame["eps"]):
|
||||
break
|
||||
if self._cur_frame_inflight > 0:
|
||||
# The previous frame still has in-flight episodes.
|
||||
return []
|
||||
self._append_log(
|
||||
{
|
||||
"event": "frame_done",
|
||||
"frame": self._cur_frame_idx,
|
||||
"size": len(frame["eps"]),
|
||||
}
|
||||
)
|
||||
self._cur_frame_idx += 1
|
||||
if self._cur_frame_idx >= len(self._frames):
|
||||
return []
|
||||
frame = self._frames[self._cur_frame_idx]
|
||||
out: list[list] = []
|
||||
while frame["next_slot"] < len(frame["eps"]):
|
||||
ep = frame["eps"][frame["next_slot"]]
|
||||
frame["next_slot"] += 1
|
||||
self._cur_frame_inflight += 1
|
||||
self._in_progress[ep] = -1
|
||||
self._append_log(
|
||||
{"event": "claimed", "ep_id": ep, "frame": self._cur_frame_idx}
|
||||
)
|
||||
out.append(list(ep))
|
||||
return out
|
||||
|
||||
def _claim_frame_fifo(self) -> list[list]:
|
||||
if not self._pending:
|
||||
return []
|
||||
n = min(self._batch_size, len(self._pending))
|
||||
out: list[list] = []
|
||||
for _ in range(n):
|
||||
ep = self._pending.pop(0)
|
||||
self._in_progress[ep] = -1
|
||||
self._append_log({"event": "claimed", "ep_id": ep})
|
||||
out.append(list(ep))
|
||||
return out
|
||||
|
||||
# ---- complete / fail ----
|
||||
|
||||
def complete(self, ep_id: list, result: dict) -> None:
|
||||
ep = tuple(ep_id)
|
||||
with self._lock:
|
||||
if ep in self._in_progress:
|
||||
self._in_progress.pop(ep, None)
|
||||
if self._batch_sync_mode:
|
||||
self._cur_frame_inflight = max(0, self._cur_frame_inflight - 1)
|
||||
self._completed[ep] = result
|
||||
self._append_log({"event": "completed", "ep_id": ep, **result})
|
||||
|
||||
def fail(self, ep_id: list, error: str) -> None:
|
||||
ep = tuple(ep_id)
|
||||
with self._lock:
|
||||
if ep in self._in_progress:
|
||||
self._in_progress.pop(ep, None)
|
||||
if self._batch_sync_mode:
|
||||
self._cur_frame_inflight = max(0, self._cur_frame_inflight - 1)
|
||||
self._failed[ep] = error
|
||||
self._append_log({"event": "failed", "ep_id": ep, "error": error})
|
||||
|
||||
# ---- status queries ----
|
||||
|
||||
def get_frame_inflight(self) -> int:
|
||||
with self._lock:
|
||||
return self._cur_frame_inflight if self._batch_sync_mode else 0
|
||||
|
||||
def is_drained(self) -> bool:
|
||||
with self._lock:
|
||||
if self._batch_sync_mode:
|
||||
return (
|
||||
self._cur_frame_idx >= len(self._frames)
|
||||
and self._cur_frame_inflight == 0
|
||||
)
|
||||
return len(self._pending) == 0 and len(self._in_progress) == 0
|
||||
|
||||
def progress(self) -> dict:
|
||||
with self._lock:
|
||||
base = {
|
||||
"in_progress": len(self._in_progress),
|
||||
"completed": len(self._completed),
|
||||
"failed": len(self._failed),
|
||||
}
|
||||
if self._batch_sync_mode:
|
||||
pending = sum(
|
||||
len(f["eps"]) - f["next_slot"]
|
||||
for f in self._frames[self._cur_frame_idx :]
|
||||
)
|
||||
base["pending"] = pending
|
||||
base["frame"] = f"{self._cur_frame_idx}/{len(self._frames)}"
|
||||
base["frame_inflight"] = self._cur_frame_inflight
|
||||
else:
|
||||
base["pending"] = len(self._pending)
|
||||
return base
|
||||
|
||||
def dump_final(self, report_path: str) -> None:
|
||||
with self._lock:
|
||||
per_task: dict[tuple, dict] = {}
|
||||
for ep, result in self._completed.items():
|
||||
key = (ep[0], ep[1])
|
||||
d = per_task.setdefault(
|
||||
key, {"attempted": 0, "successes": 0, "steps": []}
|
||||
)
|
||||
d["attempted"] += 1
|
||||
if result.get("success"):
|
||||
d["successes"] += 1
|
||||
if "steps" in result:
|
||||
d["steps"].append(result["steps"])
|
||||
for ep in self._failed:
|
||||
key = (ep[0], ep[1])
|
||||
d = per_task.setdefault(
|
||||
key, {"attempted": 0, "successes": 0, "steps": []}
|
||||
)
|
||||
d["attempted"] += 1
|
||||
|
||||
total_attempted = sum(d["attempted"] for d in per_task.values())
|
||||
total_successes = sum(d["successes"] for d in per_task.values())
|
||||
overall_rate = total_successes / max(1, total_attempted)
|
||||
|
||||
report = {
|
||||
"overall": {
|
||||
"attempted": total_attempted,
|
||||
"successes": total_successes,
|
||||
"success_rate": overall_rate,
|
||||
"failed": len(self._failed),
|
||||
},
|
||||
"per_task": {
|
||||
f"{suite}_t{task_idx}": {
|
||||
**d,
|
||||
"success_rate": d["successes"] / max(1, d["attempted"]),
|
||||
"avg_steps": (
|
||||
(sum(d["steps"]) / max(1, len(d["steps"])))
|
||||
if d["steps"]
|
||||
else None
|
||||
),
|
||||
}
|
||||
for (suite, task_idx), d in sorted(per_task.items())
|
||||
},
|
||||
}
|
||||
with open(report_path, "w") as f:
|
||||
json.dump(report, f, indent=2, ensure_ascii=False)
|
||||
self._append_log(
|
||||
{
|
||||
"event": "session_end",
|
||||
"report_path": report_path,
|
||||
"overall_success_rate": overall_rate,
|
||||
}
|
||||
)
|
||||
Reference in New Issue
Block a user