[v5] 复用 perf 与 metric 统一性能观测

src/twinkle/metric/__init__.py

@@ -10,5 +10,5 @@
 from .loss import LossMetric
 from .ppo import PPOValueMetric
 from .reporting import MetricsReporter, create_metrics_reporter
-from .train_metric import TrainMetric
+from .train_metric import ThroughputMetric, TrainMetric, memory_metrics
 from .types import MetricRecord

src/twinkle/metric/reporting.py

@@ -2,8 +2,6 @@
 """Asynchronous JSONL and SwanLab metric reporting."""
 
 from __future__ import annotations
-
-import json
 import logging
 import math
 import os
@@ -16,6 +14,8 @@
 from pathlib import Path
 from typing import Any
 
+import json
+
 from .types import MetricRecord
 
 logger = logging.getLogger(__name__)
@@ -119,6 +119,16 @@ def add(self, record: MetricRecord) -> None:
         if record.stage == 'run':
             self.run_status = record.status
             self.result = {**record.values, **record.attributes}
+            # A producer may report terminal totals instead of individual sample/group events.
+            if record.context_key:
+                for name, counter in (
+                    ('rollout_groups', self.context_rollout_groups),
+                    ('rollout_samples', self.context_rollout_samples),
+                    ('trained_samples', self.context_trained_samples),
+                ):
+                    total = _finite_number(record.values.get(name))
+                    if total is not None:
+                        counter[record.context_key] = int(total)
 
     def as_dict(self, backend_health: Mapping[str, Any]) -> dict[str, Any]:
         wall_time = _finite_number(self.result.get('wall_time_s'))
@@ -325,10 +335,10 @@ def _close_sink(self) -> None:
 
 class _JSONLBackend(_QueuedBackend):
 
-    def __init__(self, path: str | Path, **kwargs: Any):
+    def __init__(self, path: str | Path, *, append: bool = False, **kwargs: Any):
         self.path = Path(path)
         self.path.parent.mkdir(parents=True, exist_ok=True)
-        self._stream = self.path.open('w', encoding='utf-8')
+        self._stream = self.path.open('a' if append else 'w', encoding='utf-8')
         super().__init__('jsonl', **kwargs)
 
     def _write_batch(self, batch: Sequence[dict[str, Any]]) -> None:
@@ -465,7 +475,7 @@ def _normalize(self, record: MetricRecord, sequence: int) -> MetricRecord:
         for name, value in record.values.items():
             number = _finite_number(value)
             if number is None:
-                non_numeric[name] = value
+                non_numeric[name] = None if isinstance(value, float) and not math.isfinite(value) else value
             else:
                 values[name] = number
         attributes = dict(record.attributes)
@@ -523,7 +533,8 @@ def create_metrics_reporter(config: Mapping[str, Any] | None, *, run_id: str) ->
     }
     if bool(jsonl_config.get('enabled', True)):
         try:
-            backends.append(_JSONLBackend(jsonl_config['path'], **backend_defaults))
+            backends.append(_JSONLBackend(
+                jsonl_config['path'], append=bool(jsonl_config.get('append', False)), **backend_defaults))
         except Exception as exc:
             backend_errors['jsonl'] = exc
             logger.warning('JSONL metrics backend could not start: %s', exc)

src/twinkle/metric/train_metric.py

@@ -60,3 +60,119 @@ def calculate(self):
             results['speed'] = f'{speed:.2f} iters/s'
         self.reset()
         return results
