From 96eecb4acae9e76df9552913eab4f87c89bf0290 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 01/14] Add a deterministic_random reward type Add a `deterministic_random` reward that hashes the sample tokens + response to produce a stable pseudo-random 0/1 reward, used for reproducible fault-tolerance / CI tests. - rm_hub/__init__.py (+ test). --- miles/rollout/rm_hub/__init__.py | 5 +++++ tests/fast/rollout/rm_hub/test_rm_hub.py | 25 ++++++++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/miles/rollout/rm_hub/__init__.py b/miles/rollout/rm_hub/__init__.py index 6264e17c304..962c4217558 100644 --- a/miles/rollout/rm_hub/__init__.py +++ b/miles/rollout/rm_hub/__init__.py @@ -1,4 +1,5 @@ import asyncio +import hashlib import random import aiohttp @@ -62,6 +63,10 @@ async def async_rm(args, sample: Sample, **kwargs): return compute_ifbench_reward(response, label, metadata=metadata) elif rm_type == "random": return random.randint(0, 1) + elif rm_type == "deterministic_random": + content = str(sample.tokens) + response + content_hash = hashlib.sha256(content.encode()).digest() + return int(content_hash[0]) % 2 elif rm_type: raise NotImplementedError(f"Rule-based RM for {rm_type} is not implemented.") else: diff --git a/tests/fast/rollout/rm_hub/test_rm_hub.py b/tests/fast/rollout/rm_hub/test_rm_hub.py index a3dadbdaf00..eb97afa8cc5 100644 --- a/tests/fast/rollout/rm_hub/test_rm_hub.py +++ b/tests/fast/rollout/rm_hub/test_rm_hub.py @@ -51,6 +51,31 @@ def test_random_rm(self, mock_args): reward = run(async_rm(mock_args, sample)) assert reward in [0, 1] + def test_deterministic_random_rm_returns_binary(self, mock_args): + mock_args.rm_type = "deterministic_random" + sample = Sample(prompt="", response="hello", label="", tokens=[1, 2, 3]) + reward = run(async_rm(mock_args, sample)) + assert reward in [0, 1] + + def test_deterministic_random_rm_is_deterministic(self, mock_args): + mock_args.rm_type = "deterministic_random" + sample = Sample(prompt="", response="hello world", label="", tokens=[10, 20]) + rewards = [run(async_rm(mock_args, sample)) for _ in range(5)] + assert len(set(rewards)) == 1 + + def test_deterministic_random_rm_differs_by_response(self, mock_args): + mock_args.rm_type = "deterministic_random" + samples = [Sample(prompt="", response=f"response_{i}", label="", tokens=[1, 2, 3]) for i in range(20)] + rewards = [run(async_rm(mock_args, s)) for s in samples] + assert 0 in rewards and 1 in rewards + + def test_deterministic_random_rm_differs_by_tokens(self, mock_args): + """Same response with different tokens yields both reward values across many samples.""" + mock_args.rm_type = "deterministic_random" + samples = [Sample(prompt="", response="same", label="", tokens=[i, i + 1, i + 2]) for i in range(20)] + rewards = [run(async_rm(mock_args, s)) for s in samples] + assert 0 in rewards and 1 in rewards + def test_rm_type_from_metadata(self, mock_args): mock_args.rm_type = None sample = Sample(prompt="", response=r"\boxed{42}", label="42", metadata={"rm_type": "math"}) From 2c1d4a74049aac27508151b96434a6358796ac6d Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 02/14] Add an inplace_modify_args context manager Add `inplace_modify_args`, a context manager that temporarily overrides args attributes and restores them on exit (asserting they weren't clobbered), used to scope per-attempt argument overrides in the fault-tolerant trainer. - argparse_utils.py (+ test). --- miles/utils/argparse_utils.py | 20 ++++++++++ tests/fast/utils/test_argparse_utils.py | 51 ++++++++++++++++++++++++- 2 files changed, 70 insertions(+), 1 deletion(-) diff --git a/miles/utils/argparse_utils.py b/miles/utils/argparse_utils.py index 13b8f2b5cb8..b3c7f8f2f14 100644 --- a/miles/utils/argparse_utils.py +++ b/miles/utils/argparse_utils.py @@ -9,6 +9,8 @@ import argparse import dataclasses import types +from collections.abc import Iterator +from contextlib import contextmanager from pathlib import Path from typing import Generic, TypeVar, get_type_hints @@ -129,3 +131,21 @@ def to_cli_args(self, instance: T) -> str: parts.append(f"{flag} {value}") return " ".join(parts) + + +@contextmanager +def inplace_modify_args(args: argparse.Namespace, overrides: dict[str, object]) -> Iterator[None]: + """Temporarily set attributes on ``args``, restoring the originals on exit.""" + old_values = {key: getattr(args, key) for key in overrides} + for key, value in overrides.items(): + setattr(args, key, value) + try: + yield + finally: + for key, old_value in old_values.items(): + current = getattr(args, key) + assert current == overrides[key], ( + f"args.{key} was modified inside the inplace_modify_args block " + f"(expected {overrides[key]!r}, found {current!r}); restoring would clobber it" + ) + setattr(args, key, old_value) diff --git a/tests/fast/utils/test_argparse_utils.py b/tests/fast/utils/test_argparse_utils.py index 4221e24c6c9..2fe6b8e995d 100644 --- a/tests/fast/utils/test_argparse_utils.py +++ b/tests/fast/utils/test_argparse_utils.py @@ -7,7 +7,7 @@ import pytest -from miles.utils.argparse_utils import DataclassArgparseBridge +from miles.utils.argparse_utils import DataclassArgparseBridge, inplace_modify_args @dataclasses.dataclass(frozen=True) @@ -262,3 +262,52 @@ class _Bad: with pytest.raises(TypeError, match="Unsupported field type"): bridge.register_on_parser(parser) + + +class TestInplaceModifyArgs: + def test_overrides_inside_and_restores_on_exit(self) -> None: + """Overridden attributes are visible inside the block and restored afterwards.""" + args = argparse.Namespace(no_load_optim=True, finetune=True, lr=1.0) + + with inplace_modify_args(args, dict(no_load_optim=False, finetune=False)): + assert args.no_load_optim is False + assert args.finetune is False + assert args.lr == 1.0 + + assert args.no_load_optim is True + assert args.finetune is True + + def test_restores_on_exception(self) -> None: + """Originals are restored even when the block raises.""" + args = argparse.Namespace(flag=True) + + with pytest.raises(RuntimeError): + with inplace_modify_args(args, dict(flag=False)): + raise RuntimeError("boom") + + assert args.flag is True + + def test_empty_overrides_is_noop(self) -> None: + """An empty override dict changes nothing.""" + args = argparse.Namespace(flag=True) + + with inplace_modify_args(args, {}): + assert args.flag is True + + assert args.flag is True + + def test_unknown_attribute_raises(self) -> None: + """Overriding an attribute the namespace does not have fails loudly.""" + args = argparse.Namespace() + + with pytest.raises(AttributeError): + with inplace_modify_args(args, dict(missing=1)): + pass + + def test_mutation_inside_block_fails_on_exit(self) -> None: + """An attribute mutated inside the block is detected instead of silently clobbered.""" + args = argparse.Namespace(flag=True) + + with pytest.raises(AssertionError, match="modified inside"): + with inplace_modify_args(args, dict(flag=False)): + args.flag = True From 7b00c7567c580b97b96bd95b7e6b1d1712473b72 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 03/14] Add fault-tolerance support tweaks to shared utilities Small shared-utility additions used by the fault-tolerant trainer: hash non-contiguous tensors safely (reshape before viewing as bytes), an `enable_experimental_ft_trainer` env flag, forward NCCL_DEBUG/NCCL_DEBUG_FILE to worker environments, and a `filter_keys` helper. - ci_utils.py / environ.py / external_utils/command_utils.py / misc.py. --- miles/backends/megatron_utils/ci_utils.py | 2 +- miles/utils/environ.py | 15 ++++++ miles/utils/external_utils/command_utils.py | 6 ++- miles/utils/misc.py | 13 +++++ .../backends/megatron_utils/test_ci_utils.py | 50 +++++++++++++++++++ tests/fast/utils/test_misc.py | 37 +++++++++++++- 6 files changed, 120 insertions(+), 3 deletions(-) create mode 100644 tests/fast/backends/megatron_utils/test_ci_utils.py diff --git a/miles/backends/megatron_utils/ci_utils.py b/miles/backends/megatron_utils/ci_utils.py index 3ecf15eab0f..5837358bb73 100644 --- a/miles/backends/megatron_utils/ci_utils.py +++ b/miles/backends/megatron_utils/ci_utils.py @@ -36,7 +36,7 @@ def _hash_tensor_bytes(tensor: torch.Tensor) -> bytes: data = data.cpu() if not data.is_contiguous(): data = data.contiguous() - return data.view(torch.uint8).numpy().tobytes() + return data.reshape(-1).view(torch.uint8).numpy().tobytes() def compute_model_hashes_by_layer(model: Sequence[DDP]) -> dict[str, str]: diff --git a/miles/utils/environ.py b/miles/utils/environ.py index 35d1f350eed..1da32045b78 100644 --- a/miles/utils/environ.py +++ b/miles/utils/environ.py @@ -12,3 +12,18 @@ def enable_experimental_rollout_refactor() -> bool: _printed_experimental_rollout_refactor = True return result + + +_printed_experimental_ft_trainer = False + + +def enable_experimental_ft_trainer() -> bool: + raw = os.environ.get("MILES_EXPERIMENTAL_FT_TRAINER", "0").lower() + result = raw in ("1", "true", "on", "yes") + + global _printed_experimental_ft_trainer + if result and not _printed_experimental_ft_trainer: + print("MILES_EXPERIMENTAL_FT_TRAINER=1 is enabled (experimental feature)") + _printed_experimental_ft_trainer = True + + return result diff --git a/miles/utils/external_utils/command_utils.py b/miles/utils/external_utils/command_utils.py index d016e01ac34..7489d89ec99 100644 --- a/miles/utils/external_utils/command_utils.py +++ b/miles/utils/external_utils/command_utils.py @@ -156,7 +156,11 @@ def execute_train( } ), "NCCL_NVLS_ENABLE": os.environ.get("NCCL_NVLS_ENABLE", str(int(check_has_nvlink()))), - **{k: os.environ[k] for k in ("NCCL_SOCKET_IFNAME", "GLOO_SOCKET_IFNAME") if k in os.environ}, + **{ + k: os.environ[k] + for k in ("NCCL_SOCKET_IFNAME", "GLOO_SOCKET_IFNAME", "NCCL_DEBUG", "NCCL_DEBUG_FILE") + if k in os.environ + }, "no_proxy": f"127.0.0.1,{master_addr}", # This is needed by megatron / torch distributed in multi-node setup "MASTER_ADDR": master_addr, diff --git a/miles/utils/misc.py b/miles/utils/misc.py index 313a60dda90..3c160cc5cbc 100644 --- a/miles/utils/misc.py +++ b/miles/utils/misc.py @@ -1,14 +1,19 @@ import asyncio import importlib +import logging import re import subprocess +from collections.abc import Sequence from contextlib import contextmanager +from typing import Any import ray from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy from miles.utils.http_utils import is_port_available +logger = logging.getLogger(__name__) + # Mainly used for test purpose where `load_function` needs to load many in-flight generated functions class FunctionRegistry: @@ -199,3 +204,11 @@ def should_run_periodic_action( async def as_completed_async(tasks): for coro in asyncio.as_completed(tasks): yield await coro + + +def filter_keys(d: dict[str, Any], interest_keys: Sequence[str]) -> dict[str, Any]: + try: + return {k: d[k] for k in interest_keys} + except Exception: + logger.error(f"filter_keys d.keys={list(d)} {interest_keys=}", exc_info=True) + raise diff --git a/tests/fast/backends/megatron_utils/test_ci_utils.py b/tests/fast/backends/megatron_utils/test_ci_utils.py new file mode 100644 index 00000000000..21eee55fb87 --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_ci_utils.py @@ -0,0 +1,50 @@ +import pytest + +pytest.importorskip("megatron.core.distributed") + +import torch + +from miles.backends.megatron_utils.ci_utils import _hash_tensor_bytes + + +def test_hash_tensor_bytes_contiguous_float32_returns_raw_buffer_bytes() -> None: + """A contiguous 2D float32 tensor hashes to its exact raw little-endian buffer bytes.""" + tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3) + result = _hash_tensor_bytes(tensor) + assert isinstance(result, bytes) + assert len(result) == 6 * 4 + assert result == tensor.reshape(-1).contiguous().numpy().tobytes() + + +def test_hash_tensor_bytes_int64_byte_length_is_eight_per_element() -> None: + """An int64 tensor yields 8 bytes per element with bit-exact buffer contents.""" + tensor = torch.arange(6, dtype=torch.int64).reshape(2, 3) + result = _hash_tensor_bytes(tensor) + assert len(result) == 6 * 8 + assert result == tensor.reshape(-1).contiguous().numpy().tobytes() + + +def test_hash_tensor_bytes_noncontiguous_transpose_does_not_raise() -> None: + """A non-contiguous transposed view is hashed without raising (regression guard).""" + base = torch.arange(6, dtype=torch.float32).reshape(2, 3) + transposed = base.t() + assert not transposed.is_contiguous() + result = _hash_tensor_bytes(transposed) + assert isinstance(result, bytes) + assert len(result) == 6 * 4 + + +def test_hash_tensor_bytes_noncontiguous_matches_contiguous_row_major_bytes() -> None: + """The non-contiguous path returns the contiguous row-major bytes of the transposed values.""" + base = torch.arange(6, dtype=torch.float32).reshape(2, 3) + transposed = base.t() + result = _hash_tensor_bytes(transposed) + assert result == transposed.contiguous().reshape(-1).numpy().tobytes() + + +def test_hash_tensor_bytes_distinct_values_produce_distinct_bytes() -> None: + """Different tensor contents of the same shape and dtype produce different bytes.""" + a = torch.zeros(2, 3, dtype=torch.float32) + b = torch.ones(2, 3, dtype=torch.float32) + assert _hash_tensor_bytes(a) != _hash_tensor_bytes(b) + assert len(_hash_tensor_bytes(a)) == len(_hash_tensor_bytes(b)) == 6 * 4 diff --git a/tests/fast/utils/test_misc.py b/tests/fast/utils/test_misc.py index 810c2b67c75..1cc77107868 100644 --- a/tests/fast/utils/test_misc.py +++ b/tests/fast/utils/test_misc.py @@ -1,8 +1,9 @@ +import logging import os import pytest -from miles.utils.misc import FunctionRegistry, function_registry, load_function +from miles.utils.misc import FunctionRegistry, filter_keys, function_registry, load_function def _fn_a(): @@ -57,3 +58,37 @@ def test_registry_takes_precedence(self): with function_registry.temporary("os.path.join", _fn_b): assert load_function("os.path.join") is _fn_b assert load_function("os.path.join") is os.path.join + + +class TestFilterKeys: + def test_projects_dict_by_keys(self): + """filter_keys returns only the requested keys with their values.""" + d = {"a": 1, "b": 2, "c": 3} + assert filter_keys(d, ["a", "c"]) == {"a": 1, "c": 3} + + def test_empty_interest_keys_returns_empty_dict(self): + """An empty interest list yields an empty dict regardless of input.""" + assert filter_keys({"a": 1, "b": 2}, []) == {} + + def test_preserves_interest_keys_order(self): + """Result key order follows interest_keys, not the source dict order.""" + d = {"a": 1, "b": 2, "c": 3} + assert list(filter_keys(d, ["c", "a"]).keys()) == ["c", "a"] + + def test_full_subset_returns_all_entries(self): + """Requesting every key returns the whole projection.""" + d = {"x": 10, "y": 20} + assert filter_keys(d, ["x", "y"]) == {"x": 10, "y": 20} + + def test_duplicate_interest_key_collapses_to_single_entry(self): + """A repeated interest key produces a single dict entry.""" + d = {"a": 1, "b": 2} + assert filter_keys(d, ["a", "a"]) == {"a": 1} + + def test_missing_key_raises_key_error_and_logs(self, caplog): + """A missing key raises KeyError and logs the error with context.""" + d = {"a": 1} + with caplog.at_level(logging.ERROR, logger="miles.utils.misc"): + with pytest.raises(KeyError): + filter_keys(d, ["a", "missing"]) + assert any("filter_keys" in record.message for record in caplog.records) From 635aa50615b31f83aa9e98c5747912297d9c3a8e Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 04/14] Preserve the process-group backend across reload Thread the original backend through ReloadableProcessGroup so that, when a process group is rebuilt (e.g. after a reconfigure/heal), it is recreated with the same backend instead of hard-coding NCCL. - reloadable_process_group.py: carry `backend` in the reload group info. --- miles/utils/reloadable_process_group.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/miles/utils/reloadable_process_group.py b/miles/utils/reloadable_process_group.py index d9f5412953f..295dd10e999 100644 --- a/miles/utils/reloadable_process_group.py +++ b/miles/utils/reloadable_process_group.py @@ -42,7 +42,8 @@ def new_group(*args, **kwargs): if len(ranks) == 1: return group - group = ReloadableProcessGroup(group, ranks) + backend = args[2] if len(args) >= 3 else kwargs.get("backend") + group = ReloadableProcessGroup(group, ranks, backend=backend) return group dist.new_group = new_group @@ -112,7 +113,7 @@ def convert(arg): class ReloadableProcessGroup(torch.distributed.ProcessGroup): GROUPS = {} - def __init__(self, group, ranks): + def __init__(self, group, ranks, backend=None): super().__init__( rank=dist.get_rank(group), size=dist.get_world_size(group), @@ -120,6 +121,8 @@ def __init__(self, group, ranks): self.group = group self.group_info = { "ranks": ranks, + # None = inherit the default backend at reload time. + "backend": backend, } pid = os.getpid() if pid not in ReloadableProcessGroup.GROUPS: @@ -155,7 +158,9 @@ def reload_process_groups(): for reloadable_group in reloadable_groups: if reloadable_group.group is not None: continue - group = old_new_group(ranks=reloadable_group.group_info["ranks"], backend="nccl") + group = old_new_group( + ranks=reloadable_group.group_info["ranks"], backend=reloadable_group.group_info["backend"] + ) reloadable_group.group = group def rank(self) -> int: From 7a33a63f7df671bd8df8d1a6246f91ceeafbb4a3 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 05/14] Add fault-tolerance foundation utilities Add small foundation utilities used across the fault-tolerance trainer: a strict pydantic base model, a retry helper, a tensor checksum helper, a per-cell megatron world-size computation, the TrainStepOutcome enum, and the IndepDPInfo dataclass describing a cell's independent-DP identity. - pydantic_utils.py / retry_utils.py / checksum_utils.py / megatron_args_utils.py / types.py / indep_dp.py and tests. --- miles/backends/megatron_utils/types.py | 6 + miles/utils/checksum_utils.py | 34 +++ miles/utils/indep_dp.py | 26 +++ miles/utils/megatron_args_utils.py | 7 + miles/utils/pydantic_utils.py | 9 + miles/utils/retry_utils.py | 38 +++ tests/fast/utils/test_checksum_utils.py | 67 ++++++ tests/fast/utils/test_retry_utils.py | 299 ++++++++++++++++++++++++ 8 files changed, 486 insertions(+) create mode 100644 miles/backends/megatron_utils/types.py create mode 100644 miles/utils/checksum_utils.py create mode 100644 miles/utils/indep_dp.py create mode 100644 miles/utils/megatron_args_utils.py create mode 100644 miles/utils/pydantic_utils.py create mode 100644 miles/utils/retry_utils.py create mode 100644 tests/fast/utils/test_checksum_utils.py create mode 100644 tests/fast/utils/test_retry_utils.py diff --git a/miles/backends/megatron_utils/types.py b/miles/backends/megatron_utils/types.py new file mode 100644 index 00000000000..9a15742ea67 --- /dev/null +++ b/miles/backends/megatron_utils/types.py @@ -0,0 +1,6 @@ +from enum import StrEnum, auto + + +class TrainStepOutcome(StrEnum): + NORMAL = auto() + DISCARDED_SHOULD_RETRY = auto() diff --git a/miles/utils/checksum_utils.py b/miles/utils/checksum_utils.py new file mode 100644 index 00000000000..64a5f1e1f0b --- /dev/null +++ b/miles/utils/checksum_utils.py @@ -0,0 +1,34 @@ +from typing import Any + +InferenceEngineChecksums = dict[str, str] + + +def flatten_inference_engine_checksums(check_weights_result: Any) -> list[InferenceEngineChecksums]: + engine_bodies = _flatten_to_inference_engine_bodies(check_weights_result) + surviving = [body for body in engine_bodies if body is not None] + assert surviving, ( + f"check_weights('checksum') returned no non-None engine bodies " + f"(got {len(engine_bodies)} entries, all None): {check_weights_result!r}" + ) + return [_merge_inference_engine_ranks(body) for body in surviving] + + +def _flatten_to_inference_engine_bodies(check_weights_result: Any) -> list[Any]: + return [engine_body for server in check_weights_result for server_group in server for engine_body in server_group] + + +def _merge_inference_engine_ranks(engine_body: dict[str, Any]) -> InferenceEngineChecksums: + # Ranks arrive in non-deterministic (zmq) order under TP>1; sort and prefix each tensor + # name with rank{r}/ so distinct shards' identically-named tensors never clobber. + assert engine_body.get("success", False), f"check_weights engine reported failure: {engine_body!r}" + ranks: list[dict[str, Any]] = engine_body.get("ranks", []) or [] + assert ranks, f"check_weights engine body has no ranks: {engine_body!r}" + + ranks_sorted = sorted(ranks, key=lambda r: r["parallelism_info"]["rank"]) + + merged: InferenceEngineChecksums = {} + for rank_info in ranks_sorted: + rank = rank_info["parallelism_info"]["rank"] + for name, value in rank_info["checksums"].items(): + merged[f"rank{rank}/{name}"] = value + return merged diff --git a/miles/utils/indep_dp.py b/miles/utils/indep_dp.py new file mode 100644 index 00000000000..4055add45b7 --- /dev/null +++ b/miles/utils/indep_dp.py @@ -0,0 +1,26 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class IndepDPInfo: + cell_index: int + num_cells: int + alive_rank: int + alive_size: int + quorum_id: int + alive_cell_indices: list[int] + + @classmethod + def create_trivial(cls) -> "IndepDPInfo": + return cls( + cell_index=0, + num_cells=1, + alive_rank=0, + alive_size=1, + quorum_id=0, + alive_cell_indices=[0], + ) + + def __post_init__(self): + assert self.alive_rank == self.alive_cell_indices.index(self.cell_index) + assert self.alive_size == len(self.alive_cell_indices) diff --git a/miles/utils/megatron_args_utils.py b/miles/utils/megatron_args_utils.py new file mode 100644 index 00000000000..2b58b1b8644 --- /dev/null +++ b/miles/utils/megatron_args_utils.py @@ -0,0 +1,7 @@ +""" +Utils for megatron arguments, but not related to megatron core logic +""" + + +def compute_megatron_world_size_except_dp(args) -> int: + return args.tensor_model_parallel_size * args.pipeline_model_parallel_size * args.context_parallel_size diff --git a/miles/utils/pydantic_utils.py b/miles/utils/pydantic_utils.py new file mode 100644 index 00000000000..43e9763b8e3 --- /dev/null +++ b/miles/utils/pydantic_utils.py @@ -0,0 +1,9 @@ +from pydantic import BaseModel, ConfigDict + + +class StrictBaseModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class FrozenStrictBaseModel(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) diff --git a/miles/utils/retry_utils.py b/miles/utils/retry_utils.py new file mode 100644 index 00000000000..ef98f6ea896 --- /dev/null +++ b/miles/utils/retry_utils.py @@ -0,0 +1,38 @@ +import asyncio +import logging +from collections.abc import Awaitable, Callable +from typing import Any + +logger = logging.getLogger(__name__) + +_DEFAULT_INITIAL_DELAY = 1.0 +_DEFAULT_MAX_DELAY = 60.0 +_DEFAULT_BACKOFF_FACTOR = 2.0 + + +async def retry( + fn: Callable[[int], Awaitable[Any]], + *, + initial_delay: float = _DEFAULT_INITIAL_DELAY, + max_delay: float = _DEFAULT_MAX_DELAY, + backoff_factor: float = _DEFAULT_BACKOFF_FACTOR, + sleep_fn: Callable[[float], Awaitable[None]] = asyncio.sleep, + max_attempts: int | None = None, +) -> None: + """Retry until ``fn`` does not throw, with exponential backoff.""" + assert max_attempts is None or max_attempts >= 1 + + attempt = 0 + delay = initial_delay + while True: + try: + await fn(attempt) + return + except Exception: + attempt += 1 + if max_attempts is not None and attempt >= max_attempts: + logger.warning(f"retry: attempt {attempt} failed, giving up (max_attempts={max_attempts})") + raise + logger.warning(f"retry: attempt {attempt} failed, retrying in {delay:.1f}s", exc_info=True) + await sleep_fn(delay) + delay = min(delay * backoff_factor, max_delay) diff --git a/tests/fast/utils/test_checksum_utils.py b/tests/fast/utils/test_checksum_utils.py new file mode 100644 index 00000000000..f5e0801c836 --- /dev/null +++ b/tests/fast/utils/test_checksum_utils.py @@ -0,0 +1,67 @@ +"""Tests for engine_weight_checksum.flatten_inference_engine_checksums.""" + +from typing import Any + +import pytest + +from miles.utils.checksum_utils import flatten_inference_engine_checksums + + +def _engine_body(*, success: bool, ranks: list[dict[str, Any]] | None) -> dict[str, Any]: + body: dict[str, Any] = {"success": success, "message": "ok"} + if ranks is not None: + body["ranks"] = ranks + return body + + +def _rank(rank: int, checksums: dict[str, str]) -> dict[str, Any]: + return {"checksums": checksums, "parallelism_info": {"rank": rank}} + + +class TestFlattenInferenceEngineChecksums: + def test_single_server_group_engine_single_rank(self) -> None: + """One server/group/engine with one rank yields one prefixed checksum dict.""" + result = [[[_engine_body(success=True, ranks=[_rank(0, {"w": "aaa"})])]]] + assert flatten_inference_engine_checksums(result) == [{"rank0/w": "aaa"}] + + def test_multiple_engines_flattened_in_order(self) -> None: + """Engines across servers/groups are flattened into a single ordered list.""" + result = [ + [[_engine_body(success=True, ranks=[_rank(0, {"w": "e0"})])]], + [[_engine_body(success=True, ranks=[_rank(0, {"w": "e1"})])]], + ] + assert flatten_inference_engine_checksums(result) == [{"rank0/w": "e0"}, {"rank0/w": "e1"}] + + def test_none_node_rank_payloads_filtered(self) -> None: + """None engine bodies (non-zero node ranks) are dropped before indexing.""" + result = [[[_engine_body(success=True, ranks=[_rank(0, {"w": "e0"})]), None]]] + assert flatten_inference_engine_checksums(result) == [{"rank0/w": "e0"}] + + def test_all_none_fails_loud(self) -> None: + """A None-only result means the checksum action did nothing, so fail loud.""" + result = [[[None, None]]] + with pytest.raises(AssertionError, match="no non-None engine bodies"): + flatten_inference_engine_checksums(result) + + def test_multi_rank_merged_with_rank_prefix(self) -> None: + """Multiple ranks of one engine merge into one dict, prefixed by rank to avoid clobber.""" + result = [[[_engine_body(success=True, ranks=[_rank(0, {"w": "r0"}), _rank(1, {"w": "r1"})])]]] + assert flatten_inference_engine_checksums(result) == [{"rank0/w": "r0", "rank1/w": "r1"}] + + def test_ranks_out_of_order_sorted_by_parallelism_rank(self) -> None: + """Ranks arriving out of order (zmq) are sorted by parallelism rank deterministically.""" + out_of_order = [[[_engine_body(success=True, ranks=[_rank(1, {"w": "r1"}), _rank(0, {"w": "r0"})])]]] + in_order = [[[_engine_body(success=True, ranks=[_rank(0, {"w": "r0"}), _rank(1, {"w": "r1"})])]]] + assert flatten_inference_engine_checksums(out_of_order) == flatten_inference_engine_checksums(in_order) + + def test_engine_failure_fails_loud(self) -> None: + """An engine reporting success=False fails loud rather than silently dropping it.""" + result = [[[_engine_body(success=False, ranks=[_rank(0, {"w": "aaa"})])]]] + with pytest.raises(AssertionError, match="reported failure"): + flatten_inference_engine_checksums(result) + + def test_engine_without_ranks_fails_loud(self) -> None: + """A success body with no ranks fails loud (nothing to compare).""" + result = [[[_engine_body(success=True, ranks=None)]]] + with pytest.raises(AssertionError, match="no ranks"): + flatten_inference_engine_checksums(result) diff --git a/tests/fast/utils/test_retry_utils.py b/tests/fast/utils/test_retry_utils.py new file mode 100644 index 00000000000..e3eb15f8431 --- /dev/null +++ b/tests/fast/utils/test_retry_utils.py @@ -0,0 +1,299 @@ +import pytest + +from miles.utils.retry_utils import retry + +pytestmark = pytest.mark.asyncio + + +class _FakeSleep: + """Records sleep calls without actually sleeping.""" + + def __init__(self) -> None: + self.delays: list[float] = [] + + async def __call__(self, delay: float) -> None: + self.delays.append(delay) + + +class TestRetryBasic: + async def test_succeeds_immediately(self): + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + + await retry(fn, sleep_fn=fake_sleep) + + assert call_count == 1 + assert fake_sleep.delays == [] + + async def test_retries_then_succeeds(self): + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count < 4: + raise ValueError("not yet") + + await retry(fn, initial_delay=1.0, sleep_fn=fake_sleep) + + assert call_count == 4 + assert len(fake_sleep.delays) == 3 + + async def test_single_retry(self): + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count < 2: + raise RuntimeError("fail once") + + await retry(fn, initial_delay=1.0, sleep_fn=fake_sleep) + + assert call_count == 2 + assert len(fake_sleep.delays) == 1 + + async def test_fn_receives_correct_attempt_number(self): + """First call gets attempt=0, first retry gets attempt=1, etc.""" + received_attempts: list[int] = [] + fake_sleep = _FakeSleep() + + async def fn(attempt): + received_attempts.append(attempt) + if len(received_attempts) < 4: + raise ValueError("not yet") + + await retry(fn, initial_delay=1.0, sleep_fn=fake_sleep) + + assert received_attempts == [0, 1, 2, 3] + + +class TestRetryLogging: + async def test_logs_on_retry(self, caplog): + call_count = 0 + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count < 3: + raise ValueError("boom") + + with caplog.at_level("WARNING"): + await retry(fn, initial_delay=1.0, sleep_fn=_FakeSleep()) + + retry_messages = [r for r in caplog.records if "retrying" in r.message] + assert len(retry_messages) == 2 + + async def test_no_log_on_first_success(self, caplog): + async def fn(_attempt): + pass + + with caplog.at_level("WARNING"): + await retry(fn, sleep_fn=_FakeSleep()) + + assert not any("retrying" in r.message for r in caplog.records) + + async def test_logs_include_exc_info(self, caplog): + call_count = 0 + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count < 2: + raise ValueError("detail") + + with caplog.at_level("WARNING"): + await retry(fn, initial_delay=1.0, sleep_fn=_FakeSleep()) + + retry_records = [r for r in caplog.records if "retrying" in r.message] + assert len(retry_records) == 1 + assert retry_records[0].exc_info is not None + + async def test_log_message_includes_delay(self, caplog): + call_count = 0 + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count < 2: + raise RuntimeError("fail") + + with caplog.at_level("WARNING"): + await retry(fn, initial_delay=2.5, sleep_fn=_FakeSleep()) + + retry_records = [r for r in caplog.records if "retrying" in r.message] + assert len(retry_records) == 1 + assert "2.5s" in retry_records[0].message + + +class TestRetryMaxAttempts: + async def test_raises_after_max_attempts(self): + """The last exception propagates once max_attempts calls have all failed.""" + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + raise ValueError(f"fail {call_count}") + + with pytest.raises(ValueError, match="fail 3"): + await retry(fn, initial_delay=1.0, sleep_fn=fake_sleep, max_attempts=3) + + assert call_count == 3 + assert len(fake_sleep.delays) == 2 + + async def test_succeeds_on_last_allowed_attempt(self): + """No exception when fn succeeds exactly at the max_attempts-th call.""" + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count < 3: + raise RuntimeError("not yet") + + await retry(fn, initial_delay=1.0, sleep_fn=fake_sleep, max_attempts=3) + + assert call_count == 3 + + async def test_max_attempts_one_never_retries(self): + """max_attempts=1 means a single call with no retry.""" + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + raise RuntimeError("fail") + + with pytest.raises(RuntimeError): + await retry(fn, sleep_fn=fake_sleep, max_attempts=1) + + assert call_count == 1 + assert fake_sleep.delays == [] + + async def test_default_is_unlimited(self): + """Without max_attempts, retry keeps going far beyond any small cap.""" + call_count = 0 + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count < 50: + raise RuntimeError("fail") + + await retry(fn, initial_delay=0.0, sleep_fn=_FakeSleep()) + + assert call_count == 50 + + async def test_invalid_max_attempts_rejected(self): + """max_attempts below 1 is a programming error.""" + + async def fn(_attempt): + pass + + with pytest.raises(AssertionError): + await retry(fn, sleep_fn=_FakeSleep(), max_attempts=0) + + async def test_gives_up_log_message(self, caplog): + """The final failure logs a giving-up warning instead of a retrying one.""" + + async def fn(_attempt): + raise RuntimeError("fail") + + with caplog.at_level("WARNING"): + with pytest.raises(RuntimeError): + await retry(fn, initial_delay=1.0, sleep_fn=_FakeSleep(), max_attempts=2) + + assert any("giving up" in r.message for r in caplog.records) + assert len([r for r in caplog.records if "retrying" in r.message]) == 1 + + +class TestRetryBackoff: + async def test_delay_doubles_each_retry(self): + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count <= 4: + raise RuntimeError("fail") + + await retry(fn, initial_delay=1.0, max_delay=100.0, backoff_factor=2.0, sleep_fn=fake_sleep) + + assert fake_sleep.delays == [1.0, 2.0, 4.0, 8.0] + + async def test_delay_capped_at_max(self): + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count <= 5: + raise RuntimeError("fail") + + await retry(fn, initial_delay=1.0, max_delay=3.0, backoff_factor=2.0, sleep_fn=fake_sleep) + + assert fake_sleep.delays == [1.0, 2.0, 3.0, 3.0, 3.0] + + async def test_custom_backoff_factor(self): + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count <= 3: + raise RuntimeError("fail") + + await retry(fn, initial_delay=1.0, max_delay=100.0, backoff_factor=3.0, sleep_fn=fake_sleep) + + assert fake_sleep.delays == [1.0, 3.0, 9.0] + + async def test_zero_initial_delay(self): + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count < 3: + raise RuntimeError("fail") + + await retry(fn, initial_delay=0.0, sleep_fn=fake_sleep) + + assert call_count == 3 + assert fake_sleep.delays == [0.0, 0.0] + + async def test_default_params_are_reasonable(self): + from miles.utils.retry_utils import _DEFAULT_BACKOFF_FACTOR, _DEFAULT_INITIAL_DELAY, _DEFAULT_MAX_DELAY + + assert _DEFAULT_INITIAL_DELAY == 1.0 + assert _DEFAULT_MAX_DELAY == 60.0 + assert _DEFAULT_BACKOFF_FACTOR == 2.0 + + async def test_many_retries_stay_capped(self): + """After hitting max_delay, all subsequent delays remain at max.""" + call_count = 0 + fake_sleep = _FakeSleep() + + async def fn(_attempt): + nonlocal call_count + call_count += 1 + if call_count <= 8: + raise RuntimeError("fail") + + await retry(fn, initial_delay=1.0, max_delay=5.0, backoff_factor=2.0, sleep_fn=fake_sleep) + + # 1, 2, 4, 5, 5, 5, 5, 5 + assert fake_sleep.delays == [1.0, 2.0, 4.0, 5.0, 5.0, 5.0, 5.0, 5.0] From 8dc8ddd9cc49d993e48f10d78c7d8fd606790286 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 06/14] Add structured logfmt logging helper Add a `log_structured` helper that emits logfmt-style key/value log lines, used by the fault-tolerance components for greppable structured logs. - structured_log.py (+ test). --- miles/utils/structured_log.py | 125 ++++++++++++ tests/fast/utils/test_logging_utils.py | 2 +- tests/fast/utils/test_structured_log.py | 255 ++++++++++++++++++++++++ 3 files changed, 381 insertions(+), 1 deletion(-) create mode 100644 miles/utils/structured_log.py create mode 100644 tests/fast/utils/test_structured_log.py diff --git a/miles/utils/structured_log.py b/miles/utils/structured_log.py new file mode 100644 index 00000000000..b74849afbd6 --- /dev/null +++ b/miles/utils/structured_log.py @@ -0,0 +1,125 @@ +import asyncio +import functools +import inspect +import json +import logging +import time +from collections.abc import Callable +from typing import Any + +_PRUNE_CAP = 160 + + +def log_structured(log_fn: Callable[..., None], *, exc_info: bool = False, **fields: Any) -> None: + log_fn("ft " + _to_logfmt(fields), stacklevel=2, exc_info=exc_info) + + +def with_logs(func: Callable[..., Any]) -> Callable[..., Any]: + fn_name = func.__name__ + method_logger = logging.getLogger(func.__module__) + + def log_start(args: tuple[Any, ...]) -> tuple[str, float]: + cls = type(args[0]).__name__ if args else "" + log_structured(method_logger.info, cls=cls, fn=fn_name, phase="start") + return cls, time.monotonic() + + def log_end(cls: str, start: float) -> None: + log_structured(method_logger.info, cls=cls, fn=fn_name, phase="end", ok=True, elapsed_s=_elapsed(start)) + + def log_fail(cls: str, start: float) -> None: + log_structured( + method_logger.error, cls=cls, fn=fn_name, phase="end", ok=False, elapsed_s=_elapsed(start), exc_info=True + ) + + def log_cancelled(cls: str, start: float) -> None: + log_structured( + method_logger.info, cls=cls, fn=fn_name, phase="end", ok=False, elapsed_s=_elapsed(start), cancelled=True + ) + + if inspect.iscoroutinefunction(func): + + @functools.wraps(func) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + cls, start = log_start(args) + try: + result = await func(*args, **kwargs) + except asyncio.CancelledError: + log_cancelled(cls, start) + raise + except BaseException: + log_fail(cls, start) + raise + log_end(cls, start) + return result + + return async_wrapper + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + cls, start = log_start(args) + try: + result = func(*args, **kwargs) + except BaseException: + log_fail(cls, start) + raise + log_end(cls, start) + return result + + return wrapper + + +def _elapsed(start: float) -> float: + return round(time.monotonic() - start, 1) + + +def _to_logfmt(fields: dict[str, Any]) -> str: + return " ".join(f"{key}={_format_value(value)}" for key, value in fields.items()) + + +def _format_value(value: Any) -> str: + if value is None: + return "" + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, (int, float)): + return str(value) + if isinstance(value, (list, tuple)): + return _maybe_quote(",".join(_format_scalar(item) for item in value)) + if isinstance(value, dict): + return _quote(json.dumps(value, separators=(",", ":"), default=str)) + return _maybe_quote(str(value)) + + +def _format_scalar(value: Any) -> str: + if isinstance(value, bool): + return "true" if value else "false" + return str(value) + + +def _maybe_quote(text: str) -> str: + if text and any(ch in text for ch in (" ", "=", '"', "\\", "\n", "\r", "\t")): + return _quote(text) + return text + + +def _quote(text: str) -> str: + escaped = ( + text.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n").replace("\r", "\\r").replace("\t", "\\t") + ) + return f'"{escaped}"' + + +def prune_for_log(value: Any, cap: int = _PRUNE_CAP) -> Any: + if len(_compact_json(value)) <= cap: + return value + if isinstance(value, dict): + return {key: prune_for_log(item, cap) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return f"" + if isinstance(value, str): + return f"" + return f"<{type(value).__name__}>" + + +def _compact_json(value: Any) -> str: + return json.dumps(value, separators=(",", ":"), default=str) diff --git a/tests/fast/utils/test_logging_utils.py b/tests/fast/utils/test_logging_utils.py index b346d372fdd..0bfbc713d55 100644 --- a/tests/fast/utils/test_logging_utils.py +++ b/tests/fast/utils/test_logging_utils.py @@ -30,7 +30,7 @@ def _run_snippet(code: str) -> subprocess.CompletedProcess: [sys.executable, "-c", textwrap.dedent(code)], capture_output=True, text=True, - timeout=10, + timeout=60, ) diff --git a/tests/fast/utils/test_structured_log.py b/tests/fast/utils/test_structured_log.py new file mode 100644 index 00000000000..c712ea6e459 --- /dev/null +++ b/tests/fast/utils/test_structured_log.py @@ -0,0 +1,255 @@ +import asyncio +import logging + +import pytest + +from miles.utils.structured_log import _format_value, _to_logfmt, log_structured, prune_for_log, with_logs + + +class TestFormatValue: + def test_int_renders_bare(self): + """An int renders as a bare number.""" + assert _format_value(7) == "7" + + def test_float_renders_bare(self): + """A float renders as a bare number.""" + assert _format_value(49.3) == "49.3" + + def test_true_and_false_render_lowercase(self): + """Bools render as lowercase true/false (and take precedence over the int branch).""" + assert _format_value(True) == "true" + assert _format_value(False) == "false" + + def test_none_renders_as_empty(self): + """None renders as an empty value.""" + assert _format_value(None) == "" + + def test_empty_string_renders_as_empty(self): + """An empty string renders as an empty value with no quotes.""" + assert _format_value("") == "" + + def test_plain_string_renders_bare(self): + """A string with no space, '=' or quote renders bare.""" + assert _format_value("train") == "train" + + def test_string_with_space_is_quoted(self): + """A string containing a space is double-quoted so it stays one token.""" + assert _format_value("survivors normal") == '"survivors normal"' + + def test_string_with_equals_is_quoted(self): + """A string containing '=' is quoted so it can't be mis-read as a new key.""" + assert _format_value("a=b") == '"a=b"' + + def test_string_with_quote_is_escaped(self): + """A string containing a double-quote is quoted with the inner quote escaped.""" + assert _format_value('say "hi"') == '"say \\"hi\\""' + + def test_string_with_backslash_and_space_escapes_backslash(self): + """When quoting is triggered, backslashes are escaped too.""" + assert _format_value("a\\b c") == '"a\\\\b c"' + + def test_string_with_newline_is_quoted_and_escaped(self): + """A literal newline is escaped so the entry stays on one log line.""" + assert _format_value("a\nb") == '"a\\nb"' + + def test_string_with_carriage_return_is_quoted_and_escaped(self): + """A carriage return is escaped so it can't corrupt the log stream.""" + assert _format_value("a\rb") == '"a\\rb"' + + def test_string_with_tab_is_quoted_and_escaped(self): + """A tab is escaped so the value stays one unambiguous token.""" + assert _format_value("a\tb") == '"a\\tb"' + + def test_string_with_only_backslash_is_quoted_and_escaped(self): + """A backslash alone (no space) still triggers quoting so it round-trips.""" + assert _format_value("a\\b") == '"a\\\\b"' + + def test_string_mixing_all_control_chars(self): + """Newline, quote, tab and backslash all escape together.""" + assert _format_value('a\n"b"\t\\c\r') == '"a\\n\\"b\\"\\t\\\\c\\r"' + + def test_backslash_escaped_before_other_escapes(self): + """A pre-escaped-looking string round-trips: the backslash doubles, the newline escapes.""" + assert _format_value("\\n\n") == '"\\\\n\\n"' + + def test_list_element_with_newline_is_quoted(self): + """A newline inside a list element quotes and escapes the joined value.""" + assert _format_value(["a\nb", "c"]) == '"a\\nb,c"' + + def test_list_of_ints_is_comma_joined_no_space(self): + """A list renders comma-joined with no spaces.""" + assert _format_value([0, 1, 2]) == "0,1,2" + + def test_list_of_bools_lowercased(self): + """List elements that are bools also render lowercase.""" + assert _format_value([True, False]) == "true,false" + + def test_empty_list_renders_as_empty(self): + """An empty list renders as an empty value.""" + assert _format_value([]) == "" + + def test_tuple_behaves_like_list(self): + """A tuple renders the same as the equivalent list.""" + assert _format_value((0, 1)) == "0,1" + + def test_list_whose_joined_form_has_a_space_is_quoted(self): + """If a list element introduces a space, the whole joined value is quoted.""" + assert _format_value(["a b", "c"]) == '"a b,c"' + + def test_dict_renders_as_quoted_compact_json(self): + """A dict value renders as quoted compact JSON.""" + assert _format_value({"a": 1}) == '"{\\"a\\":1}"' + + +class TestToLogfmt: + def test_fields_join_with_space_in_insertion_order(self): + """Fields render as space-separated key=value in insertion order.""" + assert _to_logfmt({"cell": 1, "fn": "train", "ok": True}) == "cell=1 fn=train ok=true" + + def test_empty_fields_render_as_empty_string(self): + """No fields renders as an empty string.""" + assert _to_logfmt({}) == "" + + def test_list_and_empty_fields_keep_one_token_each(self): + """A list field stays one token; an empty field renders as a trailing 'key='.""" + assert _to_logfmt({"alive": [0, 1], "pending": []}) == "alive=0,1 pending=" + + +class TestPruneForLog: + def test_small_payload_kept_verbatim(self): + """A payload under the cap is returned unchanged.""" + payload = {"quorum_id": 1, "healed": [0]} + assert prune_for_log(payload, cap=160) == payload + + def test_large_list_field_summarized_small_siblings_kept(self): + """An oversized list field becomes a length summary; small siblings stay inline.""" + pruned = prune_for_log({"quorum_id": 1, "checksums": list(range(500))}, cap=80) + assert pruned == {"quorum_id": 1, "checksums": ""} + + def test_large_string_field_summarized(self): + """An oversized string field becomes a char-count summary.""" + assert prune_for_log({"blob": "x" * 1000}, cap=80) == {"blob": ""} + + def test_nested_dict_prunes_only_the_big_subfield(self): + """Recursion summarizes only the oversized nested field, keeping small ones.""" + pruned = prune_for_log({"a": 1, "b": {"big": list(range(500)), "small": 2}}, cap=60) + assert pruned == {"a": 1, "b": {"big": "", "small": 2}} + + def test_large_tuple_summarized_as_list(self): + """An oversized tuple is summarized like a list.""" + assert prune_for_log(tuple(range(500)), cap=80) == "" + + def test_value_at_cap_kept_just_over_summarized(self): + """A value whose compact-JSON length is <= cap is kept; one char over is summarized.""" + assert prune_for_log("x" * 8, cap=10) == "x" * 8 # json '"xxxxxxxx"' is exactly 10 chars + assert prune_for_log("x" * 9, cap=10) == "" # json is 11 chars + + +class TestLogStructured: + def test_emits_ft_prefixed_logfmt_via_given_method(self, caplog): + """log_structured emits one 'ft '-prefixed logfmt line through the passed logger method.""" + logger = logging.getLogger("t_emit") + with caplog.at_level(logging.INFO, logger="t_emit"): + log_structured(logger.info, op="execute", phase="start", cell=1, fn="train") + assert caplog.messages == ["ft op=execute phase=start cell=1 fn=train"] + + def test_uses_the_level_of_the_passed_method(self, caplog): + """The record's level is that of the bound method passed (warning here).""" + logger = logging.getLogger("t_level") + with caplog.at_level(logging.DEBUG, logger="t_level"): + log_structured(logger.warning, op="x") + assert caplog.records[0].levelno == logging.WARNING + + def test_exc_info_is_forwarded(self, caplog): + """exc_info=True attaches the active exception to the record.""" + logger = logging.getLogger("t_exc") + with caplog.at_level(logging.ERROR, logger="t_exc"): + try: + raise ValueError("boom") + except ValueError: + log_structured(logger.error, op="x", exc_info=True) + assert caplog.records[0].exc_info[0] is ValueError + + def test_stacklevel_points_to_caller_not_helper(self, caplog): + """stacklevel=2 makes the record's filename the caller's, not structured_log.py.""" + logger = logging.getLogger("t_stack") + with caplog.at_level(logging.INFO, logger="t_stack"): + log_structured(logger.info, op="x") + assert caplog.records[0].filename == "test_structured_log.py" + + +class FakeActor: + @with_logs + def train(self): + return 42 + + @with_logs + def update_weights(self): + raise ValueError("boom") + + @with_logs + async def send_ckpt(self): + return "done" + + @with_logs + async def wait_forever(self): + await asyncio.Event().wait() + + +class TestWithLogs: + def test_sync_method_emits_start_then_end_with_class_and_method(self, caplog): + """A sync call emits phase=start then phase=end carrying the auto-detected class and method.""" + with caplog.at_level(logging.INFO): + assert FakeActor().train() == 42 + assert caplog.messages[0] == "ft cls=FakeActor fn=train phase=start" + assert caplog.messages[1].startswith("ft cls=FakeActor fn=train phase=end ok=true elapsed_s=") + + def test_class_is_read_from_the_runtime_instance(self, caplog): + """The cls field reflects the actual runtime instance, not where the method was defined.""" + + class SubActor(FakeActor): + pass + + with caplog.at_level(logging.INFO): + SubActor().train() + assert caplog.messages[0] == "ft cls=SubActor fn=train phase=start" + + def test_exception_logs_end_not_ok_with_exc_info_and_reraises(self, caplog): + """On exception the end line is ok=false with the traceback, and the error propagates.""" + with caplog.at_level(logging.INFO): + with pytest.raises(ValueError, match="boom"): + FakeActor().update_weights() + assert caplog.messages[0] == "ft cls=FakeActor fn=update_weights phase=start" + end = caplog.records[1] + assert end.getMessage().startswith("ft cls=FakeActor fn=update_weights phase=end ok=false elapsed_s=") + assert end.levelno == logging.ERROR + assert end.exc_info[0] is ValueError + + def test_async_method_emits_start_then_end_ok(self, caplog): + """An async call emits the same start/end pair after the coroutine completes.""" + with caplog.at_level(logging.INFO): + assert asyncio.run(FakeActor().send_ckpt()) == "done" + assert caplog.messages[0] == "ft cls=FakeActor fn=send_ckpt phase=start" + assert caplog.messages[1].startswith("ft cls=FakeActor fn=send_ckpt phase=end ok=true elapsed_s=") + + def test_async_cancellation_logs_cancelled_at_info_and_reraises(self, caplog): + """Cancelling an async call emits an info end line with cancelled=true and no traceback.""" + + async def run() -> None: + task = asyncio.ensure_future(FakeActor().wait_forever()) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + with caplog.at_level(logging.INFO): + asyncio.run(run()) + end = caplog.records[1] + assert end.getMessage().startswith("ft cls=FakeActor fn=wait_forever phase=end ok=false elapsed_s=") + assert end.getMessage().endswith("cancelled=true") + assert end.levelno == logging.INFO + assert not end.exc_info + + def test_preserves_wrapped_function_metadata(self): + """functools.wraps keeps the original __name__ so Ray/introspection still sees it.""" + assert FakeActor.train.__name__ == "train" From 96d220f35ade97e33a58aa4539270c96ce6b9626 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 07/14] Add a Clock abstraction with a fake clock for tests Add a small `Clock` interface (`RealClock` plus a controllable fake clock) so time-dependent fault-tolerance code (health checks, heartbeats) can be driven deterministically in tests. - miles/utils/clock.py and tests. --- miles/utils/clock.py | 80 +++++++++++ tests/fast/utils/test_clock.py | 256 +++++++++++++++++++++++++++++++++ 2 files changed, 336 insertions(+) create mode 100644 miles/utils/clock.py create mode 100644 tests/fast/utils/test_clock.py diff --git a/miles/utils/clock.py b/miles/utils/clock.py new file mode 100644 index 00000000000..73986d8d2e8 --- /dev/null +++ b/miles/utils/clock.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import abc +import asyncio +import heapq +import time +from typing import NamedTuple + + +class Clock(abc.ABC): + @abc.abstractmethod + def time(self) -> float: ... + + @abc.abstractmethod + async def sleep(self, seconds: float) -> None: ... + + +class RealClock(Clock): + def time(self) -> float: + return time.time() + + async def sleep(self, seconds: float) -> None: + await asyncio.sleep(seconds) + + +_DRAIN_ITERATIONS = 20 + + +class _Waiter(NamedTuple): + target: float + seq: int + future: asyncio.Future[None] + + +class FakeClock(Clock): + """Deterministic clock for testing async time-dependent code. + + ``sleep()`` suspends the caller until ``elapse()`` advances the clock + past the target time. This gives tests precise control over which + sleeps resolve and when. + """ + + def __init__(self, start: float = 0.0) -> None: + self._now = start + self._waiters: list[_Waiter] = [] + self._counter: int = 0 + + def time(self) -> float: + return self._now + + async def sleep(self, seconds: float) -> None: + if seconds < 0: + return + + target = self._now + seconds + future: asyncio.Future[None] = asyncio.get_running_loop().create_future() + self._counter += 1 + heapq.heappush(self._waiters, _Waiter(target=target, seq=self._counter, future=future)) + self._resolve_ready() + await future + + async def elapse(self, seconds: float) -> None: + assert seconds >= 0, f"Cannot elapse negative time: {seconds}" + self._now += seconds + self._resolve_ready() + # Drain: yield enough times for resolved coroutines to run through + # sync code and register their next sleep. + for _ in range(_DRAIN_ITERATIONS): + await asyncio.sleep(0) + self._resolve_ready() + + def _resolve_ready(self) -> None: + while self._waiters and self._waiters[0].target <= self._now: + waiter = heapq.heappop(self._waiters) + if not waiter.future.done(): + waiter.future.set_result(None) + + @property + def pending_count(self) -> int: + return sum(1 for w in self._waiters if not w.future.done()) diff --git a/tests/fast/utils/test_clock.py b/tests/fast/utils/test_clock.py new file mode 100644 index 00000000000..7837f6d5f33 --- /dev/null +++ b/tests/fast/utils/test_clock.py @@ -0,0 +1,256 @@ +import asyncio + +import pytest + +from miles.utils.clock import FakeClock, RealClock + + +class TestRealClock: + def test_time_returns_current_time(self): + clock = RealClock() + import time + + before = time.time() + result = clock.time() + after = time.time() + assert before <= result <= after + + async def test_sleep_actually_waits(self): + clock = RealClock() + import time + + start = time.time() + await clock.sleep(0.05) + elapsed = time.time() - start + assert elapsed >= 0.04 + + +class TestFakeClockTime: + def test_initial_time(self): + clock = FakeClock(start=100.0) + assert clock.time() == 100.0 + + def test_default_start_is_zero(self): + clock = FakeClock() + assert clock.time() == 0.0 + + async def test_elapse_advances_time(self): + clock = FakeClock(start=10.0) + await clock.elapse(5.0) + assert clock.time() == 15.0 + + async def test_multiple_elapse_accumulate(self): + clock = FakeClock() + await clock.elapse(3.0) + await clock.elapse(7.0) + assert clock.time() == 10.0 + + async def test_elapse_zero_is_allowed(self): + clock = FakeClock(start=5.0) + await clock.elapse(0.0) + assert clock.time() == 5.0 + + async def test_elapse_negative_raises(self): + clock = FakeClock() + with pytest.raises(AssertionError, match="negative"): + await clock.elapse(-1.0) + + +class TestFakeClockSleep: + async def test_sleep_zero_returns_immediately(self): + clock = FakeClock() + completed = False + + async def task() -> None: + nonlocal completed + await clock.sleep(0) + completed = True + + asyncio.create_task(task()) + await asyncio.sleep(0) + assert completed + + async def test_sleep_blocks_until_elapse(self): + clock = FakeClock() + completed = False + + async def task() -> None: + nonlocal completed + await clock.sleep(10.0) + completed = True + + asyncio.create_task(task()) + await asyncio.sleep(0) + assert not completed + + await clock.elapse(10.0) + assert completed + + async def test_sleep_does_not_resolve_before_target(self): + clock = FakeClock() + completed = False + + async def task() -> None: + nonlocal completed + await clock.sleep(10.0) + completed = True + + asyncio.create_task(task()) + await asyncio.sleep(0) + + await clock.elapse(9.9) + assert not completed + + await clock.elapse(0.1) + assert completed + + async def test_multiple_sleeps_resolve_in_order(self): + clock = FakeClock() + order: list[str] = [] + + async def task_a() -> None: + await clock.sleep(5.0) + order.append("a") + + async def task_b() -> None: + await clock.sleep(10.0) + order.append("b") + + async def task_c() -> None: + await clock.sleep(3.0) + order.append("c") + + asyncio.create_task(task_a()) + asyncio.create_task(task_b()) + asyncio.create_task(task_c()) + await asyncio.sleep(0) + + await clock.elapse(5.0) + assert order == ["c", "a"] + + await clock.elapse(5.0) + assert order == ["c", "a", "b"] + + async def test_same_target_time_all_resolve(self): + clock = FakeClock() + count = 0 + + async def task() -> None: + nonlocal count + await clock.sleep(5.0) + count += 1 + + for _ in range(3): + asyncio.create_task(task()) + await asyncio.sleep(0) + + await clock.elapse(5.0) + assert count == 3 + + async def test_overshoot_resolves_all_pending(self): + clock = FakeClock() + order: list[int] = [] + + for delay in [1, 5, 10]: + d = delay + + async def task(d: int = d) -> None: + await clock.sleep(d) + order.append(d) + + asyncio.create_task(task()) + await asyncio.sleep(0) + + await clock.elapse(100.0) + assert sorted(order) == [1, 5, 10] + + async def test_sleep_negative_returns_immediately(self): + clock = FakeClock() + completed = False + + async def task() -> None: + nonlocal completed + await clock.sleep(-1.0) + completed = True + + asyncio.create_task(task()) + await asyncio.sleep(0) + assert completed + + +class TestFakeClockPendingCount: + async def test_no_pending_initially(self): + clock = FakeClock() + assert clock.pending_count == 0 + + async def test_pending_count_tracks_sleeps(self): + clock = FakeClock() + + asyncio.create_task(clock.sleep(5.0)) + asyncio.create_task(clock.sleep(10.0)) + await asyncio.sleep(0) + assert clock.pending_count == 2 + + await clock.elapse(5.0) + assert clock.pending_count == 1 + + await clock.elapse(5.0) + assert clock.pending_count == 0 + + +class TestFakeClockChainedSleeps: + async def test_sequential_sleeps_in_coroutine(self): + """A coroutine that does multiple sleeps in sequence.""" + clock = FakeClock() + checkpoints: list[float] = [] + + async def task() -> None: + checkpoints.append(clock.time()) + await clock.sleep(5.0) + checkpoints.append(clock.time()) + await clock.sleep(3.0) + checkpoints.append(clock.time()) + + asyncio.create_task(task()) + await asyncio.sleep(0) + + assert checkpoints == [0.0] + + await clock.elapse(5.0) + assert checkpoints == [0.0, 5.0] + + await clock.elapse(3.0) + assert checkpoints == [0.0, 5.0, 8.0] + + async def test_periodic_loop(self): + """Simulates a periodic loop like SimpleHealthChecker._loop.""" + clock = FakeClock() + ticks: list[float] = [] + + async def loop() -> None: + await clock.sleep(10.0) + while len(ticks) < 3: + ticks.append(clock.time()) + await clock.sleep(5.0) + + task = asyncio.create_task(loop()) + await asyncio.sleep(0) + + # Step 1: first_wait=10 + assert ticks == [] + await clock.elapse(10.0) + assert ticks == [10.0] + + # Step 2: interval=5 + await clock.elapse(5.0) + assert ticks == [10.0, 15.0] + + # Step 3: interval=5 + await clock.elapse(5.0) + assert ticks == [10.0, 15.0, 20.0] + + task.cancel() + try: + await task + except asyncio.CancelledError: + pass From a005a2f7a7fd259d741a725b58d6b4a82ec1d73b Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 08/14] Add a fault injector test utility Add a fault-injector test utility used to deterministically exercise fault-tolerance code paths. - miles/utils/test_utils/fault_injector.py. --- miles/utils/test_utils/fault_injector.py | 41 ++++++++++++++++++++++++ 1 file changed, 41 insertions(+) create mode 100644 miles/utils/test_utils/fault_injector.py diff --git a/miles/utils/test_utils/fault_injector.py b/miles/utils/test_utils/fault_injector.py new file mode 100644 index 00000000000..35362597f08 --- /dev/null +++ b/miles/utils/test_utils/fault_injector.py @@ -0,0 +1,41 @@ +""" +Failure modes modeled after torchft's failure.py: +https://github.com/meta-pytorch/torchft/blob/main/examples/monarch/utils/failure.py +""" + +import ctypes +import logging +import os +import signal +from enum import Enum + +logger = logging.getLogger(__name__) + + +class FailureMode(Enum): + SIGKILL = "sigkill" + EXIT = "exit" + SEGFAULT = "segfault" + DEADLOCK = "deadlock" + + +def inject_fault(mode: str) -> None: + failure_mode = FailureMode(mode) + logger.warning("FaultInjector: executing %s (pid=%d)", failure_mode.value, os.getpid()) + + match failure_mode: + case FailureMode.SIGKILL: + os.kill(os.getpid(), signal.SIGKILL) + + case FailureMode.EXIT: + os._exit(1) + + case FailureMode.SEGFAULT: + crash_func = ctypes.CFUNCTYPE(None)() + crash_func() + + case FailureMode.DEADLOCK: + libc = ctypes.PyDLL(None) + libc.sleep.argtypes = (ctypes.c_uint,) + libc.sleep.restype = ctypes.c_uint + libc.sleep(600) From 836597152d198e3676894d1612d00d05d898bfd0 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 09/14] Add control-server data models Add the shared data models for the fault-tolerance control server (e.g. the `TriState` health value), used by the health checker and later by the HTTP control server. - miles/utils/control_server/models.py. --- miles/utils/control_server/__init__.py | 0 miles/utils/control_server/models.py | 85 +++++++++++++++++++++ tests/fast/utils/control_server/__init__.py | 0 3 files changed, 85 insertions(+) create mode 100644 miles/utils/control_server/__init__.py create mode 100644 miles/utils/control_server/models.py create mode 100644 tests/fast/utils/control_server/__init__.py diff --git a/miles/utils/control_server/__init__.py b/miles/utils/control_server/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/miles/utils/control_server/models.py b/miles/utils/control_server/models.py new file mode 100644 index 00000000000..00e47aadac8 --- /dev/null +++ b/miles/utils/control_server/models.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +from enum import StrEnum +from typing import Literal + +from miles.utils.pydantic_utils import StrictBaseModel +from miles.utils.test_utils.fault_injector import FailureMode + + +class TriState(StrEnum): + """K8s condition status: ``"True"``, ``"False"``, or ``"Unknown"``.""" + + TRUE = "True" + FALSE = "False" + UNKNOWN = "Unknown" + + +class _OkResponse(StrictBaseModel): + status: str = "ok" + + +class CellCondition(StrictBaseModel): + type: Literal["Allocated", "Healthy"] + status: TriState + reason: str | None = None + message: str | None = None + lastTransitionTime: str | None = None + + @classmethod + def allocated(cls, status: TriState) -> CellCondition: + return cls(type="Allocated", status=status) + + @classmethod + def healthy(cls, status: TriState, *, reason: str | None = None) -> CellCondition: + return cls(type="Healthy", status=status, reason=reason) + + +class CellStatus(StrictBaseModel): + phase: Literal["Pending", "Running", "Suspended"] + conditions: list[CellCondition] + + +class CellSpec(StrictBaseModel): + suspend: bool = False + + +class CellMetadata(StrictBaseModel): + name: str + labels: dict[str, str] + + +class Cell(StrictBaseModel): + apiVersion: Literal["miles.io/v1"] = "miles.io/v1" + kind: Literal["Cell"] = "Cell" + metadata: CellMetadata + spec: CellSpec + status: CellStatus + + +class CellList(StrictBaseModel): + apiVersion: Literal["miles.io/v1"] = "miles.io/v1" + kind: Literal["CellList"] = "CellList" + items: list[Cell] + + +class CellPatchSpec(StrictBaseModel): + suspend: bool | None = None + + +class CellPatch(StrictBaseModel): + spec: CellPatchSpec | None = None + + +class FaultInjection(StrictBaseModel): + mode: FailureMode + sub_index: int = 0 + + +class K8sStatus(StrictBaseModel): + apiVersion: Literal["v1"] = "v1" + kind: Literal["Status"] = "Status" + status: Literal["Failure"] = "Failure" + message: str + reason: str + code: int diff --git a/tests/fast/utils/control_server/__init__.py b/tests/fast/utils/control_server/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From e308821cf0580bd62e637816319951d6f3a4e536 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 10/14] Add a cell health checker and heartbeat utilities Add the periodic health checker (debounced TriState status driven by a Clock) and heartbeat utilities used to monitor train-cell liveness. - miles/utils/health_checker.py, miles/utils/heartbeat_utils.py and tests. --- miles/utils/health_checker.py | 247 ++++++++++++ miles/utils/heartbeat_utils.py | 22 + tests/fast/utils/test_health_checker.py | 493 +++++++++++++++++++++++ tests/fast/utils/test_heartbeat_utils.py | 111 +++++ 4 files changed, 873 insertions(+) create mode 100644 miles/utils/health_checker.py create mode 100644 miles/utils/heartbeat_utils.py create mode 100644 tests/fast/utils/test_health_checker.py create mode 100644 tests/fast/utils/test_heartbeat_utils.py diff --git a/miles/utils/health_checker.py b/miles/utils/health_checker.py new file mode 100644 index 00000000000..4d87435d060 --- /dev/null +++ b/miles/utils/health_checker.py @@ -0,0 +1,247 @@ +from __future__ import annotations + +import abc +import argparse +import asyncio +import logging +from collections.abc import Callable, Coroutine +from typing import Any + +from miles.utils.clock import Clock, RealClock +from miles.utils.control_server.models import TriState +from miles.utils.pydantic_utils import StrictBaseModel +from miles.utils.structured_log import log_structured + +logger = logging.getLogger(__name__) + + +class SimpleHealthCheckerConfig(StrictBaseModel): + interval: float + timeout: float + first_wait: float + failure_threshold: int + + @staticmethod + def add_arguments(parser: argparse.ArgumentParser, *, prefix: str) -> None: + parser.add_argument( + f"--{prefix}-interval", + type=float, + default=10.0, + help=f"Interval in seconds between {prefix} health checks.", + ) + parser.add_argument( + f"--{prefix}-timeout", + type=float, + default=10.0, + help=f"Timeout in seconds for a single {prefix} health check RPC.", + ) + parser.add_argument( + f"--{prefix}-first-wait", + type=float, + default=300.0, + help=f"Initial grace period (seconds) before starting {prefix} health checks.", + ) + parser.add_argument( + f"--{prefix}-failure-threshold", + type=int, + default=3, + help=( + f"Number of consecutive failed {prefix} checks before reporting unhealthy. " + "Debounces transient RPC blips so a single hiccup does not recycle a live cell." + ), + ) + + @staticmethod + def from_args(args: object, *, prefix: str) -> SimpleHealthCheckerConfig: + attr_prefix = prefix.replace("-", "_") + return SimpleHealthCheckerConfig( + interval=getattr(args, f"{attr_prefix}_interval"), + timeout=getattr(args, f"{attr_prefix}_timeout"), + first_wait=getattr(args, f"{attr_prefix}_first_wait"), + failure_threshold=getattr(args, f"{attr_prefix}_failure_threshold"), + ) + + +class BaseHealthChecker(abc.ABC): + @property + @abc.abstractmethod + def status(self) -> TriState: ... + + @abc.abstractmethod + async def start(self) -> None: ... + + @abc.abstractmethod + def stop(self) -> None: ... + + @abc.abstractmethod + def pause(self) -> None: ... + + @abc.abstractmethod + def resume(self) -> None: ... + + +class SimpleHealthChecker(BaseHealthChecker): + """Periodic async health checker. Calls *check_fn*; reports result via *on_result*. + + After each ``resume()``, waits ``first_wait`` seconds before the first check + (matching ``RolloutHealthMonitor._need_first_wait`` semantics). + """ + + def __init__( + self, + *, + name: str, + check_fn: Callable[[], Coroutine[Any, Any, None]], + on_result: Callable[[bool], None] | None = None, + config: SimpleHealthCheckerConfig, + clock: Clock | None = None, + ) -> None: + self._name = name + self._check_fn = check_fn + self._on_result = on_result + self._config = config + self._clock = clock or RealClock() + + self._status = TriState.UNKNOWN + self._paused: bool = False + self._need_first_wait: bool = True + self._consecutive_failures: int = 0 + self._task: asyncio.Task[None] | None = None + + @property + def status(self) -> TriState: + return self._status + + async def start(self) -> None: + if self._task is not None: + return + log_structured(logger.info, op="health", phase="start", name=self._name) + self._task = asyncio.create_task(self._loop()) + await asyncio.sleep(0) + + def stop(self) -> None: + if self._task is not None: + log_structured(logger.info, op="health", phase="stop", name=self._name) + self._task.cancel() + self._task = None + self._status = TriState.UNKNOWN + + def pause(self) -> None: + log_structured(logger.info, op="health", phase="pause", name=self._name) + self._paused = True + self._status = TriState.UNKNOWN + + def resume(self) -> None: + log_structured(logger.info, op="health", phase="resume", name=self._name) + self._paused = False + self._need_first_wait = True + self._status = TriState.UNKNOWN + self._consecutive_failures = 0 + + async def _loop(self) -> None: + while True: + if self._need_first_wait: + self._need_first_wait = False + log_structured( + logger.info, op="health", phase="first_wait", name=self._name, wait_s=self._config.first_wait + ) + await self._clock.sleep(self._config.first_wait) + + if not self._paused: + success = False + try: + await asyncio.wait_for(self._check_fn(), timeout=self._config.timeout) + success = True + except Exception: + log_structured(logger.error, op="health", phase="check_failed", name=self._name, exc_info=True) + + prev_status = self._status + if success: + self._consecutive_failures = 0 + self._status = TriState.TRUE + else: + self._consecutive_failures += 1 + if self._consecutive_failures >= self._config.failure_threshold: + self._status = TriState.FALSE + + log_structured( + logger.info, + op="health", + phase="poll", + name=self._name, + ok=success, + status=self._status.value, + consecutive_failures=self._consecutive_failures, + ) + + if prev_status != self._status: + log_structured( + logger.info, + op="health", + phase="status_change", + name=self._name, + from_status=prev_status.value, + to_status=self._status.value, + consecutive_failures=self._consecutive_failures, + ) + + if self._on_result is not None: + try: + self._on_result(success) + except Exception: + log_structured( + logger.error, op="health", phase="on_result_failed", name=self._name, exc_info=True + ) + + await self._clock.sleep(self._config.interval) + + +class NoopHealthChecker(BaseHealthChecker): + @property + def status(self) -> TriState: + return TriState.UNKNOWN + + async def start(self) -> None: + pass + + def stop(self) -> None: + pass + + def pause(self) -> None: + pass + + def resume(self) -> None: + pass + + +# TODO: should move when Rollout FT is implemented +def create_rollout_cell_health_checker( + *, + cell_id: str, + get_engines: Callable[[], list[Any]], + config: SimpleHealthCheckerConfig, + on_result: Callable[[bool], None] | None = None, +) -> SimpleHealthChecker: + + async def _check() -> None: + engines = get_engines() + if not engines: + raise RuntimeError("No engines") + + lead_engine = engines[0] + if lead_engine is None: + raise RuntimeError("Lead engine is None") + + await lead_engine.health_generate.remote() + + # Preserve the pre-debounce rollout semantics for now: a single failed check + # reports unhealthy immediately. The trainer cell checker uses the default + # failure_threshold; tuning rollout debouncing is left to Rollout FT work. + config = config.model_copy(update={"failure_threshold": 1}) + + return SimpleHealthChecker( + name=f"rollout-cell-{cell_id}", + check_fn=_check, + on_result=on_result, + config=config, + ) diff --git a/miles/utils/heartbeat_utils.py b/miles/utils/heartbeat_utils.py new file mode 100644 index 00000000000..9ae64fcebee --- /dev/null +++ b/miles/utils/heartbeat_utils.py @@ -0,0 +1,22 @@ +import time +from dataclasses import dataclass + + +@dataclass(frozen=True) +class HeartbeatStatus: + last_active_timestamp: float + bump_count: int + + +class SimpleHeartbeat: + def __init__(self) -> None: + self._status = HeartbeatStatus(last_active_timestamp=time.time(), bump_count=0) + + def bump(self) -> None: + self._status = HeartbeatStatus( + last_active_timestamp=time.time(), + bump_count=self._status.bump_count + 1, + ) + + def status(self) -> HeartbeatStatus: + return self._status diff --git a/tests/fast/utils/test_health_checker.py b/tests/fast/utils/test_health_checker.py new file mode 100644 index 00000000000..c403d7d78f3 --- /dev/null +++ b/tests/fast/utils/test_health_checker.py @@ -0,0 +1,493 @@ +import asyncio + + +from miles.utils.clock import FakeClock +from miles.utils.control_server.models import TriState +from miles.utils.health_checker import NoopHealthChecker, SimpleHealthChecker + + +async def _settle(clock: FakeClock) -> None: + for _ in range(1000): + if clock.pending_count >= 1: + return + await asyncio.sleep(0) + + +def _make_checker( + *, + check_fn=None, + on_result=None, + interval: float = 10.0, + timeout: float = 5.0, + first_wait: float = 0.0, + failure_threshold: int = 1, + name: str = "test", + clock: FakeClock | None = None, +) -> tuple[SimpleHealthChecker, FakeClock]: + from miles.utils.health_checker import SimpleHealthCheckerConfig + + if check_fn is None: + + async def check_fn() -> None: + pass + + c = clock or FakeClock() + checker = SimpleHealthChecker( + name=name, + check_fn=check_fn, + on_result=on_result, + config=SimpleHealthCheckerConfig( + interval=interval, timeout=timeout, first_wait=first_wait, failure_threshold=failure_threshold + ), + clock=c, + ) + return checker, c + + +class TestStartStop: + async def test_start_creates_task(self): + checker, _ = _make_checker() + assert checker._task is None + + await checker.start() + assert checker._task is not None + + checker.stop() + assert checker._task is None + + async def test_start_is_idempotent(self): + checker, _ = _make_checker() + await checker.start() + task = checker._task + + await checker.start() + assert checker._task is task + + checker.stop() + + async def test_stop_without_start_is_noop(self): + checker, _ = _make_checker() + checker.stop() + + +class TestCheckFnCalled: + async def test_check_fn_called_after_first_interval(self): + call_count = 0 + + async def check_fn() -> None: + nonlocal call_count + call_count += 1 + + checker, clock = _make_checker(check_fn=check_fn, interval=10.0) + await checker.start() + + # Step 1: first_wait=0, so first check runs immediately after task starts + await _settle(clock) + assert call_count == 1 + + # Step 2: Elapse less than interval — no second check + await clock.elapse(5.0) + assert call_count == 1 + + # Step 3: Elapse to interval — second check + await clock.elapse(5.0) + assert call_count == 2 + + checker.stop() + + async def test_first_wait_delays_first_check(self): + call_count = 0 + + async def check_fn() -> None: + nonlocal call_count + call_count += 1 + + checker, clock = _make_checker(check_fn=check_fn, first_wait=300.0, interval=10.0) + await checker.start() + + # Step 1: Elapse 100s — still in first_wait + await clock.elapse(100.0) + assert call_count == 0 + + # Step 2: Elapse to 300s — first_wait completes, first check runs + await clock.elapse(200.0) + assert call_count == 1 + + # Step 3: Elapse interval — second check + await clock.elapse(10.0) + assert call_count == 2 + + checker.stop() + + +class TestOnResult: + async def test_on_result_true_on_success(self): + results: list[bool] = [] + + checker, clock = _make_checker(on_result=lambda s: results.append(s)) + await checker.start() + await _settle(clock) + checker.stop() + + assert results == [True] + + async def test_on_result_false_on_failure(self): + results: list[bool] = [] + + async def check_fn() -> None: + raise RuntimeError("boom") + + checker, clock = _make_checker(check_fn=check_fn, on_result=lambda s: results.append(s)) + await checker.start() + await _settle(clock) + checker.stop() + + assert results == [False] + + async def test_loop_survives_on_result_raising(self, caplog): + """A raising on_result callback is logged and does not kill the check loop.""" + results: list[bool] = [] + + def on_result(success: bool) -> None: + results.append(success) + raise RuntimeError("callback boom") + + checker, clock = _make_checker(on_result=on_result, interval=5.0) + await checker.start() + + await _settle(clock) + await clock.elapse(5.0) + checker.stop() + + assert results == [True, True] + assert any("on_result_failed" in r.message for r in caplog.records) + + async def test_loop_continues_after_failure(self): + results: list[bool] = [] + + async def check_fn() -> None: + raise RuntimeError("boom") + + checker, clock = _make_checker(check_fn=check_fn, on_result=lambda s: results.append(s), interval=5.0) + await checker.start() + + await _settle(clock) + await clock.elapse(5.0) + checker.stop() + + assert results == [False, False] + + async def test_intermittent_failure(self): + call_count = 0 + results: list[bool] = [] + + async def check_fn() -> None: + nonlocal call_count + call_count += 1 + if call_count % 2 == 0: + raise RuntimeError("intermittent") + + checker, clock = _make_checker(check_fn=check_fn, on_result=lambda s: results.append(s), interval=5.0) + await checker.start() + await _settle(clock) + # first_wait=0 so first check runs immediately on start + assert results == [True] + + for _ in range(3): + await clock.elapse(5.0) + checker.stop() + + assert results == [True, False, True, False] + + +class TestPauseResume: + async def test_paused_checker_does_not_call_check_fn(self): + call_count = 0 + + async def check_fn() -> None: + nonlocal call_count + call_count += 1 + + checker, clock = _make_checker(check_fn=check_fn, interval=5.0) + checker.pause() + + await checker.start() + await clock.elapse(20.0) + checker.stop() + + assert call_count == 0 + + async def test_resume_after_pause_resumes_checking(self): + call_count = 0 + + async def check_fn() -> None: + nonlocal call_count + call_count += 1 + + checker, clock = _make_checker(check_fn=check_fn, interval=5.0) + checker.pause() + + await checker.start() + await clock.elapse(20.0) + assert call_count == 0 + + checker.resume() + await clock.elapse(5.0) + checker.stop() + + assert call_count >= 1 + + async def test_pause_resume_flags(self): + checker, _ = _make_checker() + assert not checker._paused + + checker.pause() + assert checker._paused + + checker.resume() + assert not checker._paused + + +class TestNeedFirstWait: + async def test_resume_triggers_first_wait_again(self): + """After resume, the loop waits first_wait before the next check.""" + call_count = 0 + + async def check_fn() -> None: + nonlocal call_count + call_count += 1 + + checker, clock = _make_checker(check_fn=check_fn, first_wait=100.0, interval=5.0) + await checker.start() + + # Step 1: Initial first_wait (100s) + await clock.elapse(50.0) + assert call_count == 0 + await clock.elapse(50.0) + assert call_count == 1 + + # Step 2: Normal interval (5s) + await clock.elapse(5.0) + assert call_count == 2 + + # Step 3: Pause + resume resets first_wait + checker.pause() + checker.resume() + + # Step 4: Need to elapse past the pending interval sleep first, + # then the new first_wait (100s) before next check + await clock.elapse(5.0) + assert call_count == 2 + + await clock.elapse(50.0) + assert call_count == 2 + + await clock.elapse(50.0) + assert call_count == 3 + + checker.stop() + + async def test_pause_without_resume_no_first_wait(self): + checker, clock = _make_checker(first_wait=300.0) + await checker.start() + await clock.elapse(300.0) + assert checker._need_first_wait is False + + checker.pause() + assert checker._need_first_wait is False + + checker.stop() + + +class TestTriState: + async def test_initial_status_is_unknown(self): + checker, _ = _make_checker() + assert checker.status == TriState.UNKNOWN + + async def test_healthy_after_successful_check(self): + checker, clock = _make_checker() + await checker.start() + await _settle(clock) + + assert checker.status == TriState.TRUE + checker.stop() + + async def test_unhealthy_after_failed_check(self): + async def check_fn() -> None: + raise RuntimeError("boom") + + checker, clock = _make_checker(check_fn=check_fn) + await checker.start() + await _settle(clock) + + assert checker.status == TriState.FALSE + checker.stop() + + async def test_stop_resets_to_unknown(self): + checker, clock = _make_checker() + await checker.start() + await _settle(clock) + assert checker.status == TriState.TRUE + + checker.stop() + assert checker.status == TriState.UNKNOWN + + async def test_pause_resets_to_unknown(self): + checker, clock = _make_checker() + await checker.start() + await _settle(clock) + assert checker.status == TriState.TRUE + + checker.pause() + assert checker.status == TriState.UNKNOWN + checker.stop() + + async def test_resume_resets_to_unknown(self): + checker, clock = _make_checker() + await checker.start() + await _settle(clock) + assert checker.status == TriState.TRUE + + checker.pause() + checker.resume() + assert checker.status == TriState.UNKNOWN + checker.stop() + + async def test_recovers_from_unhealthy_to_healthy(self): + call_count = 0 + + async def check_fn() -> None: + nonlocal call_count + call_count += 1 + if call_count == 1: + raise RuntimeError("transient") + + checker, clock = _make_checker(check_fn=check_fn, interval=5.0) + await checker.start() + + await _settle(clock) + assert checker.status == TriState.FALSE + + await clock.elapse(5.0) + assert checker.status == TriState.TRUE + + checker.stop() + + +class TestFailureThresholdDebounce: + """With failure_threshold > 1, transient failures must not flip the status to FALSE + until that many consecutive checks have failed; any success resets the counter.""" + + def _flaky_check_fn(self, outcomes: list[bool]): + idx = 0 + + async def check_fn() -> None: + nonlocal idx + ok = outcomes[idx] + idx += 1 + if not ok: + raise RuntimeError("boom") + + return check_fn + + async def test_below_threshold_keeps_previous_status(self): + # success, then 2 failures (< threshold 3): status stays TRUE. + check_fn = self._flaky_check_fn([True, False, False]) + checker, clock = _make_checker(check_fn=check_fn, interval=5.0, failure_threshold=3) + await checker.start() + await _settle(clock) + assert checker.status == TriState.TRUE + + await clock.elapse(5.0) + assert checker.status == TriState.TRUE + assert checker._consecutive_failures == 1 + + await clock.elapse(5.0) + assert checker.status == TriState.TRUE + assert checker._consecutive_failures == 2 + + checker.stop() + + async def test_status_flips_false_only_at_threshold(self): + check_fn = self._flaky_check_fn([False, False, False]) + checker, clock = _make_checker(check_fn=check_fn, interval=5.0, failure_threshold=3) + await checker.start() + + await _settle(clock) + assert checker.status == TriState.UNKNOWN # 1st failure: below threshold, keep initial UNKNOWN + + await clock.elapse(5.0) + assert checker.status == TriState.UNKNOWN # 2nd failure: still below threshold + + await clock.elapse(5.0) + assert checker.status == TriState.FALSE # 3rd consecutive failure: threshold reached + + checker.stop() + + async def test_success_resets_failure_counter(self): + # 2 failures, a success, then 2 more failures: never reaches 3 consecutive, stays TRUE. + check_fn = self._flaky_check_fn([False, False, True, False, False]) + checker, clock = _make_checker(check_fn=check_fn, interval=5.0, failure_threshold=3) + await checker.start() + + await _settle(clock) # fail 1 + await clock.elapse(5.0) # fail 2 + assert checker._consecutive_failures == 2 + + await clock.elapse(5.0) # success -> reset + assert checker.status == TriState.TRUE + assert checker._consecutive_failures == 0 + + await clock.elapse(5.0) # fail 1 + await clock.elapse(5.0) # fail 2 + assert checker.status == TriState.TRUE + assert checker._consecutive_failures == 2 + + checker.stop() + + async def test_on_result_reports_raw_per_check_not_debounced(self): + results: list[bool] = [] + check_fn = self._flaky_check_fn([True, False, False, False]) + checker, clock = _make_checker( + check_fn=check_fn, on_result=lambda s: results.append(s), interval=5.0, failure_threshold=3 + ) + await checker.start() + await _settle(clock) + for _ in range(3): + await clock.elapse(5.0) + checker.stop() + + assert results == [True, False, False, False] + + async def test_resume_resets_failure_counter(self): + check_fn = self._flaky_check_fn([False, False, False]) + checker, clock = _make_checker(check_fn=check_fn, interval=5.0, failure_threshold=3) + await checker.start() + await _settle(clock) + await clock.elapse(5.0) + assert checker._consecutive_failures == 2 + + checker.resume() + assert checker._consecutive_failures == 0 + checker.stop() + + +class TestRolloutHealthCheckerPreservesImmediateFailure: + def test_rollout_checker_forces_failure_threshold_one(self): + """Rollout health checker keeps pre-debounce semantics (single failure -> unhealthy) + even when handed a config with a larger threshold.""" + from miles.utils.health_checker import SimpleHealthCheckerConfig, create_rollout_cell_health_checker + + checker = create_rollout_cell_health_checker( + cell_id="c0", + get_engines=lambda: [object()], + config=SimpleHealthCheckerConfig(interval=10.0, timeout=10.0, first_wait=0.0, failure_threshold=5), + ) + + assert checker._config.failure_threshold == 1 + + +class TestNoopHealthChecker: + def test_noop_status_is_always_unknown(self): + checker = NoopHealthChecker() + assert checker.status == TriState.UNKNOWN diff --git a/tests/fast/utils/test_heartbeat_utils.py b/tests/fast/utils/test_heartbeat_utils.py new file mode 100644 index 00000000000..2fd567960b2 --- /dev/null +++ b/tests/fast/utils/test_heartbeat_utils.py @@ -0,0 +1,111 @@ +import miles.utils.heartbeat_utils as heartbeat_utils +from miles.utils.heartbeat_utils import HeartbeatStatus, SimpleHeartbeat + + +class TestSimpleHeartbeatBumpCount: + def test_initial_bump_count_is_zero(self): + """A fresh SimpleHeartbeat reports a bump_count of 0.""" + heartbeat = SimpleHeartbeat() + assert heartbeat.status().bump_count == 0 + + def test_single_bump_increments_count(self): + """One bump() raises bump_count from 0 to 1.""" + heartbeat = SimpleHeartbeat() + heartbeat.bump() + assert heartbeat.status().bump_count == 1 + + def test_bump_count_increments_monotonically(self): + """Repeated bump() calls increment bump_count by exactly one each time.""" + heartbeat = SimpleHeartbeat() + for expected in range(1, 6): + heartbeat.bump() + assert heartbeat.status().bump_count == expected + + def test_bump_count_never_resets(self): + """bump_count keeps growing and never drops back across many bumps.""" + heartbeat = SimpleHeartbeat() + previous = heartbeat.status().bump_count + for _ in range(20): + heartbeat.bump() + current = heartbeat.status().bump_count + assert current == previous + 1 + previous = current + assert heartbeat.status().bump_count == 20 + + +class TestSimpleHeartbeatStatusObject: + def test_status_returns_heartbeat_status_instance(self): + """status() yields a HeartbeatStatus dataclass instance.""" + heartbeat = SimpleHeartbeat() + assert isinstance(heartbeat.status(), HeartbeatStatus) + + def test_status_is_frozen(self): + """The returned HeartbeatStatus is frozen and rejects attribute mutation.""" + import dataclasses + + heartbeat = SimpleHeartbeat() + status = heartbeat.status() + try: + status.bump_count = 999 + except dataclasses.FrozenInstanceError: + pass + else: + raise AssertionError("expected FrozenInstanceError on frozen dataclass") + + def test_status_returns_new_object_each_bump(self): + """Each bump() replaces the status with a distinct object identity.""" + heartbeat = SimpleHeartbeat() + before = heartbeat.status() + heartbeat.bump() + after = heartbeat.status() + assert before is not after + + def test_status_stable_between_bumps(self): + """status() returns the same object identity when no bump occurs in between.""" + heartbeat = SimpleHeartbeat() + assert heartbeat.status() is heartbeat.status() + + def test_old_status_snapshot_unaffected_by_later_bump(self): + """A captured status snapshot keeps its bump_count after subsequent bumps.""" + heartbeat = SimpleHeartbeat() + snapshot = heartbeat.status() + heartbeat.bump() + heartbeat.bump() + assert snapshot.bump_count == 0 + assert heartbeat.status().bump_count == 2 + + +class TestSimpleHeartbeatTimestamp: + def test_initial_timestamp_from_time_time(self, monkeypatch): + """Construction stamps last_active_timestamp from time.time().""" + monkeypatch.setattr(heartbeat_utils.time, "time", lambda: 1000.0) + heartbeat = SimpleHeartbeat() + assert heartbeat.status().last_active_timestamp == 1000.0 + + def test_bump_updates_timestamp_from_time_time(self, monkeypatch): + """bump() refreshes last_active_timestamp to the current time.time() value.""" + clock = {"now": 1000.0} + monkeypatch.setattr(heartbeat_utils.time, "time", lambda: clock["now"]) + heartbeat = SimpleHeartbeat() + assert heartbeat.status().last_active_timestamp == 1000.0 + + clock["now"] = 2500.0 + heartbeat.bump() + status = heartbeat.status() + assert status.last_active_timestamp == 2500.0 + assert status.bump_count == 1 + + def test_consecutive_bumps_track_timestamp(self, monkeypatch): + """Each bump() captures the time.time() value observed at that call.""" + clock = {"now": 100.0} + monkeypatch.setattr(heartbeat_utils.time, "time", lambda: clock["now"]) + heartbeat = SimpleHeartbeat() + + clock["now"] = 200.0 + heartbeat.bump() + assert heartbeat.status().last_active_timestamp == 200.0 + + clock["now"] = 350.0 + heartbeat.bump() + assert heartbeat.status().last_active_timestamp == 350.0 + assert heartbeat.status().bump_count == 2 From 4a52fcf542b67f853b8b5e741bce201b7bf44ca6 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 11/14] Add the fault-tolerance dependency, CI label, and logger-config setup Add the nvidia-resiliency-ext dependency, the "ft" CI test label, the FT test fixtures in the rollout conftest, and route startup logging through configure_logger_raw. The fault-tolerance CLI arguments themselves now live with the features that consume them (distributed across the per-feature commits). --- miles/utils/arguments.py | 6 ++++-- tests/ci/labels.py | 2 ++ tests/fast/ray/rollout/conftest.py | 8 ++++++++ 3 files changed, 14 insertions(+), 2 deletions(-) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 7d7858b91fe..49d99898c76 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -667,7 +667,7 @@ def add_fault_tolerance_arguments(parser): "--use-fault-tolerance", action="store_true", default=False, - help="Whether to enable the fault tolerance function during rollout.", + help="Enable fault tolerance. Use --ft-components to select which components.", ) parser.add_argument( "--rollout-health-check-interval", @@ -1995,7 +1995,7 @@ def add_user_provided_function_arguments(parser): def parse_args(add_custom_arguments=None): # Users may call `parse_args` very early, thus we ensure logger is configured here - configure_logger() + configure_logger("main") add_miles_arguments = get_miles_extra_args_provider(add_custom_arguments) @@ -2016,6 +2016,7 @@ def parse_args(add_custom_arguments=None): args.indexer_rope_interleave = bool(getattr(hf_config, "indexer_rope_interleave", False)) logger.info(f"Setting indexer_rope_interleave: {args.indexer_rope_interleave} into args") + # TODO: unify this .rank and .world_size w/ indep_dp logics args.rank = 0 args.world_size = args.actor_num_nodes * args.actor_num_gpus_per_node args = set_default_megatron_args(args) @@ -2023,6 +2024,7 @@ def parse_args(add_custom_arguments=None): from miles.backends.experimental.fsdp_utils.arguments import load_fsdp_args args = load_fsdp_args(extra_args_provider=add_miles_arguments) + # TODO: unify this .rank and .world_size w/ indep_dp logics args.rank = 0 # Primary process rank for wandb initialization args.world_size = args.actor_num_nodes * args.actor_num_gpus_per_node diff --git a/tests/ci/labels.py b/tests/ci/labels.py index 22b5dfcbb9f..26646a1e114 100644 --- a/tests/ci/labels.py +++ b/tests/ci/labels.py @@ -26,6 +26,8 @@ "ckpt": "Checkpoint save / load tests", "lora": "LoRA training tests", "precision": "Numerical precision parity tests", + "ft-short": "Fault-tolerance trainer comparison tests (no_failure / deterministic / with_failure)", + "ft-long": "Fault-tolerance trainer soak tests (random-crash survival, realistic-gsm8k convergence)", "weight-update": "Weight update tests", "replay": "Routing / indexer replay tests", "qwen35": "Qwen3.5-35B-A3B MTP / spec-v2 e2e tests", diff --git a/tests/fast/ray/rollout/conftest.py b/tests/fast/ray/rollout/conftest.py index de99c167f3f..7e5e94fdbda 100644 --- a/tests/fast/ray/rollout/conftest.py +++ b/tests/fast/ray/rollout/conftest.py @@ -44,6 +44,7 @@ def make_args(**overrides: Any) -> Namespace: use_dynamic_global_batch_size=False, disable_rollout_trim_samples=False, balance_data=False, + delay_split_train_data_by_dp=False, # advantage / reward advantage_estimator="grpo", rewards_normalization=True, @@ -97,6 +98,13 @@ def make_args(**overrides: Any) -> Namespace: save_debug_rollout_data=None, load_debug_rollout_data=None, load_debug_rollout_data_subsample=None, + ci_inject_rollout_data_path=None, + ci_inject_rollout_data_start_rollout_id=None, + ci_inject_rollout_data_min_match_ratio=0.9, + # event checkpointing (event_logger.restore/snapshot in RolloutManager) + save_debug_event_data=None, + load=None, + save=None, # CI ci_test=False, # dumper (sglang debug dumper integration) From 2ed9fcf003971cd674ef81bc8805003edeea1379 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 12/14] Always reconnect rollout engines on weight-update setup Unconditionally disconnect-then-reconnect the model-update process group when (re)connecting rollout engines, guarding the destroy against a missing group, so a reconfigured/healed engine set can rebuild the NCCL group from scratch. - broadcast.py: drop the "only disconnect if group exists" short-circuit; guard `destroy_process_group` against None. --- .../update_weight_from_distributed/broadcast.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py index 679d31d4c92..f5ba0ae5ae1 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/broadcast.py @@ -73,8 +73,9 @@ def connect_rollout_engines( self._group_name = f"miles-pp_{pp_rank}" if self._is_source: - if (g := self._model_update_groups) is not None: - disconnect_rollout_engines_from_distributed(self.args, self._group_name, g, self.rollout_engines) + disconnect_rollout_engines_from_distributed( + self.args, self._group_name, self._model_update_groups, self.rollout_engines + ) self._model_update_groups = connect_rollout_engines_from_distributed( self.args, self._group_name, rollout_engines ) @@ -195,8 +196,11 @@ def disconnect_rollout_engines_from_distributed(args, group_name, model_update_g Destroy NCCL on training and engines. """ refs = [engine.destroy_weights_update_group.remote(group_name) for engine in rollout_engines] - dist.destroy_process_group(model_update_groups) - ray.get(refs) + try: + if model_update_groups is not None: + dist.destroy_process_group(model_update_groups) + finally: + ray.get(refs) def update_weights_from_distributed( From a838a0ba213b07b0294993ba81bf3185df3a3043 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Mon, 22 Jun 2026 18:15:16 +0800 Subject: [PATCH 13/14] Add a fault-injection RPC to train actors Expose an `inject_fault` Ray method on TrainRayActor (in its own concurrency group) that triggers a configured failure mode via the fault injector, so fault-tolerance tests can crash/hang specific actors on demand. - train_actor.py: `inject_fault` RPC. --- miles/ray/actor_group.py | 4 +++- miles/ray/train_actor.py | 6 ++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index f729d948aed..f1436cfafaa 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -89,7 +89,9 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): actor_impl = FSDPTrainRayActor - TrainRayActor = ray.remote(num_gpus=1, runtime_env={"env_vars": env_vars})(actor_impl) + TrainRayActor = ray.remote( + num_gpus=1, runtime_env={"env_vars": env_vars}, concurrency_groups={"fault_injector": 1} + )(actor_impl) # Create worker actors actor_handles = [] diff --git a/miles/ray/train_actor.py b/miles/ray/train_actor.py index 9e6659026fc..7ab2bfcd745 100644 --- a/miles/ray/train_actor.py +++ b/miles/ray/train_actor.py @@ -15,6 +15,7 @@ from miles.utils.env_report import collect_and_print_node_env_report from miles.utils.logging_utils import configure_logger from miles.utils.memory_utils import clear_memory, print_memory +from miles.utils.test_utils.fault_injector import inject_fault as _inject_fault if TYPE_CHECKING: from miles.ray.rollout.rollout_manager import EnginesAndLock @@ -53,6 +54,7 @@ def __init__(self, world_size, rank, master_addr, master_port): # os.environ["LOCAL_RANK"] = str(ray.get_gpu_ids()[0]) os.environ["LOCAL_RANK"] = str(get_local_gpu_id()) + # TODO mv the args into ctor def init(self, args, role, with_ref=False, with_opd_teacher=False): self.args = args self.role = role @@ -109,6 +111,10 @@ def init(self, args, role, with_ref=False, with_opd_teacher=False): except Exception as e: logger.info(f"Warning: Failed to set NUMA affinity: {e}") + @ray.method(concurrency_group="fault_injector") + def inject_fault(self, mode: str) -> None: + _inject_fault(mode=mode) + def clear_memory(self): print_memory("before TrainRayActor.clear_memory") clear_memory() From c7798b563c62bba207a61d340e87cd444b1e693b Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Wed, 8 Jul 2026 10:47:39 +0800 Subject: [PATCH 14/14] Delay splitting train data by DP until actor-side processing Extract the DP split into a witness-aware split_train_data_by_dp_raw helper (with unit tests; the key list also carries seq_witness_ids) and use it to split the training data on the actor side when delay_split_train_data_by_dp is set, deferring the DP split from the rollout side to actor-side processing. split_train_data_by_dp stays a thin wrapper that ray.puts each partition. - miles/ray/rollout/train_data_conversion.py (+ tests), miles/utils/data.py, actor_group.py, rollout_manager.py. --- miles/ray/actor_group.py | 4 +- miles/ray/rollout/rollout_manager.py | 8 ++- miles/ray/rollout/train_data_conversion.py | 13 +++- miles/utils/arguments.py | 5 ++ miles/utils/data.py | 18 ++++- .../rollout/real_ray/test_rollout_manager.py | 7 +- .../ray/rollout/test_train_data_conversion.py | 67 +++++++++++++++++++ 7 files changed, 111 insertions(+), 11 deletions(-) diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index f1436cfafaa..bfbfdf55e85 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -119,9 +119,9 @@ async def init(self): "init", self.args, self.role, with_ref=self.with_ref, with_opd_teacher=self.with_opd_teacher ) - async def train(self, rollout_id, rollout_data_ref): + async def train(self, rollout_id, rollout_data_pack): """Do one rollout training""" - await self._broadcast("train", rollout_id, rollout_data_ref) + await self._broadcast("train", rollout_id, rollout_data_pack["data_ref"]) async def save_model(self, rollout_id, force_sync=False): """Save actor model""" diff --git a/miles/ray/rollout/rollout_manager.py b/miles/ray/rollout/rollout_manager.py index c216b440b24..107d771233c 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -28,6 +28,7 @@ from miles.utils.logging_utils import configure_logger from miles.utils.metric_checker import MetricChecker from miles.utils.misc import load_function +from miles.utils.ray_utils import Box from miles.utils.tracking_utils import init_tracking logging.getLogger("httpx").setLevel(logging.WARNING) @@ -117,7 +118,12 @@ async def generate(self, rollout_id): custom_convert_samples_to_train_data_func=self.custom_convert_samples_to_train_data_func, custom_reward_post_process_func=self.custom_reward_post_process_func, ) - return split_train_data_by_dp(self.args, data, self.train_parallel_config["dp_size"]) + sample_indices = data.get("sample_indices") + if self.args.delay_split_train_data_by_dp: + data_ref = Box(ray.put(data)) + else: + data_ref = split_train_data_by_dp(self.args, data, self.train_parallel_config["dp_size"]) + return dict(sample_indices=sample_indices, data_ref=data_ref) async def eval(self, rollout_id): if self.args.debug_train_only: diff --git a/miles/ray/rollout/train_data_conversion.py b/miles/ray/rollout/train_data_conversion.py index 65bc8d4b6db..e67087fbb3d 100644 --- a/miles/ray/rollout/train_data_conversion.py +++ b/miles/ray/rollout/train_data_conversion.py @@ -122,6 +122,12 @@ def _post_process_rewards(args, samples: list[Sample] | list[list[Sample]], cust def split_train_data_by_dp(args, data, dp_size): + """Split the train data by data parallel size.""" + rollout_data_list = split_train_data_by_dp_raw(args, data, dp_size=dp_size) + return [Box(ray.put(rollout_data)) for rollout_data in rollout_data_list] + + +def split_train_data_by_dp_raw(args, data: dict[str, Any], *, dp_size: int) -> list[dict[str, Any]]: """Split the train data by data parallel size.""" rollout_data = {} @@ -136,7 +142,7 @@ def split_train_data_by_dp(args, data, dp_size): else: partitions = [range(i, len(total_lengths), dp_size) for i in range(dp_size)] - rollout_data_refs = [] + ans = [] for i in range(dp_size): rollout_data = {} @@ -157,6 +163,7 @@ def split_train_data_by_dp(args, data, dp_size): "prompt", "teacher_log_probs", "opd_reverse_kl", + "seq_witness_ids", "weight_versions", ]: if key not in data: @@ -172,5 +179,5 @@ def split_train_data_by_dp(args, data, dp_size): if key not in data: continue rollout_data[key] = data[key] - rollout_data_refs.append(Box(ray.put(rollout_data))) - return rollout_data_refs + ans.append(rollout_data) + return ans diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 49d99898c76..7ff03b77878 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -270,6 +270,11 @@ def add_train_arguments(parser): parser.add_argument( "--log-probs-chunk-size", type=int, default=-1, help="Chunk size to compute log probs to save memory" ) + parser.add_argument( + "--delay-split-train-data-by-dp", + action="store_true", + default=False, + ) parser.add_argument( "--allgather-cp", action="store_true", diff --git a/miles/utils/data.py b/miles/utils/data.py index 0fe7ad483fd..21d01db972b 100644 --- a/miles/utils/data.py +++ b/miles/utils/data.py @@ -8,6 +8,8 @@ import numpy as np import ray +from miles.ray.rollout.train_data_conversion import split_train_data_by_dp_raw + try: import pyarrow.parquet as pq except ImportError: @@ -272,9 +274,19 @@ def get_minimum_num_micro_batch_size(total_lengths, max_tokens_per_gpu): return len(batches) -def process_rollout_data(args, rollout_data_ref, dp_rank, dp_size): - assert len(rollout_data_ref) == dp_size - rollout_data = ray.get(rollout_data_ref[dp_rank].inner) +def process_rollout_data( + args, + rollout_data_ref, + dp_rank, + dp_size, +): + if args.delay_split_train_data_by_dp: + raw = ray.get(rollout_data_ref.inner) + raw = split_train_data_by_dp_raw(args, raw, dp_size=dp_size) + rollout_data = raw[dp_rank] + else: + assert len(rollout_data_ref) == dp_size + rollout_data = ray.get(rollout_data_ref[dp_rank].inner) partition = rollout_data.pop("partition") total_lengths = rollout_data["total_lengths"] diff --git a/tests/fast/ray/rollout/real_ray/test_rollout_manager.py b/tests/fast/ray/rollout/real_ray/test_rollout_manager.py index ab7d79b75df..962a53099dc 100644 --- a/tests/fast/ray/rollout/real_ray/test_rollout_manager.py +++ b/tests/fast/ray/rollout/real_ray/test_rollout_manager.py @@ -516,9 +516,12 @@ def fake_rollout_fn(input): assert len(captured) == 1 assert isinstance(captured[0], RolloutFnTrainInput) assert captured[0].rollout_id == 42 + # generate returns {"sample_indices": ..., "data_ref": ...}; # split_train_data_by_dp returns Box(ObjectRef) per dp rank - assert len(result) == 2 - partitions = ray.get([box.inner for box in result]) + assert set(result) == {"sample_indices", "data_ref"} + data_refs = result["data_ref"] + assert len(data_refs) == 2 + partitions = ray.get([box.inner for box in data_refs]) for partition in partitions: assert "tokens" in partition assert "rewards" in partition diff --git a/tests/fast/ray/rollout/test_train_data_conversion.py b/tests/fast/ray/rollout/test_train_data_conversion.py index 4a3da1a2883..31f3b1975db 100644 --- a/tests/fast/ray/rollout/test_train_data_conversion.py +++ b/tests/fast/ray/rollout/test_train_data_conversion.py @@ -1,7 +1,10 @@ from __future__ import annotations +from unittest.mock import MagicMock + import pytest import ray +import torch from hypothesis import HealthCheck, given, settings from hypothesis import strategies as st from tests.fast.ray.rollout.conftest import make_args, make_sample, make_samples_grouped @@ -10,6 +13,7 @@ _post_process_rewards, convert_samples_to_train_data, split_train_data_by_dp, + split_train_data_by_dp_raw, ) from miles.utils.types import Sample @@ -472,3 +476,66 @@ def test_partition_indices_form_a_partition(self): parts = [ray.get(r.inner) for r in refs] all_indices = sorted(i for p in parts for i in p["partition"]) assert all_indices == list(range(n)) + + +class TestSplitTrainDataRaw: + def test_witness_ids_split_across_dp(self) -> None: + tokens = [[1, 2, 3], [4, 5], [6, 7, 8, 9], [10, 11]] + witness_ids = [ + torch.tensor([0, 0, 0]), + torch.tensor([1, 1]), + torch.tensor([2, 2, 2, 2]), + torch.tensor([3, 3]), + ] + + data = { + "tokens": tokens, + "seq_witness_ids": witness_ids, + "response_lengths": [1, 1, 1, 1], + "loss_masks": [[0, 0, 1], [0, 1], [0, 0, 0, 1], [0, 1]], + } + + args = MagicMock() + args.balance_data = False + + result = split_train_data_by_dp_raw(args, data, dp_size=2) + + assert len(result) == 2 + assert "seq_witness_ids" in result[0] + assert "seq_witness_ids" in result[1] + assert len(result[0]["seq_witness_ids"]) == 2 + assert len(result[1]["seq_witness_ids"]) == 2 + + def test_indexer_topk_and_opd_reverse_kl_split_across_dp(self) -> None: + """Keys from the rollout-side split (rollout_indexer_topk, opd_reverse_kl) partition per sample.""" + data = { + "tokens": [[1, 2], [3, 4], [5, 6], [7, 8]], + "response_lengths": [1, 1, 1, 1], + "loss_masks": [[0, 1], [0, 1], [0, 1], [0, 1]], + "rollout_indexer_topk": [torch.tensor([i]) for i in range(4)], + "opd_reverse_kl": [[float(i)] for i in range(4)], + } + + args = MagicMock() + args.balance_data = False + + result = split_train_data_by_dp_raw(args, data, dp_size=2) + + assert len(result) == 2 + for part in result: + assert len(part["rollout_indexer_topk"]) == 2 + assert len(part["opd_reverse_kl"]) == 2 + + def test_no_witness_ids_when_absent(self) -> None: + tokens = [[1, 2], [3, 4]] + data = { + "tokens": tokens, + "response_lengths": [1, 1], + "loss_masks": [[0, 1], [0, 1]], + } + + args = MagicMock() + args.balance_data = False + + result = split_train_data_by_dp_raw(args, data, dp_size=1) + assert "seq_witness_ids" not in result[0]