diff --git a/src/twinkle/metric/__init__.py b/src/twinkle/metric/__init__.py index 3ef151464..a493c356d 100644 --- a/src/twinkle/metric/__init__.py +++ b/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 diff --git a/src/twinkle/metric/reporting.py b/src/twinkle/metric/reporting.py index de7f9265e..cda64d017 100644 --- a/src/twinkle/metric/reporting.py +++ b/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) diff --git a/src/twinkle/metric/train_metric.py b/src/twinkle/metric/train_metric.py index 8d785c38b..8610973cb 100644 --- a/src/twinkle/metric/train_metric.py +++ b/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 diff --git a/src/twinkle/model/base.py b/src/twinkle/model/base.py index 18eebd11a..37bc898c6 100644 --- a/src/twinkle/model/base.py +++ b/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]: ... diff --git a/src/twinkle/perf/__init__.py b/src/twinkle/perf/__init__.py index 4b2c2eeae..4b22a2bc7 100644 --- a/src/twinkle/perf/__init__.py +++ b/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', ] diff --git a/src/twinkle/perf/recorder.py b/src/twinkle/perf/recorder.py index 05fe1cf11..53243885d 100644 --- a/src/twinkle/perf/recorder.py +++ b/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,7 +136,7 @@ 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 @@ -141,7 +144,7 @@ def _span(self, name: str) -> Generator[None, None, None]: # 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 diff --git a/src/twinkle/perf/stats.py b/src/twinkle/perf/stats.py index 2d4cbd458..a11cb91a1 100644 --- a/src/twinkle/perf/stats.py +++ b/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 diff --git a/src/twinkle/sampler/vllm_sampler/vllm_engine.py b/src/twinkle/sampler/vllm_sampler/vllm_engine.py index 304e57049..849b26f55 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_engine.py +++ b/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 diff --git a/src/twinkle_agentic/async_rl/cluster_adapter.py b/src/twinkle_agentic/async_rl/cluster_adapter.py index f676ca8b0..840cdc86b 100644 --- a/src/twinkle_agentic/async_rl/cluster_adapter.py +++ b/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). diff --git a/src/twinkle_agentic/async_rl/streaming_driver.py b/src/twinkle_agentic/async_rl/streaming_driver.py index b35ccbaba..17e677ea9 100644 --- a/src/twinkle_agentic/async_rl/streaming_driver.py +++ b/src/twinkle_agentic/async_rl/streaming_driver.py @@ -7,7 +7,7 @@ import threading import time from collections.abc import Callable, Iterable, Iterator, Sequence -from dataclasses import dataclass +from dataclasses import asdict, dataclass from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union from .types import PartitionAdmission, RolloutPolicy @@ -186,7 +186,13 @@ def __init__(self, adapter_name: Optional[str] = None, max_steps: Optional[int] = None, resource_adapter: Optional[ClusterAdapter] = None, + on_iteration_start: Optional[Callable[[], None]] = None, + on_iteration_end: Optional[Callable[[bool], None]] = None, sleep: Optional[Callable[[float], None]] = None) -> None: + # Observers run on the driver thread. End follows ACTUAL publication, + # including the terminal case where the budget deliberately skips it. + self.on_iteration_start = on_iteration_start + self.on_iteration_end = on_iteration_end # Fail loudly on a mis-wired seam: a non-callable injected here would otherwise surface deep # inside the run loop (or silently skip a lifecycle step), far from the construction that caused it. # The group-barrier algorithm hooks: PRM scoring is the remote ``score`` shape's own channel, @@ -499,6 +505,8 @@ def run(self) -> None: # and the sampler generates from it, never an unpinned base model (None for in-place; a no-op # without version control, where the version is pinned at 0 and nothing publishes). self._register_initial_policy() + if self.on_iteration_start is not None: + self.on_iteration_start() # A second run() on the same driver must not inherit the previous run's fold snapshots. self.stats.last_pull_idle_ratio = 0.0 self._train_at_pull = 0.0 @@ -534,9 +542,12 @@ def run(self) -> None: resolved_groups: set = set() regen_queue: List[Tuple[Any, str]] = [] regen_counts: Dict[str, int] = {} + iteration_pending = False try: while True: if exhausted or self._aborted or self._default_reached_max(): + if iteration_pending and self.on_iteration_end is not None: + self.on_iteration_end(True) break # (1) Publish phase: run an EXTERNALLY-requested publication (see request_publish). The driver @@ -569,6 +580,13 @@ def run(self) -> None: # references (runs AFTER the stale scan, which released the dropped pins). self._prune() + if iteration_pending: + iteration_pending = False + if self.on_iteration_end is not None: + self.on_iteration_end(False) + if self.on_iteration_start is not None: + self.on_iteration_start() + # (2) Admission: open a fresh window while the staleness gate is open, then admit # per-sample until the stream ends or the backpressure cap is reached. if not stream_done or regen_queue: @@ -716,6 +734,7 @@ def run(self) -> None: self._reap_scores(records, blocking=True) _t = time.perf_counter() self._train_pull(records, ready, group_counts) + iteration_pending = not self._iterating or self._iteration_committed train_now = time.perf_counter() # Fold the bubble ratio at the pull boundary, where the pairing is exact: the # train delta is the PREVIOUS pull's training (already accumulated into the @@ -1077,6 +1096,7 @@ def _bridge_consume(self, records: List[BufferRecord]) -> int: """ batch = [record.sample for record in records] self._committed_steps = 0 + self._iteration_committed = False self._ack_event.clear() self._batch_queue.put(batch) # A single wait() would deadlock against abort(): abort's lone set() can land before this @@ -1090,6 +1110,14 @@ def _bridge_consume(self, records: List[BufferRecord]) -> int: # --- iterator regime (``for batch in driver:``) ------------------------------------------------------ + def performance_snapshot(self) -> dict: + """Export plain observations without exposing driver internals to recorders.""" + adapter = self._resource_adapter + return { + 'stream_stats': asdict(self.stats), + 'resource_metrics': adapter.performance_snapshot() if adapter is not None else {}, + } + def __iter__(self) -> 'StreamingDriver': """Drive the blocking :meth:`run` on a background thread and yield assembled, scored batches. @@ -1184,6 +1212,7 @@ def commit(self, optimizer_steps: Optional[int] = 1) -> None: """ # None means the caller omitted the count. An explicit 0 does not consume budget. self._committed_steps = 1 if optimizer_steps is None else int(optimizer_steps) + self._iteration_committed = True self._budget_steps += self._committed_steps self._ack_event.set() diff --git a/tests/metric/test_reporting.py b/tests/metric/test_reporting.py new file mode 100644 index 000000000..ee3b43787 --- /dev/null +++ b/tests/metric/test_reporting.py @@ -0,0 +1,69 @@ +"""Existing metric output supports resumed runs and raw observation attributes.""" + +import json +from concurrent.futures import ThreadPoolExecutor + +from twinkle.metric import MetricRecord, create_metrics_reporter + + +def reporter(path, run_id): + return create_metrics_reporter({"jsonl": { + "path": path, "summary_path": path.with_name("summary.json"), + "append": True, "batch_size": 1, + }}, run_id=run_id) + + +def test_existing_reporter_flushes_raw_observations_and_nonfinite_values(tmp_path): + path = tmp_path / "metrics.jsonl" + output = reporter(path, "run-1") + output.record(MetricRecord( + stage="train", optimizer_step=2, values={"tokens": 19, "loss": float("nan")}, + attributes={"memory_by_rank": [{"allocated_bytes": 4096}], "timing_spans": [(1, 2, 3)]}, + )) + output.flush() + row = json.loads(path.read_text()) + assert row["values"] == {"tokens": 19} + assert row["attributes"]["non_numeric_values"]["loss"] is None + assert row["attributes"]["memory_by_rank"] == [{"allocated_bytes": 4096}] + assert row["attributes"]["timing_spans"] == [[1, 2, 3]] + assert row["run_id"] == "run-1" and row["optimizer_step"] == 2 + output.close() + + +def test_append_does_not_overwrite_previous_run(tmp_path): + path = tmp_path / "metrics.jsonl" + for run_id in ("first", "second"): + output = reporter(path, run_id) + output.record(MetricRecord(stage="train", optimizer_step=1, values={"loss": 0.5})) + output.close() + rows = [json.loads(line) for line in path.read_text().splitlines()] + assert [row["run_id"] for row in rows] == ["first", "second"] + + +def test_existing_reporter_orders_concurrent_records_and_flushes_on_close(tmp_path): + path = tmp_path / "metrics.jsonl" + output = reporter(path, "run-1") + with ThreadPoolExecutor(max_workers=4) as pool: + list(pool.map(lambda step: output.record(MetricRecord(stage="train", optimizer_step=step, values={"x": step})), range(32))) + output.close() + rows = [json.loads(line) for line in path.read_text().splitlines()] + assert [row["sequence"] for row in rows] == list(range(1, 33)) + assert sorted(row["optimizer_step"] for row in rows) == list(range(32)) + assert output.health()["dropped_records"] == 0 + + +def test_terminal_totals_are_used_when_producer_reports_batch_scopes(tmp_path): + output = reporter(tmp_path / "metrics.jsonl", "run-1") + output.record(MetricRecord(stage="train", context_key="run-1", optimizer_step=1, values={"loss": 0.5})) + output.record(MetricRecord( + stage="rollout", context_key="run-1", values={"sample_count": 64}, attributes={"scope": "iteration"}, + )) + output.record(MetricRecord( + stage="run", context_key="run-1", values={"rollout_groups": 8, "rollout_samples": 64, "trained_samples": 128}, + )) + output.close() + summary = json.loads((tmp_path / "summary.json").read_text()) + assert summary["train_steps"] == 1 + assert summary["rollout_groups"] == 8 + assert summary["rollout_samples"] == 64 + assert summary["trained_samples"] == 128 diff --git a/tests/metric/test_throughput.py b/tests/metric/test_throughput.py new file mode 100644 index 000000000..5d7d1b037 --- /dev/null +++ b/tests/metric/test_throughput.py @@ -0,0 +1,92 @@ +"""Numerical contracts for rollout accounting, padding and durable records.""" + +import math + +import pytest +from twinkle.metric import ThroughputMetric, memory_metrics + + +def compute_performance_metrics(prompts, responses, timings, devices, **widths): + metric = ThroughputMetric() + metric.accumulate(prompt_lengths=prompts, response_lengths=responses, timings=timings, num_devices=devices) + return metric.calculate(**widths) + + +@pytest.mark.parametrize("widths,clip,non_aborted_clip", [((None, None), 1 / 3, 0.5), ((12, 8), 0.0, 0.0)]) +def test_variable_lengths_and_explicit_padded_width(widths, clip, non_aborted_clip): + timings = { + "step": 7.25, + "gen": 3.5, + "ref": 0.25, + "old_log_prob": 0.4, + "adv": 0.01, + "update_actor": 2, + "update_weights": 0.5, + } + actual = compute_performance_metrics( + [3, 6, 4], [0, 2, 5], timings, 8, prompt_width=widths[0], response_width=widths[1] + ) + expected = { + "perf/total_num_tokens": 20, + "perf/num_prompt_tokens": 13, + "perf/num_response_tokens": 7, + "perf/num_samples": 3, + "perf/num_gpus": 8, + "perf/time_per_step": 7.25, + "perf/throughput": 0.3448275862068966, + "timing_per_token_ms/gen": 500, + "timing_per_token_ms/ref": 12.5, + "timing_per_token_ms/adv": 0.5, + "timing_per_token_ms/update_actor": 100, + "prompt_length/mean": 13 / 3, + "prompt_length/min": 3, + "prompt_length/max": 6, + "prompt_length/clip_ratio": clip, + "response_length/mean": 7 / 3, + "response_length/min": 0, + "response_length/max": 5, + "response_length/clip_ratio": clip, + "response/aborted_ratio": 1 / 3, + "response_length_non_aborted/mean": 3.5, + "response_length_non_aborted/min": 2, + "response_length_non_aborted/max": 5, + "response_length_non_aborted/clip_ratio": non_aborted_clip, + } + for key, value in expected.items(): + assert actual[key] == pytest.approx(value), key + assert "timing_per_token_ms/old_log_prob" not in actual + + +def test_all_empty_responses_have_no_non_aborted_population(): + actual = compute_performance_metrics([2, 3], [0, 0], {"step": 7.25, "gen": 3.5, "update_actor": 2}, 8) + assert actual["perf/total_num_tokens"] == 5 + assert actual["perf/throughput"] == pytest.approx(0.08620689655172414) + assert actual["response/aborted_ratio"] == 1 + assert actual["response_length/clip_ratio"] == 1 + assert actual["timing_per_token_ms/gen"] == 0 + assert actual["timing_per_token_ms/update_actor"] == 400 + assert math.isnan(actual["response_length_non_aborted/mean"]) + + +def test_invalid_allocation_or_timing_fails_instead_of_faking_throughput(): + for n_gpus, seconds in [(0, 1), (2, 0), (2, float("nan"))]: + with pytest.raises(ValueError): + compute_performance_metrics([3], [2], {"step": seconds}, n_gpus) + + +def test_window_throughput_uses_sum_tokens_over_sum_seconds(): + metric = ThroughputMetric() + metric.accumulate(prompt_lengths=[2], response_lengths=[8], timings={"step": 1}, num_devices=2) + metric.accumulate(prompt_lengths=[2], response_lengths=[18], timings={"step": 4}, num_devices=2) + assert metric.calculate()["perf/throughput"] == 3 # 30 / 5 / 2, not mean(5, 2.5) + metric.reset() + assert metric.calculate() == {} + + +def test_memory_reduction_uses_rank_max_for_peaks_and_mean_for_host_ram(): + assert memory_metrics( + [ + {"max_memory_allocated_bytes": 3 * 1024**3, "cpu_memory_used_bytes": 8 * 1024**3}, + {"max_memory_allocated_bytes": 5 * 1024**3, "cpu_memory_used_bytes": 10 * 1024**3}, + ] + ) == {"perf/max_memory_allocated_gb": 5, "perf/cpu_memory_used_gb": 9} diff --git a/tests/perf/test_observations.py b/tests/perf/test_observations.py new file mode 100644 index 000000000..b0ae0b99b --- /dev/null +++ b/tests/perf/test_observations.py @@ -0,0 +1,36 @@ +"""Atomic observations carry units and stay independent of recipe policy.""" + +from types import SimpleNamespace + +import torch + +from twinkle import Platform +from twinkle.perf.stats import memory_stats + + +def test_memory_returns_raw_bytes_without_implicit_sync_or_reset(monkeypatch): + import psutil + + actions = [] + monkeypatch.setattr(Platform, "device_prefix", staticmethod(lambda: "cuda")) + monkeypatch.setattr(psutil, "virtual_memory", lambda: SimpleNamespace(used=123456)) + monkeypatch.setattr(torch, "cuda", SimpleNamespace( + is_available=lambda: True, + synchronize=lambda: actions.append("sync"), + max_memory_allocated=lambda: 8192, + max_memory_reserved=lambda: 16384, + )) + assert memory_stats() == { + "cpu_memory_used_bytes": 123456, + "max_memory_allocated_bytes": 8192, + "max_memory_reserved_bytes": 16384, + } + assert actions == [] + assert memory_stats(synchronize=True)["max_memory_allocated_bytes"] == 8192 + assert actions == ["sync"] + + +def test_unavailable_device_counters_are_absent(monkeypatch): + monkeypatch.setattr(Platform, "device_prefix", staticmethod(lambda: "cuda")) + monkeypatch.setattr(torch, "cuda", SimpleNamespace(is_available=lambda: False)) + assert set(memory_stats()) == {"cpu_memory_used_bytes"} diff --git a/tests/twinkle_agentic/test_iteration_observers.py b/tests/twinkle_agentic/test_iteration_observers.py new file mode 100644 index 000000000..003728d03 --- /dev/null +++ b/tests/twinkle_agentic/test_iteration_observers.py @@ -0,0 +1,71 @@ +"""Observe actual post-commit publication, including terminal and error paths.""" + +from types import SimpleNamespace + +import pytest + +from twinkle_agentic.async_rl.streaming_driver import StreamingDriver + + +def make_driver(**kwargs): + return StreamingDriver( + prompt_stream=list(range(12)), + max_staleness=0, + num_generations=1, + groups_per_partition=1, + buffer_depth=4, + need_rows=4, + group_pull=True, + version_control=True, + max_steps=2, + submit=lambda ref, index, policy: (ref.ordinal, index), + poll=lambda handles: list(handles), + collect=lambda handle: SimpleNamespace(), + consume=None, + **kwargs, + ) + + +def test_observer_runs_after_real_publication_and_before_next_batch(): + events = [] + driver = make_driver(on_iteration_start=lambda: events.append("start")) + driver.on_iteration_end = lambda terminal: events.append(("end", terminal, driver.stats.publishes)) + try: + for samples in driver: + events.append("train") + driver.request_publish(lambda window: events.append("publish")) + driver.commit(1) + finally: + driver.abort() + assert events == ["start", "train", "publish", ("end", False, 1), "start", "train", ("end", True, 1)] + + +def test_failed_uncommitted_batch_is_not_a_completed_iteration(): + completions = [] + driver = make_driver(on_iteration_end=lambda terminal: completions.append(terminal)) + next(iter(driver)) + driver.abort() # recipe threw before commit + assert completions == [] + + +def test_observer_failure_propagates_to_recipe(): + def failed(terminal): + raise ValueError("cannot persist iteration") + + driver = make_driver(on_iteration_end=failed) + try: + with pytest.raises(ValueError, match="cannot persist"): + for samples in driver: + driver.commit(1) + finally: + driver.abort() + + +def test_performance_snapshot_is_plain_detached_data(): + driver = make_driver() + driver.stats.publishes = 2 + snapshot = driver.performance_snapshot() + assert snapshot["stream_stats"]["publishes"] == 2 + assert snapshot["resource_metrics"] == {} + snapshot["stream_stats"]["publishes"] = 99 + assert driver.stats.publishes == 2 diff --git a/tests/sampler/test_vllm_startup_lock.py b/tests/utils/test_posix_file_lock.py similarity index 58% rename from tests/sampler/test_vllm_startup_lock.py rename to tests/utils/test_posix_file_lock.py index 06e5dc47c..876013014 100644 --- a/tests/sampler/test_vllm_startup_lock.py +++ b/tests/utils/test_posix_file_lock.py @@ -5,30 +5,30 @@ from twinkle.utils.parallel import PosixFileLock -def _hold_startup_lock(lock_path: str, acquired, release) -> None: +def _hold_lock(lock_path: str, acquired, release) -> None: with PosixFileLock(lock_path): acquired.set() if not release.wait(timeout=30): - raise TimeoutError('test did not release vLLM startup lock') + raise TimeoutError("test did not release file lock") -def _acquire_startup_lock(lock_path: str, started, acquired) -> None: +def _acquire_lock(lock_path: str, started, acquired) -> None: started.set() with PosixFileLock(lock_path): acquired.set() -@pytest.mark.skipif(os.name == 'nt', reason='vLLM startup lock requires fcntl') -def test_vllm_engine_startup_is_serialized(tmp_path): - """Concurrent local sampler actors must not race vLLM's free-port probe.""" - context = multiprocessing.get_context('spawn') - lock_path = str(tmp_path / 'vllm-engine-init.lock') +@pytest.mark.skipif(os.name == "nt", reason="POSIX file lock requires fcntl") +def test_posix_file_lock_serializes_processes(tmp_path): + """A held advisory lock excludes an independent process until released.""" + context = multiprocessing.get_context("spawn") + lock_path = str(tmp_path / "process.lock") first_acquired = context.Event() release_first = context.Event() second_started = context.Event() second_acquired = context.Event() - first = context.Process(target=_hold_startup_lock, args=(lock_path, first_acquired, release_first)) - second = context.Process(target=_acquire_startup_lock, args=(lock_path, second_started, second_acquired)) + first = context.Process(target=_hold_lock, args=(lock_path, first_acquired, release_first)) + second = context.Process(target=_acquire_lock, args=(lock_path, second_started, second_acquired)) try: first.start()