+
+def memory_metrics(rows):
+    """Convert per-worker byte counters and reduce peaks by max, host usage by mean."""
+    rows = [rows] if isinstance(rows, dict) else rows
+    fields = {
+        "max_memory_allocated_bytes": "max_memory_allocated_gb",
+        "max_memory_reserved_bytes": "max_memory_reserved_gb",
+        "cpu_memory_used_bytes": "cpu_memory_used_gb",
+    }
+    result = {}
+    for raw, name in fields.items():
+        values = [row[raw] / 1024**3 for row in rows if raw in row]
+        if values:
+            result[f"perf/{name}"] = max(values) if raw.startswith("max_") else sum(values) / len(values)
+    return result
+
+
+class ThroughputMetric(Metric):
+    """Accumulate explicit token counts and phase seconds; calculate ratios of sums.
+
+    The producer defines window boundaries and supplies the physical device budget.
+    No timing, loop state, task waiting or file output occurs inside this metric.
+    """
+
+    def __init__(self, device_mesh=None, process_group=None, **kwargs):
+        super().__init__(device_mesh, process_group, **kwargs)
+        self.reset()
+
+    def reset(self):
+        from collections import Counter
+
+        self.prompts = Counter()
+        self.responses = Counter()
+        self.timings = {}
+        self.num_devices = None
+
+    def accumulate(
+        self, inputs=None, outputs=None, *, prompt_lengths, response_lengths, timings, num_devices, **kwargs
+    ):
+        import math
+
+        prompts, responses = list(prompt_lengths), list(response_lengths)
+        if not prompts or len(prompts) != len(responses):
+            raise ValueError("expected a nonempty batch of paired prompt/response lengths")
+        if any(not isinstance(n, int) or isinstance(n, bool) or n < 0 for n in prompts + responses):
+            raise ValueError("token lengths must be nonnegative integers")
+        if not isinstance(num_devices, int) or isinstance(num_devices, bool) or num_devices < 1:
+            raise ValueError("num_devices must count allocated physical devices")
+        if self.num_devices is not None and self.num_devices != num_devices:
+            raise ValueError("reset the metric when the physical allocation changes")
+        if timings.get("step", 0) <= 0 or any(not math.isfinite(t) or t < 0 for t in timings.values()):
+            raise ValueError("timings must be finite and nonnegative, with positive step time")
+        self.num_devices = num_devices
+        self.prompts.update(prompts)
+        self.responses.update(responses)
+        for name, seconds in timings.items():
+            self.timings[name] = self.timings.get(name, 0.0) + seconds
+
+    def calculate(self, *, prompt_width=None, response_width=None):
+        if not self.prompts:
+            return {}
+        prompts, responses = self.prompts, self.responses
+        count = sum(prompts.values())
+        p = sum(length * frequency for length, frequency in prompts.items())
+        r = sum(length * frequency for length, frequency in responses.items())
+        total, elapsed = p + r, self.timings["step"]
+        result = {
+            "perf/total_num_tokens": total,
+            "perf/num_prompt_tokens": p,
+            "perf/num_response_tokens": r,
+            "perf/num_samples": count,
+            "perf/num_gpus": self.num_devices,
+            "perf/time_per_step": elapsed,
+            "perf/throughput": total / elapsed / self.num_devices,
+            "perf/throughput_total": total / elapsed,
+            "perf/response_tokens_per_second": r / elapsed,
+            "perf/samples_per_second": count / elapsed,
+            **{f"timing_s/{name}": seconds for name, seconds in self.timings.items()},
+        }
+        denominators = {
+            "gen": r,
+            **dict.fromkeys(("ref", "values", "adv", "adv_pipeline", "update_critic", "update_actor"), total),
+        }
+        for name, tokens in denominators.items():
+            if name in self.timings:
+                result[f"timing_per_token_ms/{name}"] = self.timings[name] * 1000 / tokens if tokens else 0.0
+        for name, lengths, width in (
+            ("prompt_length", prompts, prompt_width),
+            ("response_length", responses, response_width),
+        ):
+            width = max(lengths) if width is None else width
+            if width < max(lengths):
+                raise ValueError("padded width cannot be smaller than an actual sequence")
+            result.update(
+                {
+                    f"{name}/mean": sum(n * frequency for n, frequency in lengths.items()) / count,
+                    f"{name}/min": min(lengths),
+                    f"{name}/max": max(lengths),
+                    f"{name}/clip_ratio": lengths.get(width, 0) / count,
+                }
+            )
+        active = {n: frequency for n, frequency in responses.items() if n > 0}
+        active_count = sum(active.values())
+        width = max(responses) if response_width is None else response_width
+        result.update(
+            {
+                "response/aborted_ratio": 1 - active_count / count,
+                "response_length_non_aborted/mean": r / active_count if active_count else float("nan"),
+                "response_length_non_aborted/min": min(active) if active else float("nan"),
+                "response_length_non_aborted/max": max(active) if active else float("nan"),
+                "response_length_non_aborted/clip_ratio": active.get(width, 0) / active_count
+                if active_count
+                else float("nan"),
+            }
+        )
+        return result

