Skip to content

Commit 124c261

Browse files
committed
[feat][SFT] Return final Fireworks export evidence
Expose the promoted model and trainer cleanup manifest to downstream lifecycle tracking. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> Signed-off-by: Bharat Mekala <155010017+bharatmekala@users.noreply.github.com>
1 parent ccf6fb3 commit 124c261

2 files changed

Lines changed: 5 additions & 3 deletions

File tree

skyrl/train/entrypoints/main_fireworks_sft.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,15 +14,15 @@
1414
from skyrl.train.fireworks_sft_trainer import FireworksSFTTrainer
1515

1616

17-
def run(cfg: SFTConfig) -> None:
17+
def run(cfg: SFTConfig) -> dict | None:
1818
"""Run the native SFT loop without attaching to Ray."""
1919

2020
validate_fireworks_sft_cfg(cfg)
2121
trainer = FireworksSFTTrainer(cfg, skyrl_cfg=build_skyrl_config_for_sft(cfg))
2222
try:
2323
trainer.setup()
2424
trainer.train()
25-
trainer.export_final_model()
25+
return trainer.export_final_model()
2626
except Exception as exc:
2727
if trainer.tracker is not None:
2828
trainer.tracker.log_exception(exc, step=trainer.global_step)

tests/train/test_fireworks_sft_trainer.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -321,14 +321,16 @@ def train(self):
321321

322322
def export_final_model(self):
323323
events.append("export")
324+
return {"model": "promoted"}
324325

325326
def shutdown(self):
326327
events.append("shutdown")
327328

328329
monkeypatch.setattr("skyrl.train.entrypoints.main_fireworks_sft.FireworksSFTTrainer", Trainer)
329330

330-
run(_cfg())
331+
result = run(_cfg())
331332

333+
assert result == {"model": "promoted"}
332334
assert events == ["init", "setup", "train", "export", "shutdown"]
333335

334336

0 commit comments

Comments
 (0)