diff --git a/rlix/pipeline/miles_model_update_service.py b/rlix/pipeline/miles_model_update_service.py
index 6371f38..9208dcc 100644
--- a/rlix/pipeline/miles_model_update_service.py
+++ b/rlix/pipeline/miles_model_update_service.py
@@ -300,6 +300,7 @@ async def _run_atomic_unit(
         # so the subsequent flush_cache returns 200 immediately.
         try:
             pause_refs = [h.pause_generation.remote(mode="retract") for h in handles.values()]
+            inflight_refs.extend(pause_refs)
             await _ray_get(pause_refs)
         except Exception as exc:  # noqa: BLE001
             logger.warning(
@@ -313,6 +314,7 @@ async def _run_atomic_unit(
             # (4.post) resume the engines so subsequent generate calls work.
             try:
                 cont_refs = [h.continue_generation.remote() for h in handles.values()]
+                inflight_refs.extend(cont_refs)
                 await _ray_get(cont_refs)
             except Exception as exc:  # noqa: BLE001
                 logger.warning(