[feat][SFT] Preserve promotable periodic checkpoints

Let hosted SFT runs retain inference-formatted checkpoints at the same cadence as resumable DCP state, so selected intermediate steps can be promoted after training without reconnecting to the trainer.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>

Author: bharat-mekala

skyrl/backends/fireworks/runtime.py

@@ -29,6 +29,15 @@ class SamplerVersion:
     snapshot_path: str
 
 
+@dataclass(frozen=True)
+class PromotableCheckpoint:
+    """Inference checkpoint persisted for later model promotion."""
+
+    snapshot_path: str
+    checkpoint_resource: str | None
+    checkpoint_type: str | None
+
+
 @dataclass(frozen=True)
 class FireworksInferenceEndpoint:
     """Native OpenAI-compatible endpoint for the managed deployment."""
@@ -72,6 +81,7 @@ def __init__(
         self.tokenizer = tokenizer
         self.config = config
         self._state_lock = threading.Condition()
+        self._provider_operation_lock = threading.RLock()
         self._publish_lock = asyncio.Lock()
         self._sampler: Any | None = None
         self._sampler_identity: SamplerVersion | None = None
@@ -238,6 +248,91 @@ def _save() -> str:
                 self._next_version = version + 1
             return identity
 
+    def save_promotable_checkpoint(
+        self,
+        checkpoint_name: str,
+        *,
+        appear_timeout_s: float = 90.0,
+        poll_s: float = 3.0,
+    ) -> PromotableCheckpoint:
+        with self._provider_operation_lock:
+            with self._state_lock:
+                if self._closed:
+                    raise RuntimeError("Fireworks runtime is closed")
+
+            result = self.training_client.save_weights_for_sampler(
+                checkpoint_name,
+                checkpoint_type="base",
+            ).result(timeout=self.config.request_timeout_s)
+            snapshot_path = str(getattr(result, "path", "") or "")
+            if not snapshot_path:
+                raise RuntimeError(f"Fireworks save_weights_for_sampler({checkpoint_name!r}) returned no path")
+
+            snapshot_id = snapshot_path.rstrip("/").rsplit("/", 1)[-1]
+            deadline = time.monotonic() + appear_timeout_s
+            matches: list[dict] = []
+            while time.monotonic() < deadline:
+                try:
+                    rows = self.service.list_checkpoints(self.trainer_job_id)
+                except Exception:
+                    time.sleep(poll_s)
+                    continue
+                matches = [row for row in rows if row.get("promotable") and _checkpoint_short_name(row) == snapshot_id]
+                if matches:
+                    break
+                time.sleep(poll_s)
+            if len(matches) > 1:
+                raise RuntimeError(
+                    f"Expected at most one promotable Fireworks checkpoint for {checkpoint_name!r}, "
+                    f"got {len(matches)}"
+                )
+            if not matches:
+                warnings.warn(
+                    f"Promotable Fireworks checkpoint {checkpoint_name!r} was saved but did not "
+                    "surface on the control plane before the visibility timeout",
+                    RuntimeWarning,
+                )
+                return PromotableCheckpoint(
+                    snapshot_path=snapshot_path,
+                    checkpoint_resource=None,
+                    checkpoint_type=None,
+                )
+            checkpoint = matches[0]
+            return PromotableCheckpoint(
+                snapshot_path=snapshot_path,
+                checkpoint_resource=checkpoint["name"],
+                checkpoint_type=checkpoint.get("checkpointType"),
+            )
+
+    async def promote_checkpoint_resource(
+        self,
+        *,
+        checkpoint: PromotableCheckpoint,
+        output_model_id: str,
+    ) -> dict[str, Any]:
+        if not checkpoint.checkpoint_resource:
+            raise ValueError("Promotable checkpoint has no control-plane resource")
+
+        def _promote() -> dict:
+            with self._provider_operation_lock:
+                with self._state_lock:
+                    if self._closed:
+                        raise RuntimeError("Fireworks runtime is closed")
+                return self.service.promote_checkpoint(
+                    name=checkpoint.checkpoint_resource,
+                    output_model_id=output_model_id,
+                    base_model=self.config.base_model,
+                )
+
+        model = await asyncio.to_thread(_promote)
+        return {
+            "sampler_path": checkpoint.snapshot_path,
+            "checkpoint_resource": checkpoint.checkpoint_resource,
+            "checkpoint_type": checkpoint.checkpoint_type,
+            "output_model_id": output_model_id,
+            "model": model,
+        }
+
     async def promote_final_model(
         self,
         *,
@@ -397,4 +492,9 @@ def _wait_for_samples() -> tuple[Any | None, int]:
                 )
             if sampler is not None:
                 await asyncio.to_thread(_close_quietly, sampler)
-            await asyncio.to_thread(_close_quietly, self.service)
+
+            def _close_service() -> None:
+                with self._provider_operation_lock:
+                    _close_quietly(self.service)
+
+            await asyncio.to_thread(_close_service)

skyrl/backends/fireworks/training_backend.py

@@ -174,14 +174,25 @@ def save_checkpoint(self, model: str, ckpt_dir: str, tokenizer=None) -> None:
         if not provider_path:
             raise RuntimeError(f"Fireworks save_state({checkpoint_name!r}) returned no checkpoint path")
 
+        promotable = None
+        if self.fireworks_config.save_promotable_checkpoints:
+            promotable = self.runtime.save_promotable_checkpoint(checkpoint_name)
+
         manifest = {
             "format_version": 1,
             "checkpoint_kind": "fireworks_dcp",
             "checkpoint_name": checkpoint_name,
             "provider_path": provider_path,
             "source_trainer_job_id": self.runtime.trainer_job_id,
+            "base_model": self.fireworks_config.base_model,
             "includes_optimizer_state": True,
         }
+        if promotable is not None:
+            manifest["promotable_checkpoint"] = {
+                "snapshot_path": promotable.snapshot_path,
+                "checkpoint_resource": promotable.checkpoint_resource,
+                "checkpoint_type": promotable.checkpoint_type,
+            }
         io.makedirs(ckpt_dir, exist_ok=True)
         manifest_path = os.path.join(ckpt_dir, self._CHECKPOINT_MANIFEST)
         with io.open_file(manifest_path, "w") as f:

skyrl/train/config/config.py

@@ -942,6 +942,8 @@ class FireworksConfig(BaseConfig):
     cleanup_on_exit: bool = True
     output_model_id: Optional[str] = None
     """Optional Fireworks model ID promoted from the final sampler checkpoint."""
+    save_promotable_checkpoints: bool = False
+    """Save a promotable inference snapshot alongside each resumable checkpoint."""
     delete_trainer_after_promotion: bool = True
     cleanup_deployment_on_close: str = "delete"
     """``"delete"`` or ``"scale_to_zero"`` for the SDK-created deployment."""

skyrl/train/config/sft_config.py

@@ -572,6 +572,8 @@ def validate_fireworks_sft_cfg(cfg: SFTConfig) -> None:
         raise ValueError("Fireworks SFT requires ckpt_path when ckpt_interval > 0")
     if cfg.resume_from == "latest" and not cfg.ckpt_path:
         raise ValueError("Fireworks SFT requires ckpt_path when resume_from='latest'")
+    if cfg.fireworks.save_promotable_checkpoints and not cfg.ckpt_path:
+        raise ValueError("Fireworks SFT requires ckpt_path when save_promotable_checkpoints=true")
     if cfg.fireworks.output_model_id:
         if not cfg.ckpt_path:
             raise ValueError("Fireworks SFT requires ckpt_path when output_model_id is set")

skyrl/train/fireworks_sft_trainer.py

@@ -6,7 +6,7 @@
 import json
 import os
 
-from skyrl.backends.fireworks.runtime import FireworksRuntime
+from skyrl.backends.fireworks.runtime import FireworksRuntime, PromotableCheckpoint
 from skyrl.backends.fireworks.sft import FireworksSFTDispatch
 from skyrl.backends.skyrl_train.utils.io import io
 from skyrl.train.sft_trainer import SFTTrainer
@@ -80,12 +80,32 @@ def export_final_model(self) -> dict | None:
         if not checkpoint_name:
             raise ValueError(f"Fireworks DCP manifest has no checkpoint_name: {dcp_manifest_path}")
 
-        result = asyncio.run(
-            self._fireworks_runtime.promote_final_model(
-                checkpoint_name=checkpoint_name,
-                output_model_id=output_model_id,
+        promotable = dcp_manifest.get("promotable_checkpoint")
+        snapshot_path = promotable.get("snapshot_path") if isinstance(promotable, dict) else None
+        checkpoint_resource = promotable.get("checkpoint_resource") if isinstance(promotable, dict) else None
+        if (
+            isinstance(snapshot_path, str)
+            and snapshot_path
+            and isinstance(checkpoint_resource, str)
+            and "/checkpoints/" in checkpoint_resource
+        ):
+            result = asyncio.run(
+                self._fireworks_runtime.promote_checkpoint_resource(
+                    checkpoint=PromotableCheckpoint(
+                        snapshot_path=snapshot_path,
+                        checkpoint_resource=checkpoint_resource,
+                        checkpoint_type=promotable.get("checkpoint_type"),
+                    ),
+                    output_model_id=output_model_id,
+                )
+            )
+        else:
+            result = asyncio.run(
+                self._fireworks_runtime.promote_final_model(
+                    checkpoint_name=checkpoint_name,
+                    output_model_id=output_model_id,
+                )
             )
-        )
         manifest = {
             "format_version": 1,
             "global_step": self.global_step,

tests/backends/fireworks/test_runtime.py

@@ -4,7 +4,7 @@
 
 import pytest
 
-from skyrl.backends.fireworks.runtime import FireworksRuntime
+from skyrl.backends.fireworks.runtime import FireworksRuntime, PromotableCheckpoint
 from skyrl.train.config import FireworksConfig
 
 
@@ -225,6 +225,71 @@ async def test_publish_reuses_one_stable_sampler() -> None:
     await runtime.close()
 
 
+def test_save_promotable_checkpoint_uses_base_sampler_and_control_plane_row() -> None:
+    service = _Service()
+    service.checkpoint_rows = [
+        {
+            "name": "accounts/test/rlorTrainerJobs/trainer/checkpoints/step-2-cafebabe",
+            "promotable": True,
+            "checkpointType": "CHECKPOINT_TYPE_INFERENCE_BASE",
+        }
+    ]
+    training = _TrainingClient()
+    runtime = _runtime(
+        service=service,
+        training_client=training,
+        config=FireworksConfig(request_timeout_s=123),
+    )
+
+    result = runtime.save_promotable_checkpoint("step-2", poll_s=0)
+
+    assert training.names == ["step-2"]
+    assert training.checkpoint_types == ["base"]
+    assert result.snapshot_path == "snapshot://step-2-cafebabe"
+    assert result.checkpoint_resource.endswith("step-2-cafebabe")
+    assert result.checkpoint_type == "CHECKPOINT_TYPE_INFERENCE_BASE"
+
+
+def test_save_promotable_checkpoint_keeps_path_when_control_plane_lags() -> None:
+    runtime = _runtime()
+
+    with pytest.warns(RuntimeWarning, match="visibility timeout"):
+        result = runtime.save_promotable_checkpoint("step-2", appear_timeout_s=0, poll_s=0)
+
+    assert result.snapshot_path == "snapshot://step-2-cafebabe"
+    assert result.checkpoint_resource is None
+    assert result.checkpoint_type is None
+
+
+@pytest.mark.asyncio
+async def test_promote_checkpoint_resource_uses_existing_control_plane_row() -> None:
+    service = _Service()
+    runtime = _runtime(
+        service=service,
+        config=FireworksConfig(base_model="accounts/fireworks/models/qwen3-4b"),
+    )
+    checkpoint = PromotableCheckpoint(
+        snapshot_path="snapshot://step-2-cafebabe",
+        checkpoint_resource="accounts/test/rlorTrainerJobs/trainer/checkpoints/step-2-cafebabe",
+        checkpoint_type="CHECKPOINT_TYPE_INFERENCE_BASE",
+    )
+
+    result = await runtime.promote_checkpoint_resource(
+        checkpoint=checkpoint,
+        output_model_id="sft-step-2",
+    )
+
+    assert service.promotions == [
+        {
+            "name": checkpoint.checkpoint_resource,
+            "output_model_id": "sft-step-2",
+            "base_model": "accounts/fireworks/models/qwen3-4b",
+        }
+    ]
+    assert result["checkpoint_resource"] == checkpoint.checkpoint_resource
+    assert result["model"]["name"] == "accounts/test/models/sft-step-2"
+
+
 @pytest.mark.asyncio
 async def test_promote_final_model_saves_base_sampler_and_uses_control_plane_row() -> None:
     service = _Service()

tests/backends/fireworks/test_training_backend.py

@@ -138,6 +138,7 @@ def test_policy_dispatch_saves_and_cross_job_loads_dcp_checkpoint(tmp_path) -> N
     manifest = json.loads((ckpt_dir / "fireworks_checkpoint.json").read_text())
     assert manifest["source_trainer_job_id"] == "source-trainer"
     assert manifest["includes_optimizer_state"] is True
+    assert "promotable_checkpoint" not in manifest
 
     runtime.trainer_job_id = "target-trainer"
     dispatch.load_checkpoint(
@@ -155,6 +156,40 @@ def test_policy_dispatch_saves_and_cross_job_loads_dcp_checkpoint(tmp_path) -> N
     ]
 
 
+def test_policy_dispatch_saves_promotable_checkpoint_when_enabled(tmp_path) -> None:
+    training_client = _TrainingClient()
+    saved_promotable = []
+
+    def save_promotable(name):
+        saved_promotable.append(name)
+        return SimpleNamespace(
+            snapshot_path=f"snapshot://{name}-cafebabe",
+            checkpoint_resource=f"accounts/test/rlorTrainerJobs/trainer/checkpoints/{name}-cafebabe",
+            checkpoint_type="CHECKPOINT_TYPE_INFERENCE_BASE",
+        )
+
+    runtime = SimpleNamespace(
+        training_client=training_client,
+        trainer_job_id="source-trainer",
+        save_promotable_checkpoint=save_promotable,
+    )
+    cfg = _cfg()
+    cfg.trainer.fireworks.save_promotable_checkpoints = True
+    dispatch = FireworksPolicyDispatch(runtime, cfg.trainer.fireworks, cfg.trainer.policy.optimizer_config)
+    ckpt_dir = tmp_path / "global_step_7" / "policy"
+
+    dispatch.save_checkpoint("policy", str(ckpt_dir), tokenizer="unused")
+
+    manifest = json.loads((ckpt_dir / "fireworks_checkpoint.json").read_text())
+    checkpoint_name = manifest["checkpoint_name"]
+    assert saved_promotable == [checkpoint_name]
+    assert manifest["promotable_checkpoint"] == {
+        "snapshot_path": f"snapshot://{checkpoint_name}-cafebabe",
+        "checkpoint_resource": (f"accounts/test/rlorTrainerJobs/trainer/checkpoints/{checkpoint_name}-cafebabe"),
+        "checkpoint_type": "CHECKPOINT_TYPE_INFERENCE_BASE",
+    }
+
+
 def test_policy_dispatch_loads_same_trainer_checkpoint_by_name(tmp_path) -> None:
     training_client = _TrainingClient()
     runtime = SimpleNamespace(training_client=training_client, trainer_job_id="same-trainer")