diff --git a/rlix/pipeline/miles_pipeline.py b/rlix/pipeline/miles_pipeline.py index 18e7195..474fe52 100644 --- a/rlix/pipeline/miles_pipeline.py +++ b/rlix/pipeline/miles_pipeline.py @@ -674,17 +674,20 @@ def _after_training(self, step: int) -> None: # Coordinator drives sync_base_weights_to_active under its # resize lock; pipeline NEVER calls service / finalize / # set_weight_version directly (F05). - ray.get( - self._coordinator_handle.sync_base_weights_to_active.remote(int(step)) - ) - # Release the actor_train allocation back to the scheduler so - # other pipelines can step. R11-F1: only flip the ledger flag - # after a SUCCESSFUL release. - released = self._notify_release_cluster_gpus( - cluster_id=self._actor_train_cluster_id, global_step=int(step) - ) - if released: - self._actor_train_allocated = False + try: + ray.get( + self._coordinator_handle.sync_base_weights_to_active.remote(int(step)) + ) + finally: + # Release the actor_train allocation back to the scheduler so + # other pipelines can step even if weight sync fails after the + # train actor has offloaded. R11-F1: only flip the ledger flag + # after a SUCCESSFUL release. + released = self._notify_release_cluster_gpus( + cluster_id=self._actor_train_cluster_id, global_step=int(step) + ) + if released: + self._actor_train_allocated = False # ------------------------------------------------------------------ # M4 minimal hard cleanup diff --git a/tests/test_miles_pipeline_after_training_cleanup.py b/tests/test_miles_pipeline_after_training_cleanup.py new file mode 100644 index 0000000..c0b2839 --- /dev/null +++ b/tests/test_miles_pipeline_after_training_cleanup.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import ast +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +def _function_def(tree: ast.AST, name: str) -> ast.FunctionDef: + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == name: + return node + raise AssertionError(f"missing function {name}") + + +def _calls_attr(node: ast.AST, attr: str) -> bool: + return any(isinstance(child, ast.Attribute) and child.attr == attr for child in ast.walk(node)) + + +def test_after_training_releases_train_gpus_when_weight_sync_fails() -> None: + source = (REPO_ROOT / "rlix" / "pipeline" / "miles_pipeline.py").read_text( + encoding="utf-8" + ) + tree = ast.parse(source) + after_training = _function_def(tree, "_after_training") + + guarded_sync = [ + node + for node in ast.walk(after_training) + if isinstance(node, ast.Try) + and any(_calls_attr(stmt, "sync_base_weights_to_active") for stmt in node.body) + ] + + assert guarded_sync, "sync_base_weights_to_active must be guarded by try/finally" + assert any( + any(_calls_attr(stmt, "_notify_release_cluster_gpus") for stmt in node.finalbody) + for node in guarded_sync + ), "actor_train release must run from the sync finally block"