APR-157: add hosted SFT CLI config

packages/prime/src/prime_cli/api/rl.py

@@ -3,7 +3,7 @@
 from datetime import datetime
 from typing import Any, Dict, List, Literal, Optional
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import AliasChoices, BaseModel, ConfigDict, Field
 
 from prime_cli.core import APIClient, APIError, ValidationError
 
@@ -75,6 +75,12 @@ class RLRun(BaseModel):
     max_steps: int = Field(..., alias="maxSteps")
     max_tokens: Optional[int] = Field(None, alias="maxTokens")
     batch_size: int = Field(..., alias="batchSize")
+    loss: Optional[str] = "rl"
+    teacher: Optional[Dict[str, Any]] = Field(
+        None,
+        validation_alias=AliasChoices("teacher", "teacherConfig"),
+        serialization_alias="teacher",
+    )
     base_model: str = Field(..., alias="baseModel")
     environments: List[Dict[str, Any]] = Field(default_factory=list)
     run_config: Optional[Dict[str, Any]] = Field(None, alias="runConfig")
@@ -208,6 +214,8 @@ def create_run(
         enable_thinking: Optional[bool] = None,
         reasoning_effort: Optional[Literal["low", "medium", "high"]] = None,
         run_config: Optional[Dict[str, Any]] = None,
+        loss: str = "rl",
+        teacher: Optional[Dict[str, Any]] = None,
     ) -> RLRun:
         """Create a new Hosted Training run."""
         try:
@@ -225,6 +233,12 @@ def create_run(
                 "secrets": secrets_list,
             }
 
+            if loss != "rl":
+                payload["loss"] = loss
+
+            if teacher:
+                payload["teacher"] = teacher
+
             if name:
                 payload["name"] = name
 

packages/prime/src/prime_cli/commands/rl.py

