From 80527a024cebae61ca7ab1582d1ee9d98a458a78 Mon Sep 17 00:00:00 2001 From: fzyzcjy Date: Thu, 30 Jul 2026 08:48:09 +0800 Subject: [PATCH] Forbid positional arguments when constructing ray actors Every actor constructor call site now passes keyword arguments only, and the constructors enforce it with keyword-only signatures. A launch layer that builds workers from a spec can then describe any constructor as a plain kwargs dict, with no positional-argument channel to carry. --- miles/dashboard/backend.py | 4 +- miles/dashboard/collector.py | 4 +- miles/dashboard/gpu_sampler.py | 2 +- miles/ray/multi_lora/controller.py | 4 +- miles/ray/placement_group.py | 2 +- miles/ray/rollout/rollout_executor.py | 2 +- miles/ray/train/actor_factory.py | 10 +- miles/ray/train_actor.py | 1 + miles/utils/http_utils.py | 4 +- .../utils/tracking_utils/prometheus_utils.py | 4 +- tests/fast-gpu/test_gpu_sampler_hw.py | 4 +- tests/fast/dashboard/test_collector.py | 76 +++- tests/fast/dashboard/test_gpu_sampler.py | 55 ++- tests/fast/ray/multi_lora/test_controller.py | 68 +++ .../rollout/real_ray/test_rollout_executor.py | 414 +++++++++++++++++- tests/fast/rollout/test_checkpoint_eval.py | 70 +++ tests/fast/utils/fake_ray_ids.py | 8 + tests/fast/utils/test_http_utils.py | 108 +++++ tests/fast/utils/tracking_utils/__init__.py | 0 .../tracking_utils/test_prometheus_utils.py | 103 +++++ 20 files changed, 897 insertions(+), 46 deletions(-) create mode 100644 tests/fast/ray/multi_lora/test_controller.py create mode 100644 tests/fast/utils/fake_ray_ids.py create mode 100644 tests/fast/utils/tracking_utils/__init__.py create mode 100644 tests/fast/utils/tracking_utils/test_prometheus_utils.py diff --git a/miles/dashboard/backend.py b/miles/dashboard/backend.py index 4ee08dd5706..7ad4f1e5306 100644 --- a/miles/dashboard/backend.py +++ b/miles/dashboard/backend.py @@ -62,7 +62,9 @@ def init_dashboard(args, *, primary: bool = True, router_addr: str | None = None node_id=ray.get_runtime_context().get_node_id(), soft=False ), ) - .remote(config, prometheus_handle_factory=_prometheus_factory if config.forward_prometheus else None) + .remote( + config=config, prometheus_handle_factory=_prometheus_factory if config.forward_prometheus else None + ) ) ray.get(_handle.ping.remote()) _handle.start.remote() diff --git a/miles/dashboard/collector.py b/miles/dashboard/collector.py index 2e8901b1fb9..6acafa08e33 100644 --- a/miles/dashboard/collector.py +++ b/miles/dashboard/collector.py @@ -98,7 +98,7 @@ def _default_spawn_sampler(node_id: str, node_ip: str, interval: float): scheduling_strategy=NodeAffinitySchedulingStrategy(node_id=node_id, soft=False), ) .remote( - _SelfGpuPush(ray.get_runtime_context().current_actor), + push=_SelfGpuPush(ray.get_runtime_context().current_actor), node=node_ip, interval=interval, push_processes=_SelfGpuProcessPush(ray.get_runtime_context().current_actor), @@ -143,8 +143,8 @@ class DashboardCollector: def __init__( self, - config: CollectorConfig, *, + config: CollectorConfig, prometheus_handle_factory=None, # () -> handle with .update.remote(dict), or None scraper_http_get=None, # test hook, forwarded to SglangScraper ): diff --git a/miles/dashboard/gpu_sampler.py b/miles/dashboard/gpu_sampler.py index 939302066f9..d7fcb5d7fc5 100644 --- a/miles/dashboard/gpu_sampler.py +++ b/miles/dashboard/gpu_sampler.py @@ -126,8 +126,8 @@ class GpuSampler: def __init__( self, - push: Callable[[str, list[GpuSample]], None], *, + push: Callable[[str, list[GpuSample]], None], node: str, interval: float = 1.0, nvml=None, diff --git a/miles/ray/multi_lora/controller.py b/miles/ray/multi_lora/controller.py index 6fb300ddd2a..e44ef8b3b2e 100644 --- a/miles/ray/multi_lora/controller.py +++ b/miles/ray/multi_lora/controller.py @@ -59,7 +59,7 @@ def _load_subclass(path: str | None, base_cls): @ray.remote(num_cpus=0) class MultiLoRAController: - def __init__(self, args, router_url: str, host: str = "0.0.0.0") -> None: + def __init__(self, *, args, router_url: str, host: str = "0.0.0.0") -> None: backend_cls = _load_subclass(getattr(args, "multi_lora_backend_path", None), MultiLoRABackend) server_cls = _load_subclass(getattr(args, "multi_lora_http_server_path", None), MultiLoRAHTTPServer) self.backend = backend_cls(args, router_url) @@ -120,4 +120,4 @@ def create_multilora_controller(args, router_url: str, host: str = "0.0.0.0"): name=CONTROLLER_NAME, namespace=CONTROLLER_NAMESPACE, **compute_ray_pin_head_options(), - ).remote(args, router_url, host) + ).remote(args=args, router_url=router_url, host=host) diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index 1d1b153759a..702fbf4225a 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -208,7 +208,7 @@ async def create_rollout_components(args, pg) -> RolloutComponents: rollout_executor = RolloutExecutor.options( num_cpus=1, num_gpus=0, **(compute_ray_pin_head_options() if args.pin_rollout_manager_to_head else {}) - ).remote(args) + ).remote(args=args) # calculate num_rollout from num_epoch num_rollout_per_epoch = None diff --git a/miles/ray/rollout/rollout_executor.py b/miles/ray/rollout/rollout_executor.py index 72aabd0f940..bb19e67f956 100644 --- a/miles/ray/rollout/rollout_executor.py +++ b/miles/ray/rollout/rollout_executor.py @@ -47,7 +47,7 @@ class RolloutExecutor: """The class to run rollout and convert rollout data to training data.""" - def __init__(self, args): + def __init__(self, *, args): event_logger_checkpoint.restore(args) configure_logger(args, source=RolloutExecutorProcessIdentity()) diff --git a/miles/ray/train/actor_factory.py b/miles/ray/train/actor_factory.py index 7e1706676a2..2ef03917012 100644 --- a/miles/ray/train/actor_factory.py +++ b/miles/ray/train/actor_factory.py @@ -93,11 +93,11 @@ def allocate_gpus_for_actor( rank_dir = os.path.join(args.offload_train_disk_dir, f"cell{cell_index}_rank{rank}") options["runtime_env"] = {"env_vars": {**env_vars, "TMS_DISK_BACKUP_DIR": rank_dir}} actor = TrainRayActor.options(**options).remote( - args, - world_size, - rank, - master_addr, - master_port, + args=args, + world_size=world_size, + rank=rank, + master_addr=master_addr, + master_port=master_port, indep_dp_store_addr=indep_dp_store_addr, role=role, cell_index=cell_index, diff --git a/miles/ray/train_actor.py b/miles/ray/train_actor.py index 248f53cca65..942a556bd48 100644 --- a/miles/ray/train_actor.py +++ b/miles/ray/train_actor.py @@ -39,6 +39,7 @@ def get_local_gpu_id(): class TrainRayActor: def __init__( self, + *, args, world_size: int, rank: int, diff --git a/miles/utils/http_utils.py b/miles/utils/http_utils.py index ec3f44ba978..6dc2b8d42d9 100644 --- a/miles/utils/http_utils.py +++ b/miles/utils/http_utils.py @@ -345,7 +345,7 @@ def _init_ray_distributed_post(args): # Define the async actor @ray.remote class _HttpPosterActor: - def __init__(self, concurrency: int): + def __init__(self, *, concurrency: int): # Lazy creation to this actor's event loop self._client = httpx.AsyncClient( limits=httpx.Limits(max_connections=max(1, concurrency)), @@ -371,7 +371,7 @@ async def do_post(self, url, payload, max_retries=60, action="post", headers=Non max_concurrency=per_actor_conc, # Use tiny CPU to schedule num_cpus=0.001, - ).remote(per_actor_conc) + ).remote(concurrency=per_actor_conc) created.append(actor) _post_actors = created diff --git a/miles/utils/tracking_utils/prometheus_utils.py b/miles/utils/tracking_utils/prometheus_utils.py index 60e9bfd0ee9..5399cdf68bd 100644 --- a/miles/utils/tracking_utils/prometheus_utils.py +++ b/miles/utils/tracking_utils/prometheus_utils.py @@ -40,7 +40,7 @@ def init_prometheus(args, start_server: bool = False): soft=False, ), ) - .remote(args) + .remote(args=args) ) ray.get(_collector_handle.ping.remote()) logger.info("Prometheus collector actor created") @@ -74,7 +74,7 @@ class _PrometheusCollector: Ray handles the RPC transparently. """ - def __init__(self, args): + def __init__(self, *, args): from prometheus_client import Gauge, start_http_server self._Gauge = Gauge diff --git a/tests/fast-gpu/test_gpu_sampler_hw.py b/tests/fast-gpu/test_gpu_sampler_hw.py index fc5982c716a..ba09e3bffc6 100644 --- a/tests/fast-gpu/test_gpu_sampler_hw.py +++ b/tests/fast-gpu/test_gpu_sampler_hw.py @@ -17,7 +17,7 @@ def __call__(self, node, batch): def test_auto_detection_picks_the_backend_matching_the_hardware(): - sampler = GpuSampler(PushSpy(), node="ci") + sampler = GpuSampler(push=PushSpy(), node="ci") assert sampler.available, "no GPU telemetry backend initialized on a GPU runner" expected = "AMD SMI" if torch.version.hip else "NVML" assert sampler._provider.name == expected @@ -26,7 +26,7 @@ def test_auto_detection_picks_the_backend_matching_the_hardware(): def test_every_device_reports_telemetry_and_processes(): push = PushSpy() push_processes = PushSpy() - sampler = GpuSampler(push, node="ci", push_processes=push_processes) + sampler = GpuSampler(push=push, node="ci", push_processes=push_processes) assert sampler.available uuids = sampler.gpu_uuids() diff --git a/tests/fast/dashboard/test_collector.py b/tests/fast/dashboard/test_collector.py index fc5077f5366..3c795e73715 100644 --- a/tests/fast/dashboard/test_collector.py +++ b/tests/fast/dashboard/test_collector.py @@ -1,7 +1,9 @@ import logging import time from pathlib import Path +from types import SimpleNamespace +import pytest from tests.fast.dashboard.dummy_telemetry import BASE_TS, dump_dummy_telemetry from miles.dashboard.collector import CollectorConfig, DashboardCollector @@ -26,7 +28,7 @@ def make_collector(tmp_path, **kwargs) -> DashboardCollector: config = kwargs.pop("config", None) or CollectorConfig( dashboard_dir=str(tmp_path / "dashboard"), run_name="collector-test", start_ts=1.0 ) - return DashboardCollector(config, **kwargs) + return DashboardCollector(config=config, **kwargs) def test_collector_satisfies_dummy_telemetry_contract(tmp_path): @@ -114,7 +116,7 @@ def test_flush_thread_persists_periodically(tmp_path): config = CollectorConfig( dashboard_dir=str(tmp_path / "dashboard"), run_name="r", start_ts=0.0, flush_interval_seconds=0.05 ) - collector = DashboardCollector(config) + collector = DashboardCollector(config=config) collector.start() collector.push_metrics(MetricsRecord(ts=1.0, step_key="rollout/step", step=0, metrics={"a": 1})) time.sleep(0.2) @@ -208,7 +210,7 @@ class FakeHandle: config = CollectorConfig( dashboard_dir=str(tmp_path / "dashboard"), run_name="r", start_ts=0.0, forward_prometheus=True ) - collector = DashboardCollector(config, prometheus_handle_factory=lambda: FakeHandle()) + collector = DashboardCollector(config=config, prometheus_handle_factory=lambda: FakeHandle()) collector.push_gpu_samples( "10.0.0.1", [GpuSample(ts=1.0, node="10.0.0.1", gpu=0, util=87, mem_mb=1000, power_w=600)] ) @@ -313,3 +315,71 @@ def test_external_engines_synthesized_from_scrapes(tmp_path): "http://10.1.0.5:15000": "regular", "http://10.1.0.6:15000": "external", } + + +class TestConstructorContract: + def test_a_positional_config_is_rejected(self, tmp_path): + """The collector is built as a ray actor, so config must be bound by keyword only.""" + config = CollectorConfig(dashboard_dir=str(tmp_path / "dashboard"), run_name="r", start_ts=1.0) + + with pytest.raises(TypeError): + DashboardCollector(config) + + +class _FakeSamplerStart: + def remote(self): + return True + + +class _FakeSamplerHandle: + def __init__(self, sampler): + self.sampler = sampler + self.start = _FakeSamplerStart() + + +class _FakeRayActorClass: + def __init__(self, cls): + self._cls = cls + self.options_kwargs = None + self.positional_args = None + self.keyword_args = None + + def options(self, **kwargs): + self.options_kwargs = kwargs + return self + + def remote(self, *args, **kwargs): + self.positional_args = args + self.keyword_args = kwargs + return _FakeSamplerHandle(self._cls(*args, **kwargs)) + + +class TestDefaultSpawnSampler: + def test_builds_the_gpu_sampler_through_keyword_arguments_only(self, monkeypatch): + """The sampler actor is constructed with keywords its keyword-only signature accepts.""" + import ray + + from miles.dashboard import collector as collector_mod + from miles.dashboard.gpu_sampler import GpuSampler + + actor_classes: list[_FakeRayActorClass] = [] + + def fake_ray_remote(cls): + actor_class = _FakeRayActorClass(cls) + actor_classes.append(actor_class) + return actor_class + + monkeypatch.setattr(ray, "remote", fake_ray_remote) + monkeypatch.setattr(ray, "get", lambda value: value) + monkeypatch.setattr(ray, "kill", lambda handle: None) + monkeypatch.setattr(ray, "get_runtime_context", lambda: SimpleNamespace(current_actor="collector-handle")) + + handle = collector_mod._default_spawn_sampler("a" * 56, "10.0.0.7", 0.25) + + [actor_class] = actor_classes + assert actor_class.positional_args == () + assert set(actor_class.keyword_args) == {"push", "node", "interval", "push_processes"} + sampler = handle.sampler + assert isinstance(sampler, GpuSampler) + assert (sampler.node, sampler.interval) == ("10.0.0.7", 0.25) + assert isinstance(sampler._push, collector_mod._SelfGpuPush) diff --git a/tests/fast/dashboard/test_gpu_sampler.py b/tests/fast/dashboard/test_gpu_sampler.py index 7322d6dd60f..a0dc9cc9ed2 100644 --- a/tests/fast/dashboard/test_gpu_sampler.py +++ b/tests/fast/dashboard/test_gpu_sampler.py @@ -117,7 +117,7 @@ def __call__(self, node, batch): def test_sample_once_converts_units(): push = PushSpy() - sampler = GpuSampler(push, node="10.0.0.1", nvml=FakeNvml(count=2)) + sampler = GpuSampler(push=push, node="10.0.0.1", nvml=FakeNvml(count=2)) assert sampler.available assert sampler.gpu_uuids() == ["GPU-fake-0", "GPU-fake-1"] @@ -133,7 +133,7 @@ def test_sample_once_converts_units(): def test_amd_sample_once_preserves_native_units_and_uuids(): push = PushSpy() - sampler = GpuSampler(push, node="amd-node", amdsmi=FakeAmdSmi(count=2)) + sampler = GpuSampler(push=push, node="amd-node", amdsmi=FakeAmdSmi(count=2)) assert sampler.available assert sampler.gpu_uuids() == ["GPU-amd-0", "GPU-amd-1"] @@ -173,7 +173,7 @@ def test_amd_unavailable_power_reports_zero_but_keeps_util_and_mem(): push = PushSpy() amdsmi = FakeAmdSmi(count=1) amdsmi.amdsmi_get_power_info = lambda handle: {"socket_power": "N/A", "current_socket_power": "N/A"} - sampler = GpuSampler(push, node="n", amdsmi=amdsmi) + sampler = GpuSampler(push=push, node="n", amdsmi=amdsmi) assert sampler.sample_once(ts=1.0) == 1 sampler.flush() @@ -183,7 +183,7 @@ def test_amd_unavailable_power_reports_zero_but_keeps_util_and_mem(): def test_flush_clears_buffer_and_skips_empty(): push = PushSpy() - sampler = GpuSampler(push, node="n", nvml=FakeNvml(count=1)) + sampler = GpuSampler(push=push, node="n", nvml=FakeNvml(count=1)) sampler.flush() # empty: no call assert push.calls == [] @@ -196,7 +196,7 @@ def test_flush_clears_buffer_and_skips_empty(): def test_nvml_init_failure_disables_sampler(caplog): push = PushSpy() with caplog.at_level(logging.WARNING): - sampler = GpuSampler(push, node="n", nvml=FakeNvml(fail_init=True)) + sampler = GpuSampler(push=push, node="n", nvml=FakeNvml(fail_init=True)) assert not sampler.available assert sampler.start() is False assert sampler.sample_once(ts=1.0) == 0 @@ -207,7 +207,7 @@ def test_nvml_init_failure_disables_sampler(caplog): @pytest.mark.parametrize("backend", ["nvml", "amdsmi"]) def test_zero_devices_disable_sampler(backend): fake = FakeNvml(count=0) if backend == "nvml" else FakeAmdSmi(count=0) - sampler = GpuSampler(PushSpy(), node="n", **{backend: fake}) + sampler = GpuSampler(push=PushSpy(), node="n", **{backend: fake}) assert not sampler.available assert sampler.gpu_uuids() == [] @@ -216,7 +216,7 @@ def test_production_auto_detection_falls_back_from_nvml_to_amdsmi(monkeypatch): monkeypatch.setitem(sys.modules, "pynvml", FakeNvml(fail_init=True)) monkeypatch.setitem(sys.modules, "amdsmi", FakeAmdSmi(count=1)) - sampler = GpuSampler(PushSpy(), node="n") + sampler = GpuSampler(push=PushSpy(), node="n") assert sampler.available assert sampler.gpu_uuids() == ["GPU-amd-0"] @@ -229,7 +229,7 @@ def amdsmi_init(self): monkeypatch.setitem(sys.modules, "amdsmi", BoomAmdSmi()) - sampler = GpuSampler(PushSpy(), node="n", nvml=FakeNvml(fail_init=True)) + sampler = GpuSampler(push=PushSpy(), node="n", nvml=FakeNvml(fail_init=True)) assert not sampler.available @@ -239,7 +239,7 @@ def test_missing_optional_backends_disable_sampler(monkeypatch, caplog): monkeypatch.setitem(sys.modules, "pynvml", None) monkeypatch.setitem(sys.modules, "amdsmi", None) with caplog.at_level(logging.WARNING): - sampler = GpuSampler(PushSpy(), node="n") + sampler = GpuSampler(push=PushSpy(), node="n") assert not sampler.available assert sampler.start() is False @@ -248,12 +248,12 @@ def test_missing_optional_backends_disable_sampler(monkeypatch, caplog): def test_only_one_backend_may_be_injected(): with pytest.raises(AssertionError, match="inject only one"): - GpuSampler(PushSpy(), node="n", nvml=FakeNvml(), amdsmi=FakeAmdSmi()) + GpuSampler(push=PushSpy(), node="n", nvml=FakeNvml(), amdsmi=FakeAmdSmi()) def test_failing_device_is_skipped_others_report(caplog): push = PushSpy() - sampler = GpuSampler(push, node="n", nvml=FakeNvml(count=3, failing_devices={1})) + sampler = GpuSampler(push=push, node="n", nvml=FakeNvml(count=3, failing_devices={1})) with caplog.at_level(logging.WARNING): assert sampler.sample_once(ts=1.0) == 2 sampler.flush() @@ -264,7 +264,7 @@ def test_failing_device_is_skipped_others_report(caplog): def test_amd_failing_device_is_skipped_while_others_report(caplog): push = PushSpy() - sampler = GpuSampler(push, node="n", amdsmi=FakeAmdSmi(count=3, failing_devices={1})) + sampler = GpuSampler(push=push, node="n", amdsmi=FakeAmdSmi(count=3, failing_devices={1})) with caplog.at_level(logging.WARNING): assert sampler.sample_once(ts=1.0) == 2 sampler.flush() @@ -285,7 +285,7 @@ def test_amd_degraded_metric_value_skips_device_while_others_report(method, payl amdsmi = FakeAmdSmi(count=2) original = getattr(amdsmi, method) setattr(amdsmi, method, lambda handle: payload if handle == 0 else original(handle)) - sampler = GpuSampler(push, node="n", amdsmi=amdsmi) + sampler = GpuSampler(push=push, node="n", amdsmi=amdsmi) with caplog.at_level(logging.WARNING): assert sampler.sample_once(ts=1.0) == 1 @@ -298,7 +298,7 @@ def test_amd_uses_smi_visible_order_without_refiltering_process_env(monkeypatch) monkeypatch.setenv("HIP_VISIBLE_DEVICES", "2") monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "2") push = PushSpy() - sampler = GpuSampler(push, node="n", amdsmi=FakeAmdSmi(count=3)) + sampler = GpuSampler(push=push, node="n", amdsmi=FakeAmdSmi(count=3)) assert sampler.sample_once(ts=1.0) == 3 sampler.flush() @@ -307,7 +307,7 @@ def test_amd_uses_smi_visible_order_without_refiltering_process_env(monkeypatch) def test_thread_lifecycle_flushes_on_stop(): push = PushSpy() - sampler = GpuSampler(push, node="n", interval=0.01, nvml=FakeNvml(count=1)) + sampler = GpuSampler(push=push, node="n", interval=0.01, nvml=FakeNvml(count=1)) assert sampler.start() is True time.sleep(0.08) sampler.stop() @@ -319,7 +319,7 @@ def test_thread_lifecycle_flushes_on_stop(): def test_sample_processes_once_converts_units(): push = PushSpy() push_processes = ProcessPushSpy() - sampler = GpuSampler(push, node="n", nvml=FakeNvml(count=2), push_processes=push_processes) + sampler = GpuSampler(push=push, node="n", nvml=FakeNvml(count=2), push_processes=push_processes) assert sampler.sample_processes_once(ts=5.0) == 2 sampler.flush() [(node, batch)] = push_processes.calls @@ -332,7 +332,7 @@ def test_sample_processes_once_converts_units(): def test_amd_processes_convert_bytes_fall_back_on_name_and_drop_zero_vram(): push_processes = ProcessPushSpy() - sampler = GpuSampler(PushSpy(), node="n", amdsmi=FakeAmdSmi(count=2), push_processes=push_processes) + sampler = GpuSampler(push=PushSpy(), node="n", amdsmi=FakeAmdSmi(count=2), push_processes=push_processes) assert sampler.sample_processes_once(ts=5.0) == 2 sampler.flush() @@ -348,7 +348,9 @@ def test_amd_processes_convert_bytes_fall_back_on_name_and_drop_zero_vram(): def test_failing_device_skipped_for_process_sampling(caplog): push = PushSpy() push_processes = ProcessPushSpy() - sampler = GpuSampler(push, node="n", nvml=FakeNvml(count=3, failing_devices={1}), push_processes=push_processes) + sampler = GpuSampler( + push=push, node="n", nvml=FakeNvml(count=3, failing_devices={1}), push_processes=push_processes + ) with caplog.at_level(logging.WARNING): assert sampler.sample_processes_once(ts=1.0) == 2 sampler.flush() @@ -360,7 +362,7 @@ def test_failing_device_skipped_for_process_sampling(caplog): def test_amd_failing_device_is_skipped_for_process_sampling(caplog): push_processes = ProcessPushSpy() sampler = GpuSampler( - PushSpy(), + push=PushSpy(), node="n", amdsmi=FakeAmdSmi(count=3, failing_devices={1}), push_processes=push_processes, @@ -375,7 +377,7 @@ def test_amd_failing_device_is_skipped_for_process_sampling(caplog): def test_process_batch_dropped_silently_without_push_processes(): push = PushSpy() - sampler = GpuSampler(push, node="n", nvml=FakeNvml(count=1)) + sampler = GpuSampler(push=push, node="n", nvml=FakeNvml(count=1)) assert sampler.sample_processes_once(ts=1.0) == 1 sampler.flush() assert push.calls == [] @@ -383,7 +385,7 @@ def test_process_batch_dropped_silently_without_push_processes(): def test_interval_must_be_positive(): with pytest.raises(AssertionError): - GpuSampler(lambda n, b: None, node="n", interval=0, nvml=FakeNvml()) + GpuSampler(push=lambda n, b: None, node="n", interval=0, nvml=FakeNvml()) def test_real_nvml_when_gpus_present(): @@ -398,7 +400,7 @@ def test_real_nvml_when_gpus_present(): push = PushSpy() push_processes = ProcessPushSpy() - sampler = GpuSampler(push, node="local", nvml=pynvml, push_processes=push_processes) + sampler = GpuSampler(push=push, node="local", nvml=pynvml, push_processes=push_processes) assert sampler.available assert sampler.sample_once(ts=1.0) >= 1 # idle test GPUs may have zero compute processes — asserting >= 0 just @@ -419,7 +421,7 @@ def test_real_amdsmi_when_gpus_present(): amdsmi = pytest.importorskip("amdsmi") push = PushSpy() push_processes = ProcessPushSpy() - sampler = GpuSampler(push, node="local", amdsmi=amdsmi, push_processes=push_processes) + sampler = GpuSampler(push=push, node="local", amdsmi=amdsmi, push_processes=push_processes) if not sampler.available: pytest.skip("no usable AMD SMI device") @@ -434,3 +436,10 @@ def test_real_amdsmi_when_gpus_present(): if push_processes.calls: proc_sample = push_processes.calls[0][1][0] assert proc_sample.pid > 0 and proc_sample.mem_mb > 0 and proc_sample.name + + +class TestConstructorContract: + def test_a_positional_push_is_rejected(self): + """The sampler is built as a ray actor, so push must be bound by keyword only.""" + with pytest.raises(TypeError): + GpuSampler(PushSpy(), node="n", nvml=FakeNvml()) diff --git a/tests/fast/ray/multi_lora/test_controller.py b/tests/fast/ray/multi_lora/test_controller.py new file mode 100644 index 00000000000..aa8e6b059f7 --- /dev/null +++ b/tests/fast/ray/multi_lora/test_controller.py @@ -0,0 +1,68 @@ +from types import SimpleNamespace + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=60, suite="stage-a-cpu") + +import pytest + +import miles.ray.multi_lora.controller as controller_mod +from miles.ray.multi_lora.backend import MultiLoRABackend +from miles.ray.multi_lora.controller import CONTROLLER_NAME, CONTROLLER_NAMESPACE, MultiLoRAController + + +def make_args() -> SimpleNamespace: + return SimpleNamespace(multi_lora_n_adapters=4) + + +class TestConstructorContract: + def test_keyword_arguments_build_the_backend_and_the_http_server(self): + """Keyword construction is the supported form and wires up both owned components.""" + args = make_args() + + controller = MultiLoRAController.__ray_actor_class__(args=args, router_url="http://router:1", host="10.0.0.9") + + assert isinstance(controller.backend, MultiLoRABackend) + assert controller.backend.router_url == "http://router:1" + assert controller.server.host == "10.0.0.9" + + def test_positional_arguments_are_rejected(self): + """The actor is launched by keyword only, so positional construction must fail loudly.""" + with pytest.raises(TypeError): + MultiLoRAController.__ray_actor_class__(make_args(), "http://router:1") + + +class _RecordingActorClass: + def __init__(self): + self.options_kwargs = None + self.positional_args = None + self.keyword_args = None + + def options(self, **kwargs): + self.options_kwargs = kwargs + return self + + def remote(self, *args, **kwargs): + self.positional_args = args + self.keyword_args = kwargs + return "controller-handle" + + +class TestCreateMultiLoRAController: + def test_launches_a_named_head_pinned_actor_with_keyword_arguments(self, monkeypatch): + """The factory names the actor, pins it to the head node and passes every argument by keyword.""" + recorder = _RecordingActorClass() + monkeypatch.setattr(controller_mod, "MultiLoRAController", recorder) + monkeypatch.setattr(controller_mod, "compute_ray_pin_head_options", lambda: {"scheduling_strategy": "head"}) + args = make_args() + + handle = controller_mod.create_multilora_controller(args, "http://router:1", "10.0.0.9") + + assert handle == "controller-handle" + assert recorder.options_kwargs == { + "name": CONTROLLER_NAME, + "namespace": CONTROLLER_NAMESPACE, + "scheduling_strategy": "head", + } + assert recorder.positional_args == () + assert recorder.keyword_args == {"args": args, "router_url": "http://router:1", "host": "10.0.0.9"} diff --git a/tests/fast/ray/rollout/real_ray/test_rollout_executor.py b/tests/fast/ray/rollout/real_ray/test_rollout_executor.py index 133030b2389..e3d4a38bc54 100644 --- a/tests/fast/ray/rollout/real_ray/test_rollout_executor.py +++ b/tests/fast/ray/rollout/real_ray/test_rollout_executor.py @@ -1,11 +1,14 @@ from __future__ import annotations +import logging +from types import SimpleNamespace from unittest.mock import MagicMock import pytest import ray from tests.fast.ray.rollout.conftest import make_args, make_samples_grouped +from miles.ray.rollout.debug_data import save_debug_rollout_data from miles.ray.rollout.rollout_executor import RolloutExecutor from miles.rollout.base_types import ( BaseRolloutFn, @@ -14,6 +17,7 @@ RolloutFnTrainInput, RolloutFnTrainOutput, ) +from miles.rollout.checkpoint_eval import CheckpointEvalFn from miles.utils.types import WeightVersionSpan, WeightVersionsPerCall @@ -40,7 +44,7 @@ def patch_low_level(monkeypatch, http_client_calls): def _make_executor(args): - return RolloutExecutor.__ray_actor_class__(args) + return RolloutExecutor.__ray_actor_class__(args=args) def _make_test_args(**overrides): @@ -215,6 +219,18 @@ def load(self, rollout_id: int | None) -> None: self._log.append((self._name, "load", rollout_id)) +class _RecordingEventLoggerCheckpoint: + def __init__(self) -> None: + self.restored: list = [] + self.snapshots: list[tuple] = [] + + def restore(self, args) -> None: + self.restored.append(args) + + def snapshot(self, args, rollout_id) -> None: + self.snapshots.append((args, rollout_id)) + + @pytest.mark.asyncio class TestCheckpointing: async def test_save_and_load_reach_every_distinct_rollout_function( @@ -319,6 +335,29 @@ async def test_legacy_function_path_does_not_get_save_load( executor.data_source.load.assert_called_once_with(1) + async def test_save_snapshots_the_event_log_for_the_rollout_id( + self, + ray_local_mode, + patch_low_level, + monkeypatch, + ): + """The audit event log is checkpointed alongside the rollout state it explains.""" + import miles.ray.rollout.rollout_executor as rexec + + recorder = _RecordingEventLoggerCheckpoint() + monkeypatch.setattr(rexec, "event_logger_checkpoint", recorder) + args = _make_test_args(rollout_global_dataset=False) + + executor = _make_executor(args) + executor.use_legacy_rollout_v1 = False + executor.generate_rollout = _RecordingRolloutFn("train", []) + executor.eval_generate_rollout = None + executor.data_source = MagicMock() + + executor.save(rollout_id=9) + + assert recorder.snapshots == [(args, 9)] + @pytest.mark.asyncio class TestEval: @@ -355,6 +394,379 @@ async def test_skipped_in_debug_train_only_mode(self, ray_local_mode, patch_low_ assert called == [] +class _NamedRolloutFn(BaseRolloutFn): + def __init__(self, path: str) -> None: + self.path = path + + def __call__(self, input): + raise AssertionError("not exercised by the loading tests") + + +@pytest.mark.asyncio +class TestRolloutFunctionLoading: + async def test_matching_train_and_eval_paths_share_one_rollout_instance( + self, ray_local_mode, patch_low_level, monkeypatch + ): + """One configured path means one stateful object, so train and eval share its state.""" + import miles.ray.rollout.rollout_executor as rexec + + monkeypatch.setattr(rexec, "load_rollout_function", lambda input, path: _NamedRolloutFn(path)) + args = _make_test_args(rollout_function_path="pkg.same_fn", eval_function_path="pkg.same_fn") + + executor = _make_executor(args) + + assert executor.eval_generate_rollout is executor.generate_rollout + + async def test_distinct_train_and_eval_paths_get_their_own_instance( + self, ray_local_mode, patch_low_level, monkeypatch + ): + """Two configured paths must each be loaded, so eval never silently runs the train function.""" + import miles.ray.rollout.rollout_executor as rexec + + monkeypatch.setattr(rexec, "load_rollout_function", lambda input, path: _NamedRolloutFn(path)) + args = _make_test_args(rollout_function_path="pkg.train_fn", eval_function_path="pkg.eval_fn") + + executor = _make_executor(args) + + assert executor.generate_rollout.path == "pkg.train_fn" + assert executor.eval_generate_rollout.path == "pkg.eval_fn" + + +@pytest.mark.asyncio +class TestCustomHooks: + async def test_get_converts_the_batch_with_the_configured_conversion_hook( + self, ray_local_mode, patch_low_level, monkeypatch + ): + """The configured conversion hook receives the generated samples and its output is what ships.""" + import miles.ray.rollout.rollout_executor as rexec + + seen_samples: list = [] + + def conversion_hook(args, samples): + seen_samples.append(samples) + return dict(tokens=[[1, 2], [3, 4]], sample_indices=[70, 71]) + + monkeypatch.setattr( + rexec, + "load_function", + lambda path: conversion_hook if path == "pkg.convert" else (lambda *a, **kw: None), + ) + args = _make_test_args(global_batch_size=4, custom_convert_samples_to_train_data_path="pkg.convert") + + executor = _make_executor(args) + executor.set_train_parallel_config({"dp_size": 2}) + executor.generate_rollout = lambda input: RolloutFnTrainOutput( + samples=[make_samples_grouped(n_groups=1, group_size=4)], metrics={} + ) + + result = await executor.get(rollout_id=1) + + assert [sample.index for sample in seen_samples[0]] == [0, 1, 2, 3] + assert result["sample_indices"] == [70, 71] + partitions = ray.get([box.inner for box in result["data_ref"]]) + assert [partition["tokens"] for partition in partitions] == [[[1, 2]], [[3, 4]]] + + async def test_get_scores_the_batch_with_the_configured_reward_hook( + self, ray_local_mode, patch_low_level, monkeypatch + ): + """The configured reward post-process hook replaces the built-in reward math.""" + import miles.ray.rollout.rollout_executor as rexec + + seen_samples: list = [] + + def reward_hook(args, samples): + seen_samples.append(samples) + return [1.0] * len(samples), [9.0] * len(samples) + + monkeypatch.setattr( + rexec, + "load_function", + lambda path: reward_hook if path == "pkg.reward" else (lambda *a, **kw: None), + ) + args = _make_test_args(global_batch_size=4, custom_reward_post_process_path="pkg.reward") + + executor = _make_executor(args) + executor.set_train_parallel_config({"dp_size": 2}) + executor.generate_rollout = lambda input: RolloutFnTrainOutput( + samples=[make_samples_grouped(n_groups=1, group_size=4)], metrics={} + ) + + result = await executor.get(rollout_id=1) + + assert [sample.index for sample in seen_samples[0]] == [0, 1, 2, 3] + partitions = ray.get([box.inner for box in result["data_ref"]]) + assert [partition["rewards"] for partition in partitions] == [[9.0, 9.0], [9.0, 9.0]] + + +@pytest.mark.asyncio +class TestWeightVersion: + async def test_get_threads_the_latest_weight_version_into_the_train_input(self, ray_local_mode, patch_low_level): + """The rollout function is told which engine weight version this batch is generated under.""" + args = _make_test_args(global_batch_size=4, indep_dp=False) + + executor = _make_executor(args) + executor.set_train_parallel_config({"dp_size": 1}) + captured: list = [] + + def fake_rollout_fn(input): + captured.append(input) + return RolloutFnTrainOutput(samples=[make_samples_grouped(n_groups=1, group_size=4)], metrics={}) + + executor.generate_rollout = fake_rollout_fn + executor.set_weight_version(3) + executor.set_weight_version(7) + + await executor.get(rollout_id=1) + + assert captured[0].weight_version == 7 + + async def test_a_decreasing_weight_version_is_rejected(self, ray_local_mode, patch_low_level): + """Weight versions only move forward, so a lower one signals a broken update path.""" + executor = _make_executor(_make_test_args(indep_dp=False)) + executor.set_weight_version(7) + + with pytest.raises(AssertionError, match="went backwards"): + executor.set_weight_version(3) + + assert executor.weight_version == 7 + + async def test_independent_dp_accepts_a_decreasing_weight_version_with_a_warning( + self, ray_local_mode, patch_low_level, caplog + ): + """Independent-DP fault tolerance may rewind a replica, so the rewind warns instead of failing.""" + executor = _make_executor(_make_test_args(indep_dp=True)) + executor.set_weight_version(7) + + with caplog.at_level(logging.WARNING): + executor.set_weight_version(3) + + assert executor.weight_version == 3 + assert any("went backwards" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +class TestDelayedDpSplit: + async def test_get_stores_one_unsplit_batch(self, ray_local_mode, patch_low_level): + """With the split delayed to the training side, one whole-batch reference is published.""" + args = _make_test_args(global_batch_size=4, delay_split_train_data_by_dp=True) + + executor = _make_executor(args) + executor.set_train_parallel_config({"dp_size": 2}) + executor.generate_rollout = lambda input: RolloutFnTrainOutput( + samples=[make_samples_grouped(n_groups=1, group_size=4)], metrics={} + ) + + result = await executor.get(rollout_id=1) + + assert not isinstance(result["data_ref"], list) + stored = ray.get(result["data_ref"].inner) + assert len(stored["tokens"]) == 4 + assert result["sample_indices"] == stored["sample_indices"] == [0, 1, 2, 3] + + +@pytest.mark.asyncio +class TestDebugRolloutData: + async def test_recorded_data_is_replayed_instead_of_generated(self, ray_local_mode, patch_low_level, tmp_path): + """Replaying a recording must not call the rollout function, and its metadata comes back verbatim.""" + template = str(tmp_path / "rollout_{rollout_id}.pt") + save_debug_rollout_data( + make_args(save_debug_rollout_data=template), + make_samples_grouped(n_groups=1, group_size=4), + rollout_id=3, + evaluation=False, + metadata={"dynamic_global_batch_size": 4}, + ) + args = _make_test_args(load_debug_rollout_data=template) + + executor = _make_executor(args) + + def unreachable(input): + raise AssertionError("the rollout function must not run when a recording is replayed") + + executor.generate_rollout = unreachable + + data, metadata, metrics = await executor._get_rollout_data(rollout_id=3) + + assert [sample.index for sample in data] == [0, 1, 2, 3] + assert metadata == {"dynamic_global_batch_size": 4} + assert metrics is None + + async def test_injected_recording_replaces_the_verified_generated_batch( + self, ray_local_mode, patch_low_level, tmp_path + ): + """CI injection swaps in the recording it verified against, and drops the generated metrics with it.""" + template = str(tmp_path / "inject_{rollout_id}.pt") + injected = make_samples_grouped(n_groups=1, group_size=4) + for sample in injected: + sample.reward = 5.0 + save_debug_rollout_data( + make_args(save_debug_rollout_data=template), + injected, + rollout_id=2, + evaluation=False, + metadata={"source": "recording"}, + ) + args = _make_test_args( + global_batch_size=4, + ci_inject_rollout_data_path=template, + ci_inject_rollout_data_start_rollout_id=2, + ) + + executor = _make_executor(args) + executor.set_train_parallel_config({"dp_size": 1}) + + mismatched = make_samples_grouped(n_groups=1, group_size=4) + for sample in mismatched: + sample.tokens[-1] = 99 + executor.generate_rollout = lambda input: RolloutFnTrainOutput( + samples=[mismatched], metrics={"generated": 1.0} + ) + + with pytest.raises(AssertionError, match="generated responses match"): + await executor._get_rollout_data(rollout_id=2) + + executor.generate_rollout = lambda input: RolloutFnTrainOutput( + samples=[make_samples_grouped(n_groups=1, group_size=4)], metrics={"generated": 1.0} + ) + data, metadata, metrics = await executor._get_rollout_data(rollout_id=2) + + assert [sample.reward for sample in data] == [5.0, 5.0, 5.0, 5.0] + assert metadata == {"source": "recording"} + assert metrics is None + + +@pytest.mark.asyncio +class TestLegacyRolloutProtocol: + async def test_train_generation_keeps_the_legacy_call_signature(self, ray_local_mode, patch_low_level): + """Without the experimental flag the train fn is still called as (args, rollout_id, data_source).""" + args = _make_test_args(global_batch_size=4) + + executor = _make_executor(args) + executor.use_legacy_rollout_v1 = True + executor.data_source = SimpleNamespace() + executor.set_train_parallel_config({"dp_size": 1}) + calls: list[tuple] = [] + + def legacy_rollout_fn(passed_args, rollout_id, data_source, evaluation): + calls.append((passed_args, rollout_id, data_source, evaluation)) + return make_samples_grouped(n_groups=1, group_size=4) + + executor.generate_rollout = legacy_rollout_fn + + await executor.get(rollout_id=11) + + assert calls == [(args, 11, executor.data_source, False)] + + async def test_eval_keeps_the_legacy_call_signature_with_the_evaluation_flag( + self, ray_local_mode, patch_low_level + ): + """Without the experimental flag eval uses the same legacy protocol, flagged as evaluation.""" + args = _make_test_args() + + executor = _make_executor(args) + executor.use_legacy_rollout_v1 = True + executor.data_source = SimpleNamespace() + calls: list[tuple] = [] + + def legacy_eval_fn(passed_args, rollout_id, data_source, evaluation): + calls.append((passed_args, rollout_id, data_source, evaluation)) + return {"my_dataset": {"rewards": [1.0]}} + + executor.eval_generate_rollout = legacy_eval_fn + + await executor.eval(rollout_id=12) + + assert calls == [(args, 12, executor.data_source, True)] + + +class _RecordingMetricChecker: + def __init__(self) -> None: + self.evaluated: list[dict] = [] + self.disposed = False + + def on_eval(self, metrics: dict) -> None: + self.evaluated.append(metrics) + + def dispose(self) -> None: + self.disposed = True + + +class _RecordingCheckpointEvalFn(CheckpointEvalFn): + def __init__(self) -> None: + self.disposed = False + + async def evaluate_checkpoint(self, checkpoint_dir, input): + raise AssertionError("not exercised by the lifecycle tests") + + def dispose(self) -> None: + self.disposed = True + + +@pytest.mark.asyncio +class TestLifecycle: + async def test_eval_reports_the_logged_metrics_to_the_metric_checker( + self, ray_local_mode, patch_low_level, monkeypatch + ): + """The CI accuracy checker is fed exactly the eval metrics that were logged.""" + import miles.ray.rollout.rollout_executor as rexec + + monkeypatch.setattr( + rexec, "log_eval_rollout_data", lambda rollout_id, args, data, metrics: {"eval/accuracy": 0.75} + ) + + executor = _make_executor(_make_test_args()) + checker = _RecordingMetricChecker() + executor._metric_checker = checker + executor.eval_generate_rollout = lambda input: RolloutFnEvalOutput( + data={"my_dataset": {"rewards": [1.0]}}, metrics={} + ) + + await executor.eval(rollout_id=4) + + assert checker.evaluated == [{"eval/accuracy": 0.75}] + + async def test_dispose_releases_every_executor_owned_resource(self, ray_local_mode, patch_low_level, monkeypatch): + """Teardown closes the data source, runs event analysis and disposes the checker and the eval fn.""" + import miles.ray.rollout.rollout_executor as rexec + + analyzed: list = [] + monkeypatch.setattr(rexec, "event_analyzer", SimpleNamespace(run_analysis_from_args=analyzed.append)) + args = _make_test_args() + + executor = _make_executor(args) + closed: list = [] + executor.data_source = SimpleNamespace(close=lambda: closed.append("closed")) + checker = _RecordingMetricChecker() + executor._metric_checker = checker + eval_fn = _RecordingCheckpointEvalFn() + executor.eval_generate_rollout = eval_fn + + executor.dispose() + + assert closed == ["closed"] + assert analyzed == [args] + assert checker.disposed + assert eval_fn.disposed + + +@pytest.mark.asyncio +class TestNumRolloutPerEpoch: + async def test_counts_only_complete_global_batches(self, ray_local_mode, patch_low_level): + """A trailing partial batch is not a rollout, so the epoch length floors the division.""" + executor = _make_executor(_make_test_args(rollout_global_dataset=True, rollout_batch_size=8)) + executor.data_source = SimpleNamespace(dataset=list(range(20))) + + assert executor.get_num_rollout_per_epoch() == 2 + + async def test_rejects_a_non_global_data_source(self, ray_local_mode, patch_low_level): + """Without a global dataset there is no epoch length to report.""" + executor = _make_executor(_make_test_args(rollout_global_dataset=False)) + executor.data_source = SimpleNamespace(dataset=list(range(20))) + + with pytest.raises(AssertionError): + executor.get_num_rollout_per_epoch() + + @pytest.mark.asyncio class TestCheckpointWithoutARolloutFunction: async def test_checkpointing_a_replay_run_touches_only_the_data_source(self, ray_local_mode, patch_low_level): diff --git a/tests/fast/rollout/test_checkpoint_eval.py b/tests/fast/rollout/test_checkpoint_eval.py index a84201198c7..91853d8094b 100644 --- a/tests/fast/rollout/test_checkpoint_eval.py +++ b/tests/fast/rollout/test_checkpoint_eval.py @@ -203,6 +203,76 @@ def eval_generate_rollout(input): assert extra is None +class TestSnapshotEvalGuards: + async def test_snapshot_eval_without_an_hf_dir_is_rejected(self, controller_env): + """Snapshot eval has no checkpoint to evaluate without a dir, so it must fail loudly.""" + fn = CheckpointFnStub() + args = make_args(hf_checkpoint="/base", eval_keep_snapshots=2) + mgr = make_manager(args, eval_fn=fn) + + with pytest.raises(AssertionError, match="checkpoint eval requires an HF snapshot dir"): + await mgr.eval(5) + + assert fn.inputs == [] + + async def test_marker_bypass_evaluates_a_dir_without_a_complete_marker(self, controller_env, tmp_path): + """A caller-supplied checkpoint was never exported here, so there is no marker to wait for.""" + snapshot = tmp_path / "step_5" + snapshot.mkdir() + + fn = CheckpointFnStub() + args = make_args(hf_checkpoint="/base", eval_hf_dir=str(tmp_path), eval_keep_snapshots=2) + mgr = make_manager(args, eval_fn=fn) + + await mgr.eval(5, hf_dir=str(snapshot), require_marker=False) + + assert len(fn.inputs) == 1 + assert fn.inputs[0].hf_dir == str(snapshot) + assert "skip" not in controller_env.logged + + +class BlockingFleet: + def __init__(self): + self.pins = [] + self.release = asyncio.Event() + + async def pin(self, checkpoint_dir, weight_version): + self.pins.append(weight_version) + await self.release.wait() + return "fleet-state" + + +class TestEvalFleetSerialization: + async def test_set_eval_fleet_serializes_concurrent_checkpoint_pins(self, controller_env, monkeypatch, tmp_path): + """One fleet holds one pinned checkpoint, so a second eval point cannot pin until the first finishes.""" + for rollout_id in (5, 6): + snapshot = tmp_path / f"step_{rollout_id}" + snapshot.mkdir() + (snapshot / ".complete").touch() + + def eval_generate_rollout(input): + return RolloutFnEvalOutput(data={"ds": {"rewards": [1.0]}}) + + monkeypatch.setattr(rollout_executor_mod, "call_rollout_function", lambda fn, input: fn(input)) + args = make_args(hf_checkpoint="/base", eval_hf_dir=str(tmp_path)) + mgr = make_manager(args, eval_fn=eval_generate_rollout) + args.eval_uses_snapshots = True + fleet = BlockingFleet() + mgr.set_eval_fleet(fleet) + + first = asyncio.create_task(mgr.eval(5, hf_dir=str(tmp_path / "step_5"))) + second = asyncio.create_task(mgr.eval(6, hf_dir=str(tmp_path / "step_6"))) + for _ in range(5): + await asyncio.sleep(0) + + assert fleet.pins == ["5"] + + fleet.release.set() + await asyncio.gather(first, second) + + assert fleet.pins == ["5", "6"] + + # ---------------- driver (train_async.EvalDispatcher) ---------------- diff --git a/tests/fast/utils/fake_ray_ids.py b/tests/fast/utils/fake_ray_ids.py new file mode 100644 index 00000000000..31ced124bc3 --- /dev/null +++ b/tests/fast/utils/fake_ray_ids.py @@ -0,0 +1,8 @@ +from __future__ import annotations + +_NODE_ID_HEX_DIGITS = 56 + + +def fake_ray_node_id(index: int) -> str: + """Ray validates node ids as hex, so a readable "node-0" is rejected before scheduling.""" + return f"{index:0{_NODE_ID_HEX_DIGITS}x}" diff --git a/tests/fast/utils/test_http_utils.py b/tests/fast/utils/test_http_utils.py index 4cb3bc78681..fe0d0aa1df8 100644 --- a/tests/fast/utils/test_http_utils.py +++ b/tests/fast/utils/test_http_utils.py @@ -22,6 +22,7 @@ """ import asyncio +import inspect import multiprocessing import socket import subprocess @@ -29,11 +30,16 @@ import threading import time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from types import SimpleNamespace +from typing import Any, NamedTuple from unittest.mock import patch import httpx import pytest +import ray +from tests.fast.utils.fake_ray_ids import fake_ray_node_id +from miles.utils import http_utils from miles.utils.http_utils import GeneralHttpClientProvider, wait_for_server_ready @@ -377,3 +383,105 @@ def __init__(self) -> None: async def handle_async_request(self, request: httpx.Request) -> httpx.Response: self.deadlines = request.extensions["timeout"] return httpx.Response(200) + + +class TestDistributedPostActors: + def test_the_poster_actor_is_constructed_with_keyword_arguments(self, monkeypatch): + """A positional handoff silently binds to the wrong parameter once the actor grows another one.""" + recorded: list[tuple[tuple, dict]] = [] + + class _FakeActorClass: + def options(self, **_options): + return self + + def remote(self, *call_args, **call_kwargs): + recorded.append((call_args, call_kwargs)) + return object() + + monkeypatch.setattr(ray, "nodes", lambda: [{"NodeID": fake_ray_node_id(0), "Alive": True}]) + monkeypatch.setattr(ray, "remote", lambda _cls: _FakeActorClass()) + monkeypatch.setattr(http_utils, "_post_actors", []) + monkeypatch.setattr(http_utils, "_client_concurrency", 7) + + http_utils._init_ray_distributed_post(SimpleNamespace(num_gpus_per_node=2)) + + assert recorded == [((), {"concurrency": 8})] * 2 + + +class _PosterActorInit(NamedTuple): + actor_class: type + calls: list[tuple[tuple, dict]] + + +class _RecordingRemoteActorClass: + def __init__(self) -> None: + self.calls: list[tuple[tuple, dict]] = [] + + def options(self, **_options: Any) -> "_RecordingRemoteActorClass": + return self + + def remote(self, *call_args: Any, **call_kwargs: Any) -> object: + self.calls.append((call_args, call_kwargs)) + return object() + + +def _run_init_ray_distributed_post( + monkeypatch: pytest.MonkeyPatch, + *, + num_gpus_per_node: int = 1, + client_concurrency: int = 7, + nodes: list[dict] | None = None, +) -> _PosterActorInit: + captured: dict[str, type] = {} + remote_actor_class = _RecordingRemoteActorClass() + + def _fake_remote(cls: type) -> _RecordingRemoteActorClass: + captured["actor_class"] = cls + return remote_actor_class + + monkeypatch.setattr(ray, "nodes", lambda: nodes or [{"NodeID": fake_ray_node_id(0), "Alive": True}]) + monkeypatch.setattr(ray, "remote", _fake_remote) + monkeypatch.setattr(http_utils, "_post_actors", []) + monkeypatch.setattr(http_utils, "_client_concurrency", client_concurrency) + + http_utils._init_ray_distributed_post(SimpleNamespace(num_gpus_per_node=num_gpus_per_node)) + + return _PosterActorInit(actor_class=captured["actor_class"], calls=remote_actor_class.calls) + + +class TestPosterActorKeywordOnlyConstruction: + def test_the_poster_actor_refuses_a_positional_concurrency(self, monkeypatch): + """Constructing the poster actor positionally must fail so a later parameter cannot silently steal the slot.""" + actor_class = _run_init_ray_distributed_post(monkeypatch).actor_class + + with pytest.raises(TypeError): + actor_class(7) + + def test_every_poster_actor_constructor_parameter_is_keyword_only(self, monkeypatch): + """No poster actor constructor parameter may be positionally bindable.""" + actor_class = _run_init_ray_distributed_post(monkeypatch).actor_class + + parameters = list(inspect.signature(actor_class.__init__).parameters.values())[1:] + + assert [parameter.kind for parameter in parameters] == [inspect.Parameter.KEYWORD_ONLY] * len(parameters) + assert parameters + + def test_the_recorded_poster_keywords_bind_to_the_actor_constructor(self, monkeypatch): + """The keywords the call site sends must name real poster actor constructor parameters.""" + init = _run_init_ray_distributed_post(monkeypatch) + signature = inspect.signature(init.actor_class.__init__) + + for call_args, call_kwargs in init.calls: + assert call_args == () + signature.bind(object(), *call_args, **call_kwargs) + + def test_the_poster_actor_is_constructed_by_keyword_on_every_node_slot(self, monkeypatch): + """Every actor created across nodes and per-node slots receives its concurrency by keyword.""" + init = _run_init_ray_distributed_post( + monkeypatch, + num_gpus_per_node=2, + client_concurrency=10, + nodes=[{"NodeID": fake_ray_node_id(0), "Alive": True}, {"NodeID": fake_ray_node_id(1), "Alive": True}], + ) + + assert init.calls == [((), {"concurrency": 6})] * 4 diff --git a/tests/fast/utils/tracking_utils/__init__.py b/tests/fast/utils/tracking_utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/fast/utils/tracking_utils/test_prometheus_utils.py b/tests/fast/utils/tracking_utils/test_prometheus_utils.py new file mode 100644 index 00000000000..6f8c2ccb1e5 --- /dev/null +++ b/tests/fast/utils/tracking_utils/test_prometheus_utils.py @@ -0,0 +1,103 @@ +import inspect +from types import SimpleNamespace +from typing import Any, NamedTuple + +import pytest +from tests.fast.utils.fake_ray_ids import fake_ray_node_id + +from miles.utils.tracking_utils import prometheus_utils + + +class TestCollectorConstruction: + def test_the_collector_actor_is_constructed_with_keyword_arguments(self, monkeypatch): + """A positional handoff silently binds to the wrong parameter once the actor grows another one.""" + recorded: dict[str, object] = {} + args = SimpleNamespace(prometheus_run_name="run") + + class _FakeActorClass: + def options(self, **_options): + return self + + def remote(self, *call_args, **call_kwargs): + recorded["call_args"] = call_args + recorded["call_kwargs"] = call_kwargs + return SimpleNamespace(ping=SimpleNamespace(remote=lambda: None)) + + monkeypatch.setattr( + prometheus_utils, + "ray", + SimpleNamespace( + remote=lambda _cls: _FakeActorClass(), + get=lambda _ref: None, + get_runtime_context=lambda: SimpleNamespace(get_node_id=lambda: fake_ray_node_id(0)), + ), + ) + monkeypatch.setattr(prometheus_utils, "_collector_handle", None) + + prometheus_utils.init_prometheus(args, start_server=True) + + assert (recorded["call_args"], recorded["call_kwargs"]) == ((), {"args": args}) + + +class _CollectorInit(NamedTuple): + actor_class: type + call_args: tuple + call_kwargs: dict + + +def _run_init_prometheus(monkeypatch: pytest.MonkeyPatch, args: SimpleNamespace) -> _CollectorInit: + captured: dict[str, Any] = {} + + class _FakeActorClass: + def options(self, **_options: Any) -> "_FakeActorClass": + return self + + def remote(self, *call_args: Any, **call_kwargs: Any) -> SimpleNamespace: + captured["call_args"] = call_args + captured["call_kwargs"] = call_kwargs + return SimpleNamespace(ping=SimpleNamespace(remote=lambda: None)) + + def _fake_remote(cls: type) -> _FakeActorClass: + captured["actor_class"] = cls + return _FakeActorClass() + + monkeypatch.setattr( + prometheus_utils, + "ray", + SimpleNamespace( + remote=_fake_remote, + get=lambda _ref: None, + get_runtime_context=lambda: SimpleNamespace(get_node_id=lambda: fake_ray_node_id(0)), + ), + ) + monkeypatch.setattr(prometheus_utils, "_collector_handle", None) + + prometheus_utils.init_prometheus(args, start_server=True) + + return _CollectorInit( + actor_class=captured["actor_class"], + call_args=captured["call_args"], + call_kwargs=captured["call_kwargs"], + ) + + +class TestCollectorKeywordOnlyConstruction: + def test_the_collector_refuses_a_positional_args_object(self): + """Constructing the collector positionally must fail so a later parameter cannot silently steal the slot.""" + with pytest.raises(TypeError): + prometheus_utils._PrometheusCollector(SimpleNamespace(prometheus_port=0)) + + def test_every_collector_constructor_parameter_is_keyword_only(self): + """No collector constructor parameter may be positionally bindable.""" + parameters = list(inspect.signature(prometheus_utils._PrometheusCollector.__init__).parameters.values())[1:] + + assert [parameter.kind for parameter in parameters] == [inspect.Parameter.KEYWORD_ONLY] * len(parameters) + assert parameters + + def test_the_recorded_collector_keywords_bind_to_the_collector_constructor(self, monkeypatch): + """The keywords the call site sends must name real collector constructor parameters.""" + init = _run_init_prometheus(monkeypatch, SimpleNamespace(prometheus_run_name="run")) + + assert init.actor_class is prometheus_utils._PrometheusCollector + assert init.call_args == () + inspect.signature(init.actor_class.__init__).bind(object(), **init.call_kwargs)