Refactor rollout.py by file decompositions

miles/ray/rollout/addr_allocator.py

@@ -0,0 +1,109 @@
+import logging
+
+import ray
+
+logger = logging.getLogger(__name__)
+
+
+def allocate_rollout_engine_addr_and_ports_normal(
+    *,
+    args,
+    rollout_engines,
+    worker_type="regular",
+    num_gpus_per_engine=None,
+    rank_offset=0,
+    base_port=15000,
+):
+    # get ports
+    # there are 4 ports we need to allocate
+    # 1. server port
+    # 2. nccl port
+    # 3. dist_init_addr port
+    # 4. other ports for dp_attention, which is of size 4 + dp_size
+    _gpus_per_engine = num_gpus_per_engine or args.rollout_num_gpus_per_engine
+    num_engines_per_node = max(1, args.num_gpus_per_node // _gpus_per_engine)
+    addr_and_ports: dict[int, dict] = {}
+
+    # Track per-node port cursors so that different server groups (called
+    # sequentially) never race for the same ports on a given node.
+    node_port_cursor: dict[int, int] = {}
+
+    visited_nodes = set()
+    for rank, engine in rollout_engines:
+        local_rank = rank - rank_offset
+        node_index = local_rank // num_engines_per_node
+        if node_index in visited_nodes:
+            continue
+        visited_nodes.add(node_index)
+        # TODO: currently when restarting engines, we will set port for all engines on this node starting with this rank.
+        # e.g. for 8 gpus, if we are restarting engine on gpu 3, we will set port for engine 3,4,5,6,7 on this node.
+        num_engines_on_this_node = num_engines_per_node - (local_rank % num_engines_per_node)
+
+        def get_addr_and_ports(engine, node_idx):
+            # use small ports to prevent ephemeral port between 32768 and 65536.
+            # also, ray uses port 10002-19999, thus we avoid near-10002 to avoid racing condition
+            start_port = node_port_cursor.get(node_idx, base_port)
+
+            def port(consecutive=1):
+                nonlocal start_port
+                _, port = ray.get(
+                    engine._get_current_node_ip_and_free_port.remote(
+                        start_port=start_port,
+                        consecutive=consecutive,
+                    )
+                )
+                start_port = port + consecutive
+                node_port_cursor[node_idx] = start_port
+                return port
+
+            def addr():
+                addr, _ = ray.get(engine._get_current_node_ip_and_free_port.remote())
+                return addr
+
+            return addr, port
+
+        get_addr, get_port = get_addr_and_ports(engine, node_index)
+
+        for i in range(num_engines_on_this_node):
+            current_rank = rank + i
+            addr_and_ports.setdefault(current_rank, {})
+            addr_and_ports[current_rank]["host"] = get_addr()
+            addr_and_ports[current_rank]["port"] = get_port()
+            addr_and_ports[current_rank]["nccl_port"] = get_port()
+            # Always allocate a unique engine_info_bootstrap_port per engine
+            addr_and_ports[current_rank]["engine_info_bootstrap_port"] = get_port()
+
+            if worker_type == "prefill":
+                addr_and_ports[current_rank]["disaggregation_bootstrap_port"] = get_port()
+
+        if _gpus_per_engine > args.num_gpus_per_node:
+            num_node_per_engine = _gpus_per_engine // args.num_gpus_per_node
+            if local_rank % num_node_per_engine == 0:
+                dist_init_addr = f"{get_addr()}:{get_port(30 + args.sglang_dp_size)}"
+                for i in range(num_node_per_engine):
+                    addr_and_ports.setdefault(rank + i, {})
+                    addr_and_ports[rank + i]["dist_init_addr"] = dist_init_addr
+        else:
+            for i in range(num_engines_on_this_node):
+                addr_and_ports[rank + i]["dist_init_addr"] = f"{get_addr()}:{get_port(30 + args.sglang_dp_size)}"
+
+    for i, _ in rollout_engines:
+        for key in ["port", "nccl_port", "dist_init_addr"]:
+            assert key in addr_and_ports[i], f"Engine {i} {key} is not set."
+        logger.info(f"Ports for engine {i}: {addr_and_ports[i]}")
+
+    return addr_and_ports, node_port_cursor
+
+
+def allocate_rollout_engine_addr_and_ports_external(args, rollout_engines):
+    addr_and_ports = {}
+    for rank, _ in rollout_engines:
+        addr = args.rollout_external_engine_addrs[rank]
+        [host, port] = addr.split(":")
+        addr_and_ports[rank] = dict(
+            dist_init_addr=addr,
+            nccl_port=None,
+            host=host,
+            port=int(port),
+        )
+    return addr_and_ports

miles/ray/rollout/debug_data.py

@@ -0,0 +1,30 @@
+import logging
+from pathlib import Path
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+# TODO extract `load_debug_rollout_data`
+
+
+# TODO: remove `self`
+def save_debug_rollout_data(self, data, rollout_id, evaluation: bool):
+    # TODO to be refactored (originally Buffer._set_data)
+    if (path_template := self.args.save_debug_rollout_data) is not None:
+        path = Path(path_template.format(rollout_id=("eval_" if evaluation else "") + str(rollout_id)))
+        logger.info(f"Save debug rollout data to {path}")
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        # TODO may improve the format
+        if evaluation:
+            dump_data = dict(
+                samples=[sample.to_dict() for dataset_name, info in data.items() for sample in info["samples"]]
+            )
+        else:
+            dump_data = dict(
+                samples=[sample.to_dict() for sample in data],
+            )
+
+        torch.save(dict(rollout_id=rollout_id, **dump_data), path)

miles/ray/rollout/metrics.py

@@ -0,0 +1,202 @@
+import logging
+from typing import Any
+
+import numpy as np
+
+from miles.utils import tracking_utils
+from miles.utils.iter_utils import group_by
+from miles.utils.metric_utils import (
+    compute_pass_rate,
+    compute_rollout_step,
+    compute_statistics,
+    dict_add_prefix,
+    has_repetition,
+)
+from miles.utils.misc import load_function
+from miles.utils.types import Sample
+
+
+logger = logging.getLogger(__name__)
+
+
+def log_eval_rollout_data(rollout_id, args, data, extra_metrics: dict[str, Any] | None = None):
+    if args.custom_eval_rollout_log_function_path is not None:
+        custom_log_func = load_function(args.custom_eval_rollout_log_function_path)
+        if custom_log_func(rollout_id, args, data, extra_metrics):
+            return
+
+    log_dict = extra_metrics or {}
+    for key in data.keys():
+        rewards = data[key]["rewards"]
+        log_dict[f"eval/{key}"] = sum(rewards) / len(rewards)
+        if (samples := data[key].get("samples")) is not None:
+            log_dict |= dict_add_prefix(_compute_metrics_from_samples(args, samples), f"eval/{key}/")
+        if "truncated" in data[key]:
+            truncated = data[key]["truncated"]
+            log_dict[f"eval/{key}-truncated_ratio"] = sum(truncated) / len(truncated)
+        if args.log_passrate:
+            log_dict |= dict_add_prefix(
+                compute_pass_rate(
+                    flat_rewards=rewards,
+                    group_size=args.n_samples_per_eval_prompt,
+                ),
+                f"eval/{key}-",
+            )
+
+    logger.info(f"eval {rollout_id}: {log_dict}")
+
+    step = compute_rollout_step(args, rollout_id)
+    log_dict["eval/step"] = step
+    tracking_utils.log(args, log_dict, step_key="eval/step")
+
+    return log_dict
+
+
+def log_rollout_data(rollout_id, args, samples, rollout_extra_metrics, rollout_time):
+    if args.custom_rollout_log_function_path is not None:
+        custom_log_func = load_function(args.custom_rollout_log_function_path)
+        if custom_log_func(rollout_id, args, samples, rollout_extra_metrics, rollout_time):
+            return
+
+    if args.load_debug_rollout_data:
+        return
+
+    log_dict = {**(rollout_extra_metrics or {})}
+    log_dict |= dict_add_prefix(_compute_metrics_from_samples(args, samples), "rollout/")
+    log_dict |= dict_add_prefix(_compute_perf_metrics_from_samples(args, samples, rollout_time), "perf/")
+    logger.info(f"perf {rollout_id}: {log_dict}")
+    step = compute_rollout_step(args, rollout_id)
+    log_dict["rollout/step"] = step
+    tracking_utils.log(args, log_dict, step_key="rollout/step")
+
+
+def _compute_metrics_from_samples(args, samples):
+    response_lengths = [sample.effective_response_length for sample in samples]
+
+    log_dict = {}
+    log_dict |= dict_add_prefix(compute_statistics(response_lengths), "response_len/")
+    log_dict |= _compute_zero_std_metrics(args, samples)
+    log_dict |= _compute_spec_metrics(args, samples)
+    log_dict |= _compute_prefix_cache_metrics(args, samples)
+    log_dict |= _compute_reward_cat_metrics(args, samples)
+    log_dict["repetition_frac"] = np.mean([int(has_repetition(s.response)) for s in samples]).item()
+    log_dict["truncated_ratio"] = np.mean([int(s.status == Sample.Status.TRUNCATED) for s in samples]).item()
+
+    oldest_versions = [s.oldest_weight_version for s in samples if s.oldest_weight_version is not None]
+    if oldest_versions:
+        log_dict |= dict_add_prefix(compute_statistics(oldest_versions), "weight_version/")
+        mixed = sum(1 for s in samples if len(set(s.weight_versions)) > 1)
+        log_dict["weight_version/mixed_version_ratio"] = mixed / len(samples)
+
+    tito_vals = [s.metadata.get("tito_session_mismatch") for s in samples]
+    tito_vals = [v for v in tito_vals if v is not None]
+    if tito_vals:
+        log_dict["tito_session_mismatch_rate"] = np.mean([len(v) > 0 for v in tito_vals]).item()
+        for mtype in ("special_token_count", "special_token_type", "non_assistant_text", "assistant_text"):
+            log_dict[f"tito_session_mismatch_rate/{mtype}"] = np.mean(
+                [any(m.get("type") == mtype for m in v) for v in tito_vals]
+            ).item()
+        if args.ci_test:
+            for strict_type in ("special_token_count", "special_token_type", "non_assistant_text"):
+                rate = log_dict.get(f"tito_session_mismatch_rate/{strict_type}", 0)
+                assert rate == 0, (
+                    f"tito_session_mismatch_rate/{strict_type}={rate:.4f} must be 0 — "
+                    "this indicates a bug in the TITO algorithm or chat template. "
+                    "Please check your tito model and chat template."
+                )
+            # assistant_text mismatch is non-critical: assistant tokens are inherited
+            # from the pretokenized prefix and may differ from canonical tokenization.
+
+    return log_dict
+
+
+def _compute_perf_metrics_from_samples(args, samples, rollout_time):
+    non_generation_time = [sample.non_generation_time for sample in samples]
+
+    log_dict = {}
+    log_dict["rollout_time"] = rollout_time
+    if max(non_generation_time) > 0:
+        log_dict |= dict_add_prefix(compute_statistics(non_generation_time), "non_generation_time/")
+
+    def token_perf(response_lengths, non_generation_time, key=""):
+        max_response_length = max(response_lengths)
+        if args.rollout_num_gpus:
+            log_dict[f"{key}tokens_per_gpu_per_sec"] = sum(response_lengths) / rollout_time / args.rollout_num_gpus
+        log_dict[f"longest_{key}sample_tokens_per_sec"] = max_response_length / rollout_time
+
+        if max(non_generation_time) == 0:
+            return
+
+        non_generation_time = [
+            t for t, length in zip(non_generation_time, response_lengths, strict=True) if length == max_response_length
+        ]
+        mean_non_generation_time = sum(non_generation_time) / len(non_generation_time)
+
+        log_dict[f"longest_{key}sample_non_generation_time"] = mean_non_generation_time
+        log_dict[f"longest_{key}sample_tokens_per_sec_without_non_generation"] = max_response_length / (
+            rollout_time - mean_non_generation_time
+        )
+
+    token_perf([sample.response_length for sample in samples], non_generation_time, key="")
+    token_perf([sample.effective_response_length for sample in samples], non_generation_time, key="effective_")
+
+    return log_dict
+
+
+def _compute_zero_std_metrics(args, all_samples: list[Sample]):
+    # only compute in GRPO-like algorithms where one prompt has multiple responses
+    if args.advantage_estimator == "ppo":
+        return {}
+
+    def _is_zero_std(samples: list[Sample]):
+        rewards = [sample.get_reward_value(args) for sample in samples]
+        return len(rewards) == 0 or all(rewards[0] == r for r in rewards)
+
+    all_sample_groups = group_by(all_samples, lambda s: s.group_index)
+    interesting_sample_groups = [g for g in all_sample_groups.values() if _is_zero_std(g)]
+
+    interesting_rewards = [str(round(g[0].get_reward_value(args), 1)) for g in interesting_sample_groups]
+
+    counts = {reward: len(items) for reward, items in group_by(interesting_rewards).items()}
+    log_dict = {f"zero_std/count_{reward}": count for reward, count in counts.items()}
+
+    # Percentages over total groups, so "too hard" (all-0) and "too easy"
+    # (all-1) rates are comparable across runs without needing to know the
+    # rollout batch size.
+    total_groups = len(all_sample_groups)
+    if total_groups > 0:
+        log_dict["zero_std/all_zero_percentage"] = counts.get("0.0", 0) / total_groups
+        log_dict["zero_std/all_one_percentage"] = counts.get("1.0", 0) / total_groups
+
+    return log_dict
+
+
+def _compute_spec_metrics(args, all_samples: list[Sample]):
+    if args.sglang_speculative_algorithm is None:
+        return {}
+    num_samples = len(all_samples)
+    metrics = {}
+    metrics["spec_accept_rate"] = sum(sample.spec_info.spec_accept_rate for sample in all_samples) / num_samples
+    metrics["spec_accept_length"] = sum(sample.spec_info.spec_accept_length for sample in all_samples) / num_samples
+    return metrics
+
+
+def _compute_prefix_cache_metrics(args, all_samples: list[Sample]):
+    num_samples = len(all_samples)
+    metrics = {}
+    total_cached_tokens = sum(sample.prefix_cache_info.cached_tokens for sample in all_samples)
+    total_prompt_tokens = sum(sample.prefix_cache_info.total_prompt_tokens for sample in all_samples)
+
+    metrics["prefix_cache_hit_rate"] = total_cached_tokens / total_prompt_tokens if total_prompt_tokens > 0 else 0.0
+    metrics["avg_cached_tokens_per_sample"] = total_cached_tokens / num_samples
+    return metrics
+
+
+def _compute_reward_cat_metrics(args, all_samples: list[Sample]):
+    reward_cat_key = args.log_reward_category
+    if reward_cat_key is None:
+        return {}
+
+    samples_of_reward_cat = group_by(all_samples, lambda s: s.reward[reward_cat_key])
+
+    return {f"error_cat/{reward_cat}": len(s) / len(all_samples) for reward_cat, s in samples_of_reward_cat.items()}