From 260ccd1e8d4e5b53d361d2ca2c016105820f5804 Mon Sep 17 00:00:00 2001 From: Zane Li Date: Mon, 24 Aug 2026 16:29:55 +0800 Subject: [PATCH 1/2] =?UTF-8?q?fix:=20=F0=9F=90=9B=20fix=20Ray=20topology?= =?UTF-8?q?=20before=20actor=20creation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/running-modes.md | 2 +- tests/tine_engine/test_cpu_accelerator.py | 10 +- tests/tine_engine/test_ddp_accelerator.py | 5 +- tests/utils/test_ray_utils_unit.py | 177 ++++++++++++++++-- tinyexp/exp_mixins/basic_mixins.py | 29 +-- .../accelerator/base_accelerator.py | 1 - tinyexp/utils/ray_utils.py | 118 +++++++----- 7 files changed, 259 insertions(+), 83 deletions(-) diff --git a/docs/running-modes.md b/docs/running-modes.md index f4cfc9f..120766f 100644 --- a/docs/running-modes.md +++ b/docs/running-modes.md @@ -99,7 +99,7 @@ Requirements: - `ray_cfg.ray_placement_timeout_s` controls how long TinyExp waits for the placement group; the default is 120 seconds. - Requests that exceed the cluster's total CPU or GPU capacity fail before placement starts. If the total capacity is sufficient but currently busy, placement waits up to the configured timeout. - GPU workers require a CUDA-enabled PyTorch installation and visible GPUs. -- TinyExp derives `RANK`, `LOCAL_RANK`, and `LOCAL_WORLD_SIZE` from the nodes where Ray actually places the actors. When actors are interleaved in Ray's returned actor list, global ranks are reassigned so ranks remain contiguous within each node (for example, node-local workers receive `0..N-1`, then the next node receives the following range). +- TinyExp reads the placement-group bundle-to-node topology before creating Ray worker actors, then derives `RANK` and `LOCAL_RANK` from that topology. Ray workers must be homogeneous: every participating node must host the same number of workers, so every node has the same local-rank range. When bundles are interleaved across nodes, global ranks are reassigned so ranks remain contiguous within each node (for example, node-local workers receive `0..N-1`, then the next node receives the following range). If `RAY_ADDRESS` already points to a reachable Ray cluster, `ray.init()` can attach to that cluster. Otherwise, Ray starts a local runtime. diff --git a/tests/tine_engine/test_cpu_accelerator.py b/tests/tine_engine/test_cpu_accelerator.py index 12d4388..c09c96d 100644 --- a/tests/tine_engine/test_cpu_accelerator.py +++ b/tests/tine_engine/test_cpu_accelerator.py @@ -3,7 +3,7 @@ from torch import nn from tinyexp.tiny_engine.accelerator import CPUAccelerator -from tinyexp.utils.ray_utils import get_num_worker_options, get_placement_group +from tinyexp.utils.ray_utils import get_num_worker_options, get_placement_group, get_placement_group_node_ids @ray.remote @@ -40,7 +40,13 @@ def test_ddp_accelerator(self, ray_session): num_gpus_per_worker=0.0, # CPU workers, so no GPUs num_cpus_per_worker=2, # Each worker gets 2 CPUs ) - options_list = get_num_worker_options(pg, num_worker=num_worker, gpu_ratio=0.0) + node_ids = get_placement_group_node_ids(pg, num_worker) + options_list = get_num_worker_options( + pg, + num_worker=num_worker, + gpu_ratio=0.0, + node_ids=node_ids, + ) # Create the remote actors. worker_group = [CPUAcceleratorProxy.options(**options).remote() for options in options_list] diff --git a/tests/tine_engine/test_ddp_accelerator.py b/tests/tine_engine/test_ddp_accelerator.py index 03fa676..064ec5d 100644 --- a/tests/tine_engine/test_ddp_accelerator.py +++ b/tests/tine_engine/test_ddp_accelerator.py @@ -5,7 +5,7 @@ import torch from tinyexp.tiny_engine.accelerator import DDPAccelerator -from tinyexp.utils.ray_utils import get_num_worker_options, get_placement_group +from tinyexp.utils.ray_utils import get_num_worker_options, get_placement_group, get_placement_group_node_ids @ray.remote @@ -44,6 +44,7 @@ def test_ddp_accelerator(self, ray_session): num_gpus_per_worker=1.0, # Each worker gets 1 GPU num_cpus_per_worker=4, ) + node_ids = get_placement_group_node_ids(pg, num_workers) gpu_per_actor = 0.2 cpus_per_actor = 1 @@ -54,6 +55,7 @@ def test_ddp_accelerator(self, ray_session): num_workers, gpu_ratio=gpu_per_actor, num_cpus_per_worker=cpus_per_actor, + node_ids=node_ids, ) worker_group1 = [DDPAcceleratorProxy.options(**options).remote() for options in options_list1] @@ -62,6 +64,7 @@ def test_ddp_accelerator(self, ray_session): num_workers, gpu_ratio=gpu_per_actor, num_cpus_per_worker=cpus_per_actor, + node_ids=node_ids, ) worker_group2 = [DDPAcceleratorProxy.options(**options).remote() for options in options_list2] diff --git a/tests/utils/test_ray_utils_unit.py b/tests/utils/test_ray_utils_unit.py index 24521ad..30d750e 100644 --- a/tests/utils/test_ray_utils_unit.py +++ b/tests/utils/test_ray_utils_unit.py @@ -1,5 +1,7 @@ from __future__ import annotations +from types import SimpleNamespace + import pytest import ray from omegaconf import OmegaConf @@ -11,7 +13,9 @@ _maybe_start_ray_redis_cache, build_ray_worker_env_vars, get_network_config, + get_num_worker_options, get_placement_group, + get_placement_group_node_ids, ) @@ -321,14 +325,13 @@ def test_build_ray_worker_env_vars_uses_actual_node_local_ranks( "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK", - "LOCAL_WORLD_SIZE", } ] * 4 - assert [(env["RANK"], env["LOCAL_RANK"], env["LOCAL_WORLD_SIZE"]) for env in env_vars] == [ - ("0", "0", "2"), - ("1", "1", "2"), - ("2", "0", "2"), - ("3", "1", "2"), + assert [(env["RANK"], env["LOCAL_RANK"]) for env in env_vars] == [ + ("0", "0"), + ("1", "1"), + ("2", "0"), + ("3", "1"), ] @@ -344,14 +347,24 @@ def test_build_ray_worker_env_vars_groups_global_ranks_by_node( master_port=12345, ) - assert [(env["RANK"], env["LOCAL_RANK"], env["LOCAL_WORLD_SIZE"]) for env in env_vars] == [ - ("0", "0", "2"), - ("2", "0", "2"), - ("1", "1", "2"), - ("3", "1", "2"), + assert [(env["RANK"], env["LOCAL_RANK"]) for env in env_vars] == [ + ("0", "0"), + ("2", "0"), + ("1", "1"), + ("3", "1"), ] +def test_build_ray_worker_env_vars_rejects_heterogeneous_worker_counts() -> None: + with pytest.raises(ValueError, match="homogeneous.*node-a=2.*node-b=1"): + build_ray_worker_env_vars( + num_worker=3, + node_ids=["node-a", "node-b", "node-a"], + master_addr="10.0.0.1", + master_port=12345, + ) + + def test_build_ray_worker_env_vars_rejects_mismatched_worker_count() -> None: with pytest.raises(ValueError, match="one node id for every Ray worker"): build_ray_worker_env_vars( @@ -362,6 +375,148 @@ def test_build_ray_worker_env_vars_rejects_mismatched_worker_count() -> None: ) +def test_get_placement_group_node_ids_reads_bundle_assignments(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "tinyexp.utils.ray_utils.ray.util.placement_group_table", + lambda pg: {"bundles_to_node_id": {0: "node-a", 1: "node-b"}}, + ) + + assert get_placement_group_node_ids(object(), 2) == ["node-a", "node-b"] + + +def test_get_placement_group_node_ids_accepts_string_bundle_keys(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "tinyexp.utils.ray_utils.ray.util.placement_group_table", + lambda pg: {"bundles_to_node_id": {"0": "node-a", "1": "node-b"}}, + ) + + assert get_placement_group_node_ids(object(), 2) == ["node-a", "node-b"] + + +def test_get_placement_group_node_ids_rejects_missing_bundle_assignment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + "tinyexp.utils.ray_utils.ray.util.placement_group_table", + lambda pg: {"bundles_to_node_id": {0: "node-a"}}, + ) + + with pytest.raises(RuntimeError, match="bundle 1"): + get_placement_group_node_ids(object(), 2) + + +def test_get_num_worker_options_uses_final_topology_env_for_each_bundle() -> None: + placement_group = SimpleNamespace(bundle_specs=[{"CPU": 2, "GPU": 0}]) + + options_list = get_num_worker_options( + placement_group, + num_worker=4, + gpu_ratio=0.0, + num_cpus_per_worker=2, + master_addr="10.0.0.1", + master_port=12345, + node_ids=["node-a", "node-b", "node-a", "node-b"], + ) + + assert [option["scheduling_strategy"].placement_group_bundle_index for option in options_list] == [ + 0, + 1, + 2, + 3, + ] + assert [ + ( + option["runtime_env"]["env_vars"]["RANK"], + option["runtime_env"]["env_vars"]["LOCAL_RANK"], + ) + for option in options_list + ] == [("0", "0"), ("2", "0"), ("1", "1"), ("3", "1")] + assert all( + set(option["runtime_env"]["env_vars"]) == {"WORLD_SIZE", "RANK", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"} + for option in options_list + ) + + +def test_ray_cfg_run_resolves_topology_before_constructing_worker_actors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cfg = OmegaConf.create( + { + "ray_cfg": { + "ray_num_worker": 4, + "ray_num_cpus_per_worker": 1, + "ray_num_gpus_per_worker": 0, + "ray_placement_strategy": "PACK", + } + } + ) + events: list[object] = [] + + class FakeWorker: + def __init__(self, env_vars: dict[str, str]) -> None: + self.set_cfg = SimpleNamespace( + remote=lambda cfg: events.append(("set_cfg", env_vars["RANK"], cfg)), + ) + self.run = SimpleNamespace( + remote=lambda: events.append(("run", env_vars["RANK"])), + ) + + class FakeConfiguredActor: + def __init__(self, options: dict[str, object]) -> None: + self.options = options + + def remote(self) -> FakeWorker: + env_vars = self.options["runtime_env"]["env_vars"] + assert isinstance(env_vars, dict) + events.append(("actor", dict(env_vars))) + return FakeWorker(env_vars) + + class FakeRemoteClass: + def options(self, **options: object) -> FakeConfiguredActor: + return FakeConfiguredActor(options) + + monkeypatch.setattr("tinyexp.exp_mixins.basic_mixins.ray.init", lambda: None) + monkeypatch.setattr("tinyexp.exp_mixins.basic_mixins.ray.remote", lambda exp_class: FakeRemoteClass()) + monkeypatch.setattr( + "tinyexp.exp_mixins.basic_mixins.ray.cluster_resources", + lambda: {"CPU": 4.0, "GPU": 0.0}, + ) + monkeypatch.setattr("tinyexp.exp_mixins.basic_mixins.ray.is_initialized", lambda: True) + monkeypatch.setattr("tinyexp.exp_mixins.basic_mixins.ray.shutdown", lambda **kwargs: None) + monkeypatch.setattr("tinyexp.exp_mixins.basic_mixins.ray.get", lambda refs, **kwargs: refs) + monkeypatch.setattr( + "tinyexp.exp_mixins.basic_mixins.get_placement_group", + lambda **kwargs: SimpleNamespace(bundle_specs=[{"CPU": 1, "GPU": 0}]), + ) + monkeypatch.setattr( + "tinyexp.exp_mixins.basic_mixins.get_placement_group_node_ids", + lambda pg, num_worker: ( + events.append(("topology",)), + ["node-a", "node-b", "node-a", "node-b"], + )[1], + ) + monkeypatch.setattr( + "tinyexp.exp_mixins.basic_mixins.get_network_config", + lambda: ("10.0.0.1", 12345), + ) + monkeypatch.setattr( + "tinyexp.exp_mixins.basic_mixins.ray.util.remove_placement_group", + lambda pg: None, + ) + + RayCfgMixin.RayCfg.run(object, cfg) + + topology_index = events.index(("topology",)) + actor_events = [event for event in events if event[0] == "actor"] + assert topology_index < events.index(actor_events[0]) + assert [(event[1]["RANK"], event[1]["LOCAL_RANK"]) for event in actor_events] == [ + ("0", "0"), + ("2", "0"), + ("1", "1"), + ("3", "1"), + ] + + def test_get_network_config_uses_public_ray_ip_api( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tinyexp/exp_mixins/basic_mixins.py b/tinyexp/exp_mixins/basic_mixins.py index 43510af..aa9dea2 100644 --- a/tinyexp/exp_mixins/basic_mixins.py +++ b/tinyexp/exp_mixins/basic_mixins.py @@ -22,10 +22,10 @@ from ..utils.log_utils import tiny_logger_setup from ..utils.ray_utils import ( _maybe_start_ray_redis_cache, - build_ray_worker_env_vars, get_network_config, get_num_worker_options, get_placement_group, + get_placement_group_node_ids, ) @@ -77,14 +77,6 @@ def _resolve_ray_num_worker( @dataclass class RayCfgMixin: - def _get_ray_node_id(self) -> str: - """Return the Ray node hosting this worker actor.""" - return str(ray.get_runtime_context().get_node_id()) - - def _set_ray_runtime_env(self, env_vars: dict[str, str]) -> None: - """Apply launcher-provided environment variables before the experiment runs.""" - os.environ.update({key: str(value) for key, value in env_vars.items()}) - @dataclass class RayCfg: # ---------------- luancher configuration ---------------- # @@ -161,6 +153,7 @@ def run(cls, exp_class: type[Any], experiment_cfg: DictConfig) -> None: # noqa: timeout_s=placement_timeout_s, ) master_addr, master_port = get_network_config() + worker_node_ids = get_placement_group_node_ids(pg, ray_cfg.ray_num_worker) options_list = get_num_worker_options( pg, ray_cfg.ray_num_worker, @@ -168,16 +161,9 @@ def run(cls, exp_class: type[Any], experiment_cfg: DictConfig) -> None: # noqa: num_cpus_per_worker=needed_num_cpus_per_worker, master_addr=master_addr, master_port=master_port, - ) - worker_group = [remote_exp.options(**options).remote() for options in options_list] - - worker_node_ids = ray.get([worker._get_ray_node_id.remote() for worker in worker_group]) - runtime_envs = build_ray_worker_env_vars( - num_worker=ray_cfg.ray_num_worker, node_ids=worker_node_ids, - master_addr=master_addr, - master_port=master_port, ) + runtime_envs = [options["runtime_env"]["env_vars"] for options in options_list] print("==> Ray worker topology:", flush=True) node_ranks = {node_id: node_rank for node_rank, node_id in enumerate(dict.fromkeys(worker_node_ids))} topology = sorted( @@ -188,15 +174,10 @@ def run(cls, exp_class: type[Any], experiment_cfg: DictConfig) -> None: # noqa: print( f" rank={env_vars['RANK']}/{ray_cfg.ray_num_worker} node={node_id} " f"node_rank={node_ranks[node_id]} " - f"local_rank={env_vars['LOCAL_RANK']}/{env_vars['LOCAL_WORLD_SIZE']}", + f"local_rank={env_vars['LOCAL_RANK']}", flush=True, ) - ray.get( - [ - worker._set_ray_runtime_env.remote(env_vars) - for worker, env_vars in zip(worker_group, runtime_envs) - ] - ) + worker_group = [remote_exp.options(**options).remote() for options in options_list] ray.get([worker.set_cfg.remote(experiment_cfg) for worker in worker_group]) ray.get([worker.run.remote() for worker in worker_group]) finally: diff --git a/tinyexp/tiny_engine/accelerator/base_accelerator.py b/tinyexp/tiny_engine/accelerator/base_accelerator.py index daa783e..4268e4b 100644 --- a/tinyexp/tiny_engine/accelerator/base_accelerator.py +++ b/tinyexp/tiny_engine/accelerator/base_accelerator.py @@ -55,7 +55,6 @@ def __init__(self) -> None: self.rank = int(os.getenv("RANK", 0)) self.world_size = int(os.getenv("WORLD_SIZE", 1)) self.local_rank = int(os.getenv("LOCAL_RANK", 0)) - self.local_world_size = int(os.getenv("LOCAL_WORLD_SIZE", 1)) self.sync_gradients = True self._destroyed = False self._process_group_initialized = False diff --git a/tinyexp/utils/ray_utils.py b/tinyexp/utils/ray_utils.py index bf1fb3d..a80447b 100644 --- a/tinyexp/utils/ray_utils.py +++ b/tinyexp/utils/ray_utils.py @@ -5,7 +5,7 @@ from collections import Counter from collections.abc import Sequence from contextlib import suppress -from typing import Optional +from typing import Any, Optional import ray from omegaconf import DictConfig @@ -77,44 +77,65 @@ def get_placement_group( return pg -def get_worker_options(gpu_ratio, num_cpus, pg, rank, local_rank, num_worker, master_addr, master_port): - """Create options for Ray workers.""" - env_vars = _build_worker_env_vars( - num_worker=num_worker, - rank=rank, - local_rank=local_rank, - master_addr=master_addr, - master_port=master_port, - ) +def get_worker_options( + gpu_ratio: float, + num_cpus: int, + pg: Any, + bundle_index: int, + env_vars: dict[str, str], +) -> dict[str, Any]: + """Create Ray actor options for one worker and one placement-group bundle.""" return { "runtime_env": {"env_vars": env_vars}, - "scheduling_strategy": PlacementGroupSchedulingStrategy(placement_group=pg, placement_group_bundle_index=rank), + "scheduling_strategy": PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_bundle_index=bundle_index, + ), "num_cpus": num_cpus, "num_gpus": gpu_ratio, } def _build_worker_env_vars( - num_worker, - rank, - local_rank, - master_addr, - master_port, - local_world_size=1, -): + *, + num_worker: int, + rank: int, + local_rank: int, + master_addr: str, + master_port: int, +) -> dict[str, str]: env_vars = { "WORLD_SIZE": str(num_worker), "RANK": str(rank), "MASTER_ADDR": master_addr, "MASTER_PORT": str(master_port), "LOCAL_RANK": str(local_rank), - "LOCAL_WORLD_SIZE": str(local_world_size), } if os.getenv("GLOO_SOCKET_IFNAME"): env_vars["GLOO_SOCKET_IFNAME"] = os.environ["GLOO_SOCKET_IFNAME"] return env_vars +def get_placement_group_node_ids(pg: Any, num_worker: int) -> list[str]: + """Return the node hosting each placement-group bundle in bundle-index order.""" + placement_group_state = ray.util.placement_group_table(pg) + bundles_to_node_id = placement_group_state.get("bundles_to_node_id") + if not isinstance(bundles_to_node_id, dict): + raise RuntimeError("Ray did not return placement-group bundle-to-node assignments") # noqa: TRY003, TRY004 + + node_ids = [] + for bundle_index in range(num_worker): + node_id = bundles_to_node_id.get(bundle_index) + if node_id is None: + node_id = bundles_to_node_id.get(str(bundle_index)) + if node_id is None: + raise RuntimeError( # noqa: TRY003 + f"Ray did not return a node assignment for placement-group bundle {bundle_index}" + ) + node_ids.append(str(node_id)) + return node_ids + + def build_ray_worker_env_vars( num_worker: int, node_ids: Sequence[str], @@ -127,13 +148,21 @@ def build_ray_worker_env_vars( if not node_ids: return [] - local_world_sizes = Counter(node_ids) + worker_counts = Counter(node_ids) + if len(set(worker_counts.values())) != 1: + counts = ", ".join(f"{node_id}={count}" for node_id, count in worker_counts.items()) + raise ValueError( # noqa: TRY003 + "Ray distributed workers must be homogeneous: every worker node must host the same " + f"number of workers; counts: {counts}" + ) + + workers_per_node = next(iter(worker_counts.values())) node_ranks = {node_id: node_rank for node_rank, node_id in enumerate(dict.fromkeys(node_ids))} node_rank_offsets: dict[str, int] = {} rank_offset = 0 for node_id in node_ranks: node_rank_offsets[node_id] = rank_offset - rank_offset += local_world_sizes[node_id] + rank_offset += workers_per_node local_ranks: Counter[str] = Counter() env_vars = [] @@ -147,7 +176,6 @@ def build_ray_worker_env_vars( local_rank=local_rank, master_addr=master_addr, master_port=master_port, - local_world_size=local_world_sizes[node_id], ) ) return env_vars @@ -163,31 +191,35 @@ def get_network_config(): def get_num_worker_options( - pg, - num_worker, - gpu_ratio=1.0, - num_cpus_per_worker=None, - master_addr=None, - master_port=None, -): - """Create options for multiple Ray workers with GPU allocation.""" + pg: Any, + num_worker: int, + gpu_ratio: float = 1.0, + num_cpus_per_worker: int | None = None, + master_addr: str | None = None, + master_port: int | None = None, + *, + node_ids: Sequence[str], +) -> list[dict[str, Any]]: + """Create actor options after the placement-group topology has been resolved.""" if num_cpus_per_worker is None: num_cpus_per_worker = pg.bundle_specs[0].get("CPU", 0) if master_addr is None or master_port is None: master_addr, master_port = get_network_config() - options_list = [] - for i in range(num_worker): - options = get_worker_options( - gpu_ratio, - num_cpus_per_worker, - pg, - i, - i, - num_worker, - master_addr, - master_port, + runtime_envs = build_ray_worker_env_vars( + num_worker=num_worker, + node_ids=node_ids, + master_addr=master_addr, + master_port=master_port, + ) + return [ + get_worker_options( + gpu_ratio=gpu_ratio, + num_cpus=num_cpus_per_worker, + pg=pg, + bundle_index=bundle_index, + env_vars=runtime_envs[bundle_index], ) - options_list.append(options) - return options_list + for bundle_index in range(num_worker) + ] From 4748519b324e039c5540af3e877797f250f9ec1b Mon Sep 17 00:00:00 2001 From: Zane Li Date: Tue, 25 Aug 2026 08:48:48 +0800 Subject: [PATCH 2/2] =?UTF-8?q?fix:=20=F0=9F=90=9B=20fix=20preserve=20Ray?= =?UTF-8?q?=20worker=20logs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/examples/test_pi_exp_run.py | 42 ++++++++++++++++++++++++++++++ tests/utils/test_ray_utils_unit.py | 7 ++++- tinyexp/examples/pi_exp.py | 5 +++- tinyexp/exp_mixins/basic_mixins.py | 3 ++- 4 files changed, 54 insertions(+), 3 deletions(-) create mode 100644 tests/examples/test_pi_exp_run.py diff --git a/tests/examples/test_pi_exp_run.py b/tests/examples/test_pi_exp_run.py new file mode 100644 index 0000000..51b4f4a --- /dev/null +++ b/tests/examples/test_pi_exp_run.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from tinyexp.examples import pi_exp + + +def test_pi_run_prints_result_on_stdout( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + events: list[str] = [] + + class DummyAccelerator: + rank = 0 + world_size = 1 + is_main_process = True + + def reduce_sum(self, tensor): + return tensor + + def destroy(self) -> None: + events.append("destroy") + + accelerator = DummyAccelerator() + logger = SimpleNamespace() + exp = pi_exp.Exp(output_root=str(tmp_path), exp_name="pi_test") + + monkeypatch.setattr(pi_exp, "torch", SimpleNamespace(pi=3.14)) + monkeypatch.setattr("tinyexp.tiny_engine.accelerator.CPUAccelerator", lambda: accelerator) + monkeypatch.setattr(exp.logger_cfg, "build_logger", lambda **kwargs: logger) + monkeypatch.setattr(exp, "print_cfg", lambda logger: {}) + monkeypatch.setattr(exp, "_estimate_pi", lambda accelerator: 3.14) + + exp.run() + + assert events == ["destroy"] + assert "pi ~= 3.140000 (error=0.000000, samples=10000000)" in capsys.readouterr().out diff --git a/tests/utils/test_ray_utils_unit.py b/tests/utils/test_ray_utils_unit.py index 30d750e..ff81ebb 100644 --- a/tests/utils/test_ray_utils_unit.py +++ b/tests/utils/test_ray_utils_unit.py @@ -34,6 +34,7 @@ def test_ray_cfg_run_uses_explicit_resources_without_dataloader_cfg( } ) captured: dict[str, object] = {} + shutdown_kwargs: list[dict[str, object]] = [] monkeypatch.setattr("tinyexp.exp_mixins.basic_mixins.ray.init", lambda: None) monkeypatch.setattr("tinyexp.exp_mixins.basic_mixins.ray.remote", lambda exp_class: object()) @@ -42,7 +43,10 @@ def test_ray_cfg_run_uses_explicit_resources_without_dataloader_cfg( lambda: {"CPU": 3.0, "GPU": 0.0}, ) monkeypatch.setattr("tinyexp.exp_mixins.basic_mixins.ray.is_initialized", lambda: True) - monkeypatch.setattr("tinyexp.exp_mixins.basic_mixins.ray.shutdown", lambda **kwargs: None) + monkeypatch.setattr( + "tinyexp.exp_mixins.basic_mixins.ray.shutdown", + lambda **kwargs: shutdown_kwargs.append(kwargs), + ) def stop_after_resource_resolution(**kwargs): # type: ignore[no-untyped-def] captured.update(kwargs) @@ -60,6 +64,7 @@ def stop_after_resource_resolution(**kwargs): # type: ignore[no-untyped-def] assert captured["num_cpus_per_worker"] == 3 assert captured["num_gpus_per_worker"] == 0 assert captured["timeout_s"] == 7.0 + assert shutdown_kwargs == [{"_exiting_interpreter": True}] def test_ray_cfg_run_rejects_invalid_ray_worker_count_without_starting_ray( diff --git a/tinyexp/examples/pi_exp.py b/tinyexp/examples/pi_exp.py index 0bc0121..74a9085 100644 --- a/tinyexp/examples/pi_exp.py +++ b/tinyexp/examples/pi_exp.py @@ -56,7 +56,10 @@ def run(self) -> None: pi = self._estimate_pi(accelerator) if accelerator.is_main_process: - logger.info(f"pi ~= {pi:.6f} (error={abs(pi - torch.pi):.6f}, samples={self.pi_cfg.total_samples})") + print( + f"pi ~= {pi:.6f} (error={abs(pi - torch.pi):.6f}, samples={self.pi_cfg.total_samples})", + flush=True, + ) accelerator.destroy() diff --git a/tinyexp/exp_mixins/basic_mixins.py b/tinyexp/exp_mixins/basic_mixins.py index aa9dea2..9430615 100644 --- a/tinyexp/exp_mixins/basic_mixins.py +++ b/tinyexp/exp_mixins/basic_mixins.py @@ -191,7 +191,8 @@ def run(cls, exp_class: type[Any], experiment_cfg: DictConfig) -> None: # noqa: if ray.is_initialized(): with suppress(Exception): - ray.shutdown() + # Drain Ray's asynchronous worker logs before tearing down the runtime. + ray.shutdown(_exiting_interpreter=True) ray_cfg: RayCfg = field(default_factory=RayCfg)