feat(fireworks): retry + reconnect training-client RPCs on transient faults (#708)

A transient blip on the shared training client around the per-step weight
hot-load (eager_parameter_sync_step=1) crashed entire multi-hour runs: the
Fireworks SDK's forward_backward / optim_step / save_and_hotload have no retry
("failures propagate so the training loop can crash"), and a step-N
forward_backward TimeoutError at step ~17 tore down the run.

Wrap the per-step training-client RPCs (forward, forward_backward incl. aux,
optim_step, save_and_hotload) in a retry helper that, on a transient error
(timeout / connection / channel reset / 5xx markers), optionally reconnects the
client to the SAME job and retries with linear backoff. Non-transient errors
(bad data) and cancellation propagate immediately.

- Reconnect refreshes the ReconnectableClient via the SDK's own _connect path
  (rlor_mgr.wait_for_existing(job_id) -> client._use_endpoint), re-attaching to
  the same job so server-side optimizer/gradient state is preserved. Used for
  forward/forward_backward/optim_step. save_and_hotload retries without
  reconnect (the WeightSyncer holds the raw client; re-dispatch is idempotent).
- Config (fireworks_infra.common): step_max_retries (default 2, 0 disables),
  step_retry_backoff_s (default 10). Each retry logs at WARNING.

CAVEAT (documented in _run_training_op): forward_backward / optim_step mutate
server-side state. A retry assumes the failed RPC didn't commit it — true for
channel/dispatch failures; a *false* timeout could double-apply for one step.
Retries are kept low; this bounded, rare perturbation beats crashing the run.
For genuine slowness, raise step_timeout instead.

Adds tests/trainer/test_fireworks_policy_trainer_retry.py (transient
classification, retry-succeeds, retry-exhausted, non-transient-not-retried,
zero-retries-disables, reconnect on/off, reconnect-without-mgr no-op).

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: jeffreysijuntan

rllm/trainer/config/rllm/backend/fireworks.yaml

@@ -99,6 +99,11 @@ fireworks_infra:
     learning_rate: 1e-5      # placeholder for fireworks managed flow. Should be overridden in the training codes.
     max_seq_len: null        # null = use the training shape's default context length
     step_timeout: 3600       # SDK default (3600s)
+    # Transient-fault tolerance for training-client RPCs (fwd/bwd, optim, weight
+    # hot-load). A blip on the shared training client around the per-step
+    # hot-load otherwise crashes the whole run, since the SDK does not retry.
+    step_max_retries: 2      # retries on transient errors (timeout/connection); 0 disables
+    step_retry_backoff_s: 10 # base linear backoff between retries (seconds)
     weight_sync_timeout: 600
 
   deployments:

rllm/trainer/fireworks/fireworks_policy_trainer.py

@@ -21,6 +21,7 @@
 
 import tinker
 from fireworks.training.sdk import WeightSyncer
+from omegaconf import OmegaConf
 from tinker.types import AdamParams
 from training.utils.client import ReconnectableClient
 
@@ -109,6 +110,9 @@ def __init__(
         self.weight_syncer = weight_syncer
         self._rlor_mgr = rlor_mgr
         self._policy_job_id = policy_job_id
+        # Transient-fault tolerance for training-client RPCs (see _run_training_op).
+        self._step_max_retries = max(0, int(OmegaConf.select(config, "fireworks_infra.common.step_max_retries", default=2)))
+        self._step_retry_backoff_s = float(OmegaConf.select(config, "fireworks_infra.common.step_retry_backoff_s", default=10.0))
         self._resume_checkpoint_name = self.config.training.get("resume_from_dcp_checkpoint")
         self._resume_source_job_id = self.config.training.get("resume_from_fireworks_job_id") or policy_job_id
 
@@ -117,6 +121,102 @@ def __init__(
         self.algorithm_config = algorithm_config or AlgorithmConfig.from_config(self.config.rllm.algorithm)
         self.resolve_builtin_loss(self.algorithm_config)
 
+    # ------------------------------------------------------------------
+    # Transient-fault tolerance for training-client RPCs
+    # ------------------------------------------------------------------
+
+    # Substrings marking a retryable transient failure on the shared training
+    # client — channel/deadline blips (typically around the per-step weight
+    # hot-load), not data errors. Mirrors the rollout engine's transient markers.
+    _TRANSIENT_MARKERS = (
+        "timeout",
+        "timed out",
+        "deadline",
+        "unavailable",
+        "connection",
+        "connection reset",
+        "broken pipe",
+        "eof",
+        "502",
+        "503",
+        "504",
+        "goaway",
+        "rst_stream",
+    )
+
+    def _is_transient(self, exc: BaseException) -> bool:
+        """Whether *exc* looks like a transient channel/timeout fault (vs a data error)."""
+        if isinstance(exc, TimeoutError | ConnectionError):
+            return True
+        text = f"{type(exc).__name__}: {exc}".lower()
+        return any(marker in text for marker in self._TRANSIENT_MARKERS)
+
+    def _reconnect_training_client(self) -> bool:
+        """Best-effort refresh of the training client's channel to the **same** job.
+
+        Re-resolves the trainer endpoint via the job manager and rebinds the
+        ``ReconnectableClient`` to it (the SDK's own ``_connect`` path, which we
+        replicate here because ``from_training_client`` instances carry no job
+        manager). Server-side optimizer/gradient state lives with the job, which
+        is unchanged — this only replaces a stale channel. Returns True on success.
+
+        Note: this refreshes ``self.training_client`` (used by forward / fwd-bwd /
+        optim_step) only; the ``WeightSyncer`` holds the raw client, so its calls
+        are retried without reconnect.
+        """
+        mgr, job_id = self._rlor_mgr, self._policy_job_id
+        if mgr is None or not job_id:
+            logger.warning("Cannot reconnect training client: rlor_mgr/policy_job_id unset")
+            return False
+        try:
+            endpoint = mgr.wait_for_existing(job_id)
+            self.training_client._use_endpoint(endpoint)
+            logger.info("Reconnected training client to job %s", job_id)
+            return True
+        except Exception as exc:
+            logger.warning("Training client reconnect failed for job %s: %s", job_id, exc)
+            return False
+
+    async def _run_training_op(self, fn, *args, op_name: str, reconnect: bool = False, **kwargs):
+        """Run a blocking training-client RPC in a thread, retrying transient faults.
+
+        On a transient error (timeout / connection / channel reset — typically a
+        blip on the shared training client around the per-step weight hot-load)
+        we optionally reconnect to the same job and retry, up to
+        ``step_max_retries`` times with linear backoff. Non-transient errors
+        (e.g. malformed data) and cancellation propagate immediately.
+
+        CAVEAT: forward_backward / optim_step mutate server-side state (gradient
+        accumulation / optimizer). A retry assumes the failed RPC did NOT commit
+        that mutation — true for channel/dispatch failures, where the server
+        never ran it. A *false* timeout (server finished but the ack was lost)
+        could double-apply for that one step; retries are kept low and this
+        bounded, rare perturbation is preferred over crashing a multi-hour run.
+        For genuine slowness, raise ``step_timeout`` instead of relying on retry.
+        """
+        attempt = 0
+        while True:
+            try:
+                return await asyncio.to_thread(fn, *args, **kwargs)
+            except Exception as exc:  # noqa: BLE001 — re-raised below unless transient
+                if attempt >= self._step_max_retries or not self._is_transient(exc):
+                    raise
+                attempt += 1
+                backoff = self._step_retry_backoff_s * attempt
+                logger.warning(
+                    "Training op %r failed (%s: %s); %sretry %d/%d after %.0fs",
+                    op_name,
+                    type(exc).__name__,
+                    exc,
+                    "reconnect+" if reconnect else "",
+                    attempt,
+                    self._step_max_retries,
+                    backoff,
+                )
+                if reconnect:
+                    self._reconnect_training_client()
+                await asyncio.sleep(backoff)
+
     # ------------------------------------------------------------------
     # Initialization
     # ------------------------------------------------------------------
@@ -225,10 +325,13 @@ async def _sync_weights(self, name: str, checkpoint_type: str | None = None) ->
         Returns the snapshot_name on success, None on failure."""
         if self.weight_syncer is None:
             return None
-        snapshot_name = await asyncio.to_thread(
+        # reconnect=False: the syncer holds the raw client, not self.training_client;
+        # re-dispatching save_and_hotload (idempotent) is the available recovery.
+        snapshot_name = await self._run_training_op(
             self.weight_syncer.save_and_hotload,
             name,
             checkpoint_type=checkpoint_type,
+            op_name="save_and_hotload",
         )
         logger.debug("Weights synced to deployment: %s", name)
         return snapshot_name
@@ -307,10 +410,12 @@ async def _compute_proximal_logprobs(
 
         Only called when ``bypass_mode=False`` (3-policy / decoupled PPO).
         """
-        prox_fwd = await asyncio.to_thread(
+        prox_fwd = await self._run_training_op(
             self.training_client.forward,
             datums,
             "cross_entropy",
+            op_name="forward",
+            reconnect=True,
         )
         return [out["logprobs"].data for out in prox_fwd.loss_fn_outputs]
 
@@ -536,11 +641,13 @@ async def forward_backward_from_trajectory_groups(
         )
 
         kernel_loss, kernel_config = self._builtin_loss
-        fwd_bwd_result = await asyncio.to_thread(
+        fwd_bwd_result = await self._run_training_op(
             self.training_client.forward_backward,
             builtin_datums,
             kernel_loss,
             loss_fn_config=kernel_config,
+            op_name="forward_backward",
+            reconnect=True,
         )
 
         # Merge remote fwd/bwd metrics (e.g. loss) into adv_metrics
@@ -552,10 +659,12 @@ async def forward_backward_from_trajectory_groups(
         # Auxiliary-loss passes: each accumulates gradients on top of the policy
         # gradient before the (externally invoked) optim_step.
         for aux, datums in aux_passes:
-            aux_fwd_bwd = await asyncio.to_thread(
+            aux_fwd_bwd = await self._run_training_op(
                 self.training_client.forward_backward,
                 datums,
                 "cross_entropy",
+                op_name="forward_backward(aux)",
+                reconnect=True,
             )
             if hasattr(aux_fwd_bwd, "metrics") and aux_fwd_bwd.metrics:
                 for k, v in aux_fwd_bwd.metrics.items():
@@ -613,10 +722,12 @@ async def optim_step(
         # gradient, so we disable server-side grad-accumulation normalization.
         if build_aux_losses(self.algorithm_config):
             grad_norm = GradAccNormalization.NONE
-        optim_result = await asyncio.to_thread(
+        optim_result = await self._run_training_op(
             self.training_client.optim_step,
             adam_params,
             grad_accumulation_normalization=grad_norm,
+            op_name="optim_step",
+            reconnect=True,
         )
 
         metrics = {}

tests/trainer/test_fireworks_policy_trainer_retry.py

@@ -0,0 +1,143 @@
+"""Transient-fault tolerance for Fireworks training-client RPCs.
+
+Covers ``FireworksPolicyTrainer._run_training_op`` / ``_is_transient`` /
+``_reconnect_training_client`` in isolation (the trainer is built via
+``__new__`` so the heavy SDK init is skipped). Run via ``asyncio.run`` — no
+pytest-asyncio needed.
+"""
+
+import asyncio
+
+import pytest
+
+from rllm.trainer.fireworks.fireworks_policy_trainer import FireworksPolicyTrainer
+
+
+class _FakeClient:
+    def __init__(self):
+        self.endpoints = []
+
+    def _use_endpoint(self, ep):
+        self.endpoints.append(ep)
+
+
+class _FakeMgr:
+    def wait_for_existing(self, job_id):
+        return f"ep-{job_id}"
+
+
+def _trainer(*, max_retries=2, backoff=0.0, rlor_mgr=None, job_id="job-1", client=None):
+    """A trainer with only the attributes the fault-tolerance helpers touch."""
+    t = FireworksPolicyTrainer.__new__(FireworksPolicyTrainer)
+    t._step_max_retries = max_retries
+    t._step_retry_backoff_s = backoff
+    t._rlor_mgr = rlor_mgr
+    t._policy_job_id = job_id
+    t.training_client = client
+    return t
+
+
+def test_is_transient_classification():
+    t = _trainer()
+    assert t._is_transient(TimeoutError("timed out"))
+    assert t._is_transient(ConnectionError("reset"))
+    assert t._is_transient(RuntimeError("StatusCode.UNAVAILABLE"))
+    assert t._is_transient(RuntimeError("deadline exceeded"))
+    assert t._is_transient(RuntimeError("upstream returned 503"))
+    # Data / programming errors are NOT transient — must propagate.
+    assert not t._is_transient(ValueError("malformed datum"))
+    assert not t._is_transient(KeyError("missing field"))
+
+
+def test_retry_succeeds_after_transient_blips():
+    calls = []
+
+    def fn():
+        calls.append(1)
+        if len(calls) < 3:
+            raise TimeoutError("blip")
+        return "ok"
+
+    t = _trainer(max_retries=3)
+    assert asyncio.run(t._run_training_op(fn, op_name="x")) == "ok"
+    assert len(calls) == 3  # 2 failures + 1 success
+
+
+def test_retry_exhausted_reraises():
+    def fn():
+        raise TimeoutError("persistent")
+
+    t = _trainer(max_retries=2)
+    with pytest.raises(TimeoutError):
+        asyncio.run(t._run_training_op(fn, op_name="x"))
+
+
+def test_non_transient_not_retried():
+    calls = []
+
+    def fn():
+        calls.append(1)
+        raise ValueError("bad data")
+
+    t = _trainer(max_retries=5)
+    with pytest.raises(ValueError):
+        asyncio.run(t._run_training_op(fn, op_name="x"))
+    assert len(calls) == 1  # failed once, no retry
+
+
+def test_zero_retries_disables():
+    calls = []
+
+    def fn():
+        calls.append(1)
+        raise TimeoutError("blip")
+
+    t = _trainer(max_retries=0)
+    with pytest.raises(TimeoutError):
+        asyncio.run(t._run_training_op(fn, op_name="x"))
+    assert len(calls) == 1
+
+
+def test_reconnect_invoked_between_retries_when_enabled():
+    client = _FakeClient()
+    t = _trainer(max_retries=2, rlor_mgr=_FakeMgr(), job_id="j", client=client)
+    calls = []
+
+    def fn():
+        calls.append(1)
+        if len(calls) < 2:
+            raise TimeoutError("blip")
+        return "ok"
+
+    assert asyncio.run(t._run_training_op(fn, op_name="x", reconnect=True)) == "ok"
+    assert client.endpoints == ["ep-j"]  # same-job channel refresh, once
+
+
+def test_reconnect_skipped_when_disabled():
+    client = _FakeClient()
+    t = _trainer(max_retries=2, rlor_mgr=_FakeMgr(), client=client)
+    calls = []
+
+    def fn():
+        calls.append(1)
+        if len(calls) < 2:
+            raise TimeoutError("blip")
+        return "ok"
+
+    asyncio.run(t._run_training_op(fn, op_name="x", reconnect=False))
+    assert client.endpoints == []  # never reconnected
+
+
+def test_reconnect_without_mgr_is_safe_noop():
+    # _reconnect returns False (no rlor_mgr) but the retry loop still proceeds.
+    t = _trainer(max_retries=1, rlor_mgr=None)
+    calls = []
+
+    def fn():
+        calls.append(1)
+        if len(calls) < 2:
+            raise TimeoutError("blip")
+        return "ok"
+
+    assert asyncio.run(t._run_training_op(fn, op_name="x", reconnect=True)) == "ok"
+    assert t._reconnect_training_client() is False