@@ -215,6 +215,7 @@ def generate_rl_config_template(environment: str | None = None) -> str:
 
     return f'''\
 model = "Qwen/Qwen3.5-0.8B"
+loss = "rl" # "rl" | "sft"; OPD is not yet supported on hosted runtimes
 max_steps = 100
 
 # env_files = ["secrets.env"] # optional file(s) for secrets
@@ -228,6 +229,15 @@ def generate_rl_config_template(environment: str | None = None) -> str:
 # Optional: warm-start from an existing checkpoint
 # checkpoint_id = "..."
 
+# Optional: SFT distillation teacher
+# To use SFT, change the top-level loss to "sft" and uncomment this block.
+# [teacher]
+# model = "openai/gpt-oss-120b"
+#
+# [teacher.sampling]
+# max_tokens = 2048
+# reasoning_effort = "medium"
+
 [sampling]
 max_tokens = 2048
 # temperature = 0.7
@@ -385,6 +395,35 @@ def _reasoning_controls_mutually_exclusive(self) -> "SamplingConfig":
         return self
 
 
+class TeacherSamplingConfig(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    max_tokens: int | None = None
+    temperature: float | None = None
+    extra_body: Dict[str, Any] | None = None
+    enable_thinking: bool | None = None
+    reasoning_effort: Literal["low", "medium", "high"] | None = None
+
+    @model_validator(mode="after")
+    def _reasoning_controls_mutually_exclusive(self) -> "TeacherSamplingConfig":
+        if self.enable_thinking is not None and self.reasoning_effort is not None:
+            raise ValueError("enable_thinking and reasoning_effort cannot both be set")
+        return self
+
+
+class TeacherConfig(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    model: str
+    sampling: TeacherSamplingConfig | None = None
+
+    def to_api_dict(self) -> Dict[str, Any]:
+        result: Dict[str, Any] = {"model": {"name": self.model}}
+        if self.sampling is not None:
+            result["sampling"] = self.sampling.model_dump(exclude_none=True)
+        return result
+
+
 class EvalConfig(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
@@ -589,6 +628,8 @@ class RLConfig(BaseModel):
 
     name: str | None = None
     model: str
+    loss: Literal["rl", "sft", "opd"] = "rl"
+    teacher: TeacherConfig | None = None
     max_steps: int = 100
     batch_size: int = 128
     rollouts_per_example: int = 8
@@ -614,13 +655,23 @@ class RLConfig(BaseModel):
     env_files: List[str] = Field(default_factory=list)
 
     @model_validator(mode="after")
-    def validate_max_inflight_rollouts(self) -> "RLConfig":
+    def validate_config_consistency(self) -> "RLConfig":
         if self.max_inflight_rollouts is not None and self.oversampling_factor is not None:
             raise ValueError("Only one of max_inflight_rollouts and oversampling_factor can be set")
-        if self.max_inflight_rollouts is None:
-            return self
-        if self.max_inflight_rollouts < self.rollouts_per_example:
+        if (
+            self.max_inflight_rollouts is not None
+            and self.max_inflight_rollouts < self.rollouts_per_example
+        ):
             raise ValueError("max_inflight_rollouts must be at least rollouts_per_example")
+        if self.loss == "rl" and self.teacher is not None:
+            raise ValueError("teacher can only be set when loss is 'sft' or 'opd'")
+        if self.loss == "sft" and self.teacher is None:
+            raise ValueError("teacher is required when loss is 'sft'")
+        if self.loss == "opd":
+            raise ValueError(
+                "loss='opd' is not supported for hosted runs yet; OPD requires "
+                "teacher logprob scoring support in the hosted runtime"
+            )
         return self
 
 
@@ -893,6 +944,9 @@ def _fetch_pricing() -> None:
         # Model & Environment
         console.print("[cyan]Model & Environment[/cyan]")
         console.print(f"  Model:        {cfg.model}")
+        console.print(f"  Loss:         {cfg.loss}")
+        if cfg.teacher is not None:
+            console.print(f"  Teacher:      {cfg.teacher.model}")
         console.print(f"  Environments: {', '.join(e.id for e in cfg.env)}")
         if app_config.team_id:
             console.print(f"  Team:         {app_config.team_id}")
@@ -1112,6 +1166,8 @@ def _format(list_p: Any, eff_p: Any) -> str:
             enable_thinking=cfg.sampling.enable_thinking,
             reasoning_effort=cfg.sampling.reasoning_effort,
             run_config=cfg.run_config if cfg.run_config else None,
+            loss=cfg.loss,
+            teacher=cfg.teacher.to_api_dict() if cfg.teacher else None,
         )
 
         if output == "json":

packages/prime/tests/test_rl_api.py

@@ -81,3 +81,56 @@ def test_create_run_sends_max_inflight_rollouts() -> None:
     assert api_client.posts[0][0] == "/rft/runs"
     assert api_client.posts[0][1]["max_inflight_rollouts"] == 96
     assert run.max_inflight_rollouts == 96
+
+
+def test_create_run_sends_sft_loss_and_teacher_config() -> None:
+    api_client = FakeAPIClient()
+    client = RLClient(api_client)  # type: ignore[arg-type]
+
+    client.create_run(
+        model_name="openai/gpt-oss-20b",
+        environments=[{"id": "primeintellect/reverse-text"}],
+        max_tokens=512,
+        loss="sft",
+        teacher={
+            "model": {"name": "openai/gpt-oss-120b"},
+            "sampling": {
+                "max_tokens": 2048,
+                "reasoning_effort": "medium",
+            },
+        },
+    )
+
+    assert api_client.posts[0][0] == "/rft/runs"
+    payload = api_client.posts[0][1]
+    assert payload == {
+        "model": {"name": "openai/gpt-oss-20b"},
+        "environments": [{"id": "primeintellect/reverse-text"}],
+        "rollouts_per_example": 8,
+        "max_steps": 100,
+        "batch_size": 128,
+        "secrets": [],
+        "loss": "sft",
+        "teacher": {
+            "model": {"name": "openai/gpt-oss-120b"},
+            "sampling": {
+                "max_tokens": 2048,
+                "reasoning_effort": "medium",
+            },
+        },
+        "max_tokens": 512,
+    }
+
+
+def test_create_run_omits_default_rl_loss() -> None:
+    api_client = FakeAPIClient()
+    client = RLClient(api_client)  # type: ignore[arg-type]
+
+    client.create_run(
+        model_name="Qwen/Qwen3.5-0.8B",
+        environments=[{"id": "reverse-text"}],
+    )
+
+    assert api_client.posts[0][0] == "/rft/runs"
+    assert "loss" not in api_client.posts[0][1]
+    assert "teacher" not in api_client.posts[0][1]

packages/prime/tests/test_rl_config.py

@@ -71,6 +71,36 @@ def test_generate_rl_config_template_keeps_default_surface_minimal() -> None:
         assert field not in template
 
 
+def test_generate_rl_config_template_sft_example_loads(tmp_path: Path) -> None:
+    template = generate_rl_config_template()
+    template = template.replace(
+        'loss = "rl" # "rl" | "sft"; OPD is not yet supported on hosted runtimes',
+        'loss = "sft" # "rl" | "sft"; OPD is not yet supported on hosted runtimes',
+    )
+
+    lines: list[str] = []
+    in_teacher_example = False
+    for line in template.splitlines():
+        if line == "# Optional: SFT distillation teacher":
+            in_teacher_example = True
+            lines.append(line)
+            continue
+        if in_teacher_example and line == "":
+            in_teacher_example = False
+        if in_teacher_example and line.startswith("# ") and not line.startswith("# To use"):
+            line = line[2:]
+        lines.append(line)
+
+    config_path = tmp_path / "sft-template.toml"
+    config_path.write_text("\n".join(lines) + "\n")
+
+    cfg = load_config(str(config_path))
+
+    assert cfg.loss == "sft"
+    assert cfg.teacher is not None
+    assert cfg.teacher.model == "openai/gpt-oss-120b"
+
+
 def test_flatten_config_schema_expands_optional_nested_models() -> None:
     schema = RLConfig.model_json_schema()
     rows = _flatten_config_schema(schema, schema.get("$defs", {}))
@@ -143,6 +173,72 @@ def test_load_config_rejects_max_inflight_and_oversampling(tmp_path: Path) -> No
         load_config(str(config_path))
 
 
+def test_load_config_accepts_sft_teacher(tmp_path: Path) -> None:
+    config_path = tmp_path / "sft.toml"
+    config_path.write_text(
+        'model = "openai/gpt-oss-20b"\n'
+        'loss = "sft"\n'
+        "[teacher]\n"
+        'model = "openai/gpt-oss-120b"\n'
+        "[teacher.sampling]\n"
+        "max_tokens = 2048\n"
+        'reasoning_effort = "medium"\n'
+    )
+
+    cfg = load_config(str(config_path))
+
+    assert cfg.loss == "sft"
+    assert cfg.teacher is not None
+    assert cfg.teacher.model == "openai/gpt-oss-120b"
+    assert cfg.teacher.sampling is not None
+    assert cfg.teacher.sampling.max_tokens == 2048
+    assert cfg.teacher.to_api_dict() == {
+        "model": {"name": "openai/gpt-oss-120b"},
+        "sampling": {
+            "max_tokens": 2048,
+            "reasoning_effort": "medium",
+        },
+    }
+
+
+def test_load_config_rejects_teacher_temp_scheduler(tmp_path: Path) -> None:
+    config_path = tmp_path / "sft.toml"
+    config_path.write_text(
+        'model = "openai/gpt-oss-20b"\n'
+        'loss = "sft"\n'
+        "[teacher]\n"
+        'model = "openai/gpt-oss-120b"\n'
+        "[teacher.sampling.temp_scheduler]\n"
+        'type = "linear"\n'
+        "start_temperature = 1.0\n"
+        "end_temperature = 0.1\n"
+    )
+
+    with pytest.raises(typer.Exit):
+        load_config(str(config_path))
+
+
+def test_load_config_rejects_sft_without_teacher(tmp_path: Path) -> None:
+    config_path = tmp_path / "sft.toml"
+    config_path.write_text('model = "openai/gpt-oss-20b"\nloss = "sft"\n')
+
+    with pytest.raises(typer.Exit):
+        load_config(str(config_path))
+
+
+def test_load_config_rejects_opd_until_hosted_scoring_exists(tmp_path: Path) -> None:
+    config_path = tmp_path / "opd.toml"
+    config_path.write_text(
+        'model = "openai/gpt-oss-20b"\n'
+        'loss = "opd"\n'
+        "[teacher]\n"
+        'model = "openai/gpt-oss-120b"\n'
+    )
+
+    with pytest.raises(typer.Exit):
+        load_config(str(config_path))
+
+
 def test_load_config_rejects_max_inflight_below_rollouts_per_example(tmp_path: Path) -> None:
     config_path = tmp_path / "rl.toml"
     config_path.write_text('model = "dummy"\nrollouts_per_example = 8\nmax_inflight_rollouts = 4\n')