src/twinkle/model/base.py

@@ -297,6 +297,13 @@ def apply_patch(self, patch_cls: Union[Patch, Type[Patch], str], **kwargs) -> No
     def add_metric(self, metric_cls: Union[Metric, str], is_training: Optional[bool] = None, **kwargs) -> None:
         ...
 
+    @remote_function(dispatch='all', collect='none', lazy_collect=False)
+    def memory_stats(self, *, synchronize: bool = False) -> dict[str, int]:
+        """Raw per-rank allocator peaks and host RAM in bytes; no cross-rank reduction."""
+        from twinkle.perf.stats import memory_stats
+
+        return memory_stats(synchronize=synchronize)
+
     @abstractmethod
     def calculate_metric(self, is_training: bool, **kwargs) -> Dict[str, Any]:
         ...

src/twinkle/perf/__init__.py

@@ -37,7 +37,7 @@
 from .counters import CounterRegistry, default_counters
 from .module_profiler import ModuleProfiler, activation_stats, build_module_profiler, param_stats
 from .recorder import TimeRecorder, default_recorder, disable_perf, enable_perf, span, timed
-from .stats import describe_shape, primary_tensor, tensor_moments, tensor_numel
+from .stats import describe_shape, memory_stats, primary_tensor, tensor_moments, tensor_numel
 from .timeline import chrome_trace, dag_report
 
 __all__ = [
@@ -60,4 +60,5 @@
     'tensor_moments',
     'tensor_numel',
     'describe_shape',
+    'memory_stats',
 ]

src/twinkle/perf/recorder.py

