diff --git a/rlix/pipeline/miles_model_update_service.py b/rlix/pipeline/miles_model_update_service.py index 6371f38..e3e8093 100644 --- a/rlix/pipeline/miles_model_update_service.py +++ b/rlix/pipeline/miles_model_update_service.py @@ -190,13 +190,17 @@ async def sync_selected_workers( ) cpu_serialize_set = target - broadcast_set - # R06-F1 fix: track every Ray ObjectRef issued by this atomic - # unit so an asyncio.wait_for timeout (or any other cancellation) - # can fan out ray.cancel(force=True) on them. Without this, the - # local coroutine cancels but the Ray actor methods on - # cache_owner / rollout_manager keep running and continue to - # hold cache_owner._cache_lock — the next sync_selected_workers - # call would queue behind the still-executing prior method. + # R06-F1: track every Ray ObjectRef issued by this atomic unit so + # an asyncio.wait_for timeout (or any other cancellation) can fan + # out ray.cancel on them. Scope of what cancel can actually do + # (Ray actor-task semantics): QUEUED tasks are cancelled and never + # start; a task already EXECUTING on the sync cache_owner / + # rollout_manager actors cannot be interrupted (Ray only sets a + # cooperative is_canceled() flag that miles does not poll), so a + # running run_sync_session keeps holding cache_owner._cache_lock + # until it finishes and the next sync queues behind it. Preventing + # the queued follow-up RPCs (finalize / set_weight_version / + # try_put) from firing after a timeout is the effective win here. inflight_refs: list = [] async def _run() -> int: @@ -235,7 +239,14 @@ def _cancel_inflight(self, inflight_refs: list, *, reason: str) -> None: ) for ref in inflight_refs: try: - ray.cancel(ref, force=True) + # force=False is the ONLY mode Ray allows for actor tasks — + # force=True raises ValueError before cancelling anything + # (docs: "Only force=False is allowed for an Actor Task"), + # which made this whole fan-out a silent no-op. force=False + # cancels queued tasks; executing tasks on sync actors run + # to completion (see the R06-F1 note in + # sync_selected_workers). + ray.cancel(ref, force=False) except Exception as exc: # noqa: BLE001 logger.warning( "ray.cancel failed pipeline_id=%s reason=%s exc=%r", diff --git a/tests/test_miles_service_cancel_inflight.py b/tests/test_miles_service_cancel_inflight.py new file mode 100644 index 0000000..176543f --- /dev/null +++ b/tests/test_miles_service_cancel_inflight.py @@ -0,0 +1,90 @@ +"""Regression test for _cancel_inflight's Ray actor-task cancel semantics. + +Every ref in ``inflight_refs`` is an actor-method ObjectRef, and Ray only +allows ``force=False`` for actor tasks — ``force=True`` raises ValueError +before cancelling anything, which made the R06-F1 fan-out a silent no-op +(each ValueError was swallowed by the per-ref try/except). Pin that every +cancel goes out with ``force=False`` and actually reaches ray.cancel. +""" + +from __future__ import annotations + +import importlib +import sys +import types +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[1] +RLIX_ROOT = REPO_ROOT / "rlix" + + +def _install_import_stubs(monkeypatch: pytest.MonkeyPatch) -> tuple[list, types.ModuleType]: + for module_name in list(sys.modules): + if module_name == "ray" or module_name.startswith("rlix"): + monkeypatch.delitem(sys.modules, module_name, raising=False) + + ray_stub = types.ModuleType("ray") + cancel_calls: list = [] + + def _remote(*args, **kwargs): + # Support both bare ``@ray.remote`` (class passed directly) and + # ``@ray.remote(...)`` / ``ray.remote(cls)`` forms. + if len(args) == 1 and callable(args[0]) and not kwargs: + return args[0] + + def _decorate(obj): + return obj + + return _decorate + + def _cancel(ref, *, force=False, recursive=True): + # Mirror Ray's actor-task contract: force=True raises ValueError + # before any cancellation happens. + if force: + raise ValueError("force=True is not allowed for an Actor Task") + cancel_calls.append((ref, force)) + + ray_stub.remote = _remote + ray_stub.cancel = _cancel + ray_stub.get = lambda ref, timeout=None: ref + ray_stub.get_actor = lambda *args, **kwargs: None + monkeypatch.setitem(sys.modules, "ray", ray_stub) + + package_roots = { + "rlix": RLIX_ROOT, + "rlix.pipeline": RLIX_ROOT / "pipeline", + "rlix.protocol": RLIX_ROOT / "protocol", + "rlix.utils": RLIX_ROOT / "utils", + } + for module_name, module_path in package_roots.items(): + package_module = types.ModuleType(module_name) + package_module.__path__ = [str(module_path)] # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, module_name, package_module) + + return cancel_calls, ray_stub + + +def test_cancel_inflight_uses_force_false_and_cancels_every_ref( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cancel_calls, _ = _install_import_stubs(monkeypatch) + module = importlib.import_module("rlix.pipeline.miles_model_update_service") + + service = module.MilesModelUpdateService( + pipeline_id="miles_test", + cache_owner_actor=object(), + rollout_manager=object(), + ) + + refs = ["ref_a", "ref_b", "ref_c"] + inflight = list(refs) + service._cancel_inflight(inflight, reason="unit-test") + + # Every ref reached ray.cancel with force=False (force=True would have + # raised inside the stub and never recorded the call — the pre-fix + # silent-no-op behavior). + assert cancel_calls == [(ref, False) for ref in refs] + # The list is drained so a second timeout cannot double-cancel. + assert inflight == []