@@ -57,13 +57,16 @@ def __init__(self, name: str, parent: Optional['_Span'], start_ns: int, tid: int
 class TimeRecorder:
     """Aggregate per-name span timings into a summary and a metric-ready dict.
 
+    ``clock_ns`` optionally supplies a monotonic nanosecond clock for replay/tests.
+
     Enable/disable is checked at CALL time (not decoration time), so decorating a
     function once and toggling the recorder later behaves correctly, and a disabled
     recorder takes a no-op fast path with no timing or allocation.
     """
 
-    def __init__(self, *, enabled: bool = True) -> None:
+    def __init__(self, *, enabled: bool = True, clock_ns: Callable[[], int] | None = None) -> None:
         self._enabled = enabled
+        self._clock_ns = clock_ns or time.perf_counter_ns
         self._lock = threading.Lock()
         # name -> [count, total_ns, self_ns]
         self._stats: Dict[str, list] = {}
@@ -133,15 +136,15 @@ def _span(self, name: str) -> Generator[None, None, None]:
         partial duration); the exception is NOT suppressed.
         """
         parent = self._current.get()
-        current = _Span(name, parent, time.perf_counter_ns(), threading.get_ident())
+        current = _Span(name, parent, self._clock_ns(), threading.get_ident())
         token = self._current.set(current)
         try:
             yield
         finally:
             # Reset BEFORE recording so a raising body cannot leave a stale span on
             # the chain for subsequent code in this context.
             self._current.reset(token)
-            self._record(current, time.perf_counter_ns())
+            self._record(current, self._clock_ns())
 
     def _record(self, current: _Span, end_ns: int) -> None:
         total_ns = end_ns - current.start_ns

src/twinkle/perf/stats.py

@@ -1,7 +1,7 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
-"""Tensor statistics helpers for the module profiler.
+"""Raw tensor and process observations for performance instrumentation.
 
-These are plain functions over ``torch.Tensor`` (and the nested containers a
+Tensor helpers are plain functions over ``torch.Tensor`` (and nested containers a
 forward may return) so both the built-in :func:`activation_stats` /
 :func:`param_stats` collectors and user collectors can reuse them. Every
 reduction ends in ``.item()``, which forces a device sync -- that is
@@ -16,7 +16,7 @@
 
 import torch
 
-__all__ = ['iter_tensors', 'primary_tensor', 'tensor_moments', 'tensor_numel', 'describe_shape']
+__all__ = ['iter_tensors', 'primary_tensor', 'tensor_moments', 'tensor_numel', 'describe_shape', 'memory_stats']
 
 
 def iter_tensors(obj: Any) -> Iterator[torch.Tensor]:
@@ -84,3 +84,27 @@ def describe_shape(obj: Any) -> str:
     if len(tensors) == 1:
         return 'x'.join(str(dim) for dim in tensors[0].shape)
     return '(' + ', '.join('x'.join(str(dim) for dim in t.shape) for t in tensors) + ')'
+
+
+def memory_stats(*, synchronize: bool = False) -> dict[str, int]:
+    """Raw per-process allocator peaks and host-wide used RAM, in bytes.
+
+    Counters are lifetime peaks, not step peaks. Host usage is not process RSS.
+    Reading never resets counters. Device synchronization is an explicit opt-in
+    measurement barrier, disabled by default.
+    Unsupported counters are absent rather than invented zeroes.
+    """
+    import psutil
+
+    from twinkle import Platform
+
+    result = {"cpu_memory_used_bytes": psutil.virtual_memory().used}
+    device = getattr(torch, Platform.device_prefix(), None)
+    if device is not None and getattr(device, "is_available", lambda: False)():
+        if synchronize:
+            device.synchronize()
+        for name in ("max_memory_allocated", "max_memory_reserved"):
+            counter = getattr(device, name, None)
+            if counter is not None:
+                result[name + "_bytes"] = counter()
+    return result

src/twinkle/sampler/vllm_sampler/vllm_engine.py

@@ -14,7 +14,6 @@
 from twinkle.sampler.base_engine import BaseSamplerEngine
 from twinkle.utils import Platform
 from twinkle.utils.framework import Torch
-from twinkle.utils.parallel import PosixFileLock
 from twinkle.utils.zmq_utils import configure_zmq_socket, get_timeout_s_from_env
 
 logger = get_logger()
@@ -403,11 +402,10 @@ def _create_engine(self):
         engine_args = AsyncEngineArgs(**filtered_engine_config)
         vllm_config = engine_args.create_engine_config(usage_context=UsageContext.OPENAI_API_SERVER)
 
-        with PosixFileLock('/tmp/twinkle-vllm-engine-init.lock'):
-            engine = AsyncLLM.from_vllm_config(
-                vllm_config=vllm_config,
-                usage_context=UsageContext.OPENAI_API_SERVER,
-            )
+        engine = AsyncLLM.from_vllm_config(
+            vllm_config=vllm_config,
+            usage_context=UsageContext.OPENAI_API_SERVER,
+        )
 
         logger.info(f'VLLMEngine initialized: model={self.model_id}')
         return engine

src/twinkle_agentic/async_rl/cluster_adapter.py

@@ -967,6 +967,18 @@ def shutdown(self) -> None:
 
     # --- observation / decision ------------------------------------------------------------------------
 
+    def performance_snapshot(self) -> dict[str, float]:
+        """Commit-boundary resource counters; allocated GPU budget stays fixed."""
+        return {
+            'resource/shared_enabled': int(self._shared_enabled),
+            'resource/in_generation': int(self.in_generation),
+            'resource/flex_sampler_replicas': len(self._flex_workers),
+            'resource/free_flex_groups': len(self._free_flex_groups),
+            'resource/change_pending': int(self.change_pending),
+            'resource/last_change_step': self._last_change_step,
+            'resource/handover_seconds_total': self._counters.get('handover'),
+        }
+
     def observe(self, *, step: int, driver_stats: DriverStats) -> ResourceWindow:
         """Fold one batch boundary's stat deltas into the decision window (the commit boundary).