From c41314929dc99e1ad237c0c47b1b29d2c2ccd2f6 Mon Sep 17 00:00:00 2001 From: Alexander Paskov Date: Fri, 21 Aug 2026 18:03:46 +0300 Subject: [PATCH 1/2] =?UTF-8?q?feat(sdk):=20RL=20training=20surface=20?= =?UTF-8?q?=E2=80=94=20typed=20Rust=20core=20+=20python=20client.rl=20name?= =?UTF-8?q?space?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three layers, replacing the earlier pure-python draft after review feedback (the SDK's architecture is Rust-core + pyo3; a parallel urllib transport was the wrong shape): 1. CORE (basilica-sdk): rl.rs DTOs mirroring basilica-api's /rl/* route DTOs field-for-field — the compiler, not a runtime test, now catches wire drift. serde-flatten catch-alls on the create requests preserve unknown fields verbatim (forward-compat / the body= escape hatch). Client methods beside the transport helpers: create/get cluster + job, submit_rl_manifest. 2. BINDING (basilica-sdk-python): five rl_* pymethods with a JSON-string boundary — serde-validates against the core's typed DTOs, sends through the core transport (inheriting the FULL auth chain incl. the CLI-login token fallback), returns JSON. Deliberate boundary choice: nine nested pyclass mirrors would be pure boilerplate (the in-tree precedent). 3. PYTHON (basilica/rl.py): thin RlNamespace — ergonomic kwargs builders, wait_cluster/wait_job poll loops (wait_job RETURNS a Failed document rather than raising), client-side judge-requires-custom-reward guard. BasilicaClient gains a lazy .rl property. Bonus correctness: the core parses the API's real error envelope ({"error":{code,message,...}}), so RL errors surface the server's actionable message verbatim — the earlier urllib draft guessed the wrong envelope and would have shown raw JSON blobs. Tests: 10 contract tests against the COMPILED extension (maturin develop) with a real in-process HTTP server — exact wire bodies, ref renames, escape-hatch field survival, poll-to-terminal both outcomes, client-side serde rejection before any HTTP, error-envelope mapping to ValueError. Plus 3 core wire-shape tests (92 core tests green total); clippy clean both crates. --- .../python/basilica/__init__.py | 15 + .../basilica-sdk-python/python/basilica/rl.py | 213 +++++++++++ crates/basilica-sdk-python/src/lib.rs | 78 ++++ .../tests/test_rl_client.py | 208 +++++++++++ crates/basilica-sdk/src/client.rs | 38 ++ crates/basilica-sdk/src/lib.rs | 2 + crates/basilica-sdk/src/rl.rs | 338 ++++++++++++++++++ 7 files changed, 892 insertions(+) create mode 100644 crates/basilica-sdk-python/python/basilica/rl.py create mode 100644 crates/basilica-sdk-python/tests/test_rl_client.py create mode 100644 crates/basilica-sdk/src/rl.rs diff --git a/crates/basilica-sdk-python/python/basilica/__init__.py b/crates/basilica-sdk-python/python/basilica/__init__.py index 02bff4c33..e0be2fb63 100644 --- a/crates/basilica-sdk-python/python/basilica/__init__.py +++ b/crates/basilica-sdk-python/python/basilica/__init__.py @@ -455,6 +455,8 @@ def _build_inference_health_check(port: int) -> HealthCheckConfig: __all__ = [ # Main client "BasilicaClient", + # RL training namespace (client.rl; module: basilica.rl) + "RlNamespace", # Decorator API "deployment", "DeployedFunction", @@ -600,12 +602,25 @@ def __init__(self, base_url: Optional[str] = None, api_key: Optional[str] = None self._base_url = base_url self._client = _BasilicaClient(base_url, api_key) + self._rl = None @property def base_url(self) -> str: """The API endpoint URL.""" return self._base_url + @property + def rl(self) -> "RlNamespace": + """RL training namespace (GRPO post-training): clusters, jobs, + manifests. Thin wrapper over the compiled core's rl_* methods — + inherits the full auth chain incl. the CLI-login fallback. See + :mod:`basilica.rl`.""" + if self._rl is None: + from basilica.rl import RlNamespace + + self._rl = RlNamespace(self._client) + return self._rl + def _build_deploy_request( self, name: str, diff --git a/crates/basilica-sdk-python/python/basilica/rl.py b/crates/basilica-sdk-python/python/basilica/rl.py new file mode 100644 index 000000000..e96d7f475 --- /dev/null +++ b/crates/basilica-sdk-python/python/basilica/rl.py @@ -0,0 +1,213 @@ +"""RL training namespace: GRPO post-training on the Basilica RL Training API. + + >>> from basilica import BasilicaClient + >>> client = BasilicaClient() # BASILICA_API_TOKEN / CLI login + >>> client.rl.create_cluster( + ... name="my-pool", + ... base_model="Qwen/Qwen2.5-7B-Instruct", + ... gpu_model="H100", + ... ) + >>> client.rl.wait_cluster("my-pool") + >>> job = client.rl.create_job( + ... cluster="my-pool", max_steps=50, + ... reward_name="my-reward", reward_source=REWARD_PY, + ... dataset_name="my-data", dataset_repo="openai/gsm8k", + ... dataset_config="main", dataset_split="train", + ... prompt_column="question", answer_column="answer", + ... ) + >>> final = client.rl.wait_job(job["name"]) # {phase, step, metrics, artifactURI} + +THIN WRAPPER over the compiled core (#1509 review round): this module builds +the ergonomic kwargs into wire dicts and hands them to the Rust binding's +``rl_*`` methods, which serde-validate against the core's typed DTOs +(``basilica_sdk::rl`` — the compile-time-shared contract with the server) +and send through the core transport. That inherits the full auth chain +(explicit key, BASILICA_API_TOKEN, and the CLI-login token fallback the +earlier pure-python transport could not reach) and the core's error +mapping: non-2xx surfaces as ValueError (bad request), PermissionError +(authz), ConnectionError (transport), FileNotFoundError (not found), or +RuntimeError (server error), each carrying the server's message verbatim. + +The ``body=`` escape hatch on every create call sends a raw dict; unknown +fields survive the typed round-trip verbatim (serde-flatten catch-alls in +the core DTOs), so server-side schema additions never strand you on an SDK +release. +""" + +from __future__ import annotations + +import json +import time +from typing import Any, Optional + +_TERMINAL_JOB_PHASES = frozenset({"Succeeded", "Failed", "TimedOut"}) + + +def _drop_none(d: dict) -> dict: + return {k: v for k, v in d.items() if v is not None} + + +class RlNamespace: + """The ``client.rl`` surface. Constructed by BasilicaClient; hold no + credentials here — the compiled core owns auth and transport.""" + + def __init__(self, core: Any): + self._core = core + + # -- clusters ---------------------------------------------------------- + + def create_cluster( + self, + *, + base_model: str, + gpu_model: str, + trainer_gpus: int = 4, + rollout_gpus: int = 4, + name: Optional[str] = None, + min_memory_gb: Optional[int] = None, + idle_ttl: Optional[str] = None, + body: Optional[dict] = None, + ) -> dict: + """POST /rl/clusters — a warm trainer+rollout GPU pool. + + Certified shapes: 4+4 (H100 for <16B models, H200-class for >=16B — + admission rejects bad pairings with an actionable message). + ``body`` overrides everything (raw wire dict, escape hatch). + """ + if body is None: + + def fleet(count: int) -> dict: + return { + "replicas": 1, + "gpu": _drop_none( + {"model": gpu_model, "count": count, "minMemoryGb": min_memory_gb} + ), + } + + body = _drop_none( + { + "name": name, + "baseModel": base_model, + "trainer": fleet(trainer_gpus), + "rollout": fleet(rollout_gpus), + "idleTtl": idle_ttl, + } + ) + return json.loads(self._core.rl_create_cluster(json.dumps(body))) + + def get_cluster(self, name: str) -> dict: + return json.loads(self._core.rl_get_cluster(name)) + + def wait_cluster( + self, name: str, timeout_s: float = 1800.0, poll_s: float = 15.0 + ) -> dict: + """Poll until phase == Ready (raises TimeoutError otherwise).""" + deadline = time.monotonic() + timeout_s + while True: + cluster = self.get_cluster(name) + if cluster.get("phase") == "Ready": + return cluster + if time.monotonic() >= deadline: + raise TimeoutError( + f"cluster {name!r} not Ready after {timeout_s}s " + f"(last phase: {cluster.get('phase')!r})" + ) + time.sleep(poll_s) + + # -- jobs -------------------------------------------------------------- + + def create_job( + self, + *, + cluster: str, + max_steps: int, + name: Optional[str] = None, + algorithm: str = "grpo", + # custom reward (user: + inline source); omit for the builtin + reward_name: Optional[str] = None, + reward_source: Optional[str] = None, + judge: bool = False, + judge_model: Optional[str] = None, + # custom dataset (public HF repo + column mapping); omit for builtin + dataset_name: Optional[str] = None, + dataset_repo: Optional[str] = None, + dataset_split: Optional[str] = None, + dataset_config: Optional[str] = None, + prompt_column: Optional[str] = None, + answer_column: Optional[str] = None, + lr: Optional[str] = None, + body: Optional[dict] = None, + ) -> dict: + """POST /rl/jobs — a GRPO training job on a Ready cluster. + + The reward is any deterministic stdlib-Python + ``reward(prompt, completion, **ctx) -> float``; it runs in an + isolated credential-free pod. ``judge=True`` exposes + ``ctx["judge"](prompt)`` backed by an in-cluster judge model + (requires a custom reward). ``body`` overrides everything. + """ + if body is None: + reward = None + if reward_name is not None: + if reward_source is None: + raise ValueError("reward_source is required with reward_name") + reward = {"ref": f"user:{reward_name}", "source": reward_source} + if judge or judge_model: + reward["judge"] = _drop_none({"model": judge_model}) or {} + elif judge or judge_model: + raise ValueError( + "judge requires a custom reward (it is called from your reward code)" + ) + dataset = None + if dataset_name is not None: + dataset = { + "ref": f"user:{dataset_name}", + "hf": _drop_none( + { + "repo": dataset_repo, + "config": dataset_config, + "split": dataset_split, + "promptColumn": prompt_column, + "answerColumn": answer_column, + } + ), + } + body = _drop_none( + { + "clusterRef": cluster, + "name": name, + "algorithm": algorithm, + "maxSteps": max_steps, + "reward": reward, + "dataset": dataset, + "lr": lr, + } + ) + return json.loads(self._core.rl_create_job(json.dumps(body))) + + def get_job(self, name: str) -> dict: + return json.loads(self._core.rl_get_job(name)) + + def wait_job( + self, name: str, timeout_s: float = 6 * 3600.0, poll_s: float = 30.0 + ) -> dict: + """Poll until the job is terminal (Succeeded/Failed/TimedOut) and + return the final document either way — check ``phase`` yourself; + raising on Failed would hide the failure detail behind an + exception.""" + deadline = time.monotonic() + timeout_s + while True: + job = self.get_job(name) + if job.get("phase") in _TERMINAL_JOB_PHASES: + return job + if time.monotonic() >= deadline: + raise TimeoutError( + f"job {name!r} not terminal after {timeout_s}s " + f"(last phase: {job.get('phase')!r})" + ) + time.sleep(poll_s) + + # -- manifest (declarative: one document -> cluster and/or job) -------- + + def submit_manifest(self, manifest: dict) -> dict: + return json.loads(self._core.rl_submit_manifest(json.dumps(manifest))) diff --git a/crates/basilica-sdk-python/src/lib.rs b/crates/basilica-sdk-python/src/lib.rs index 737a997cd..49a00e5c2 100644 --- a/crates/basilica-sdk-python/src/lib.rs +++ b/crates/basilica-sdk-python/src/lib.rs @@ -99,6 +99,84 @@ impl BasilicaClient { Ok(response.into()) } + // ===== RL Training API (GRPO post-training) ===== + // + // JSON-string boundary by design (the line-317 precedent: nine nested + // pyclass mirrors would be pure boilerplate): the Python side builds the + // wire dict, this layer serde-validates it against the CORE's typed DTOs + // (basilica_sdk::rl — the compile-time-shared contract with the server) + // and sends through the core client, inheriting its auth chain including + // the CLI-login token fallback. Responses return as JSON for the thin + // Python wrapper to expose as dicts. + + /// Create a warm RL cluster. `request_json` must match + /// basilica_sdk::rl::CreateRlClusterRequest. + fn rl_create_cluster(&self, py: Python, request_json: String) -> PyResult { + let request: basilica_sdk::rl::CreateRlClusterRequest = serde_json::from_str(&request_json) + .map_err(|e| PyValueError::new_err(format!("invalid RL cluster request: {e}")))?; + let client = Arc::clone(&self.inner); + let response = py + .detach(|| { + self.runtime + .block_on(async move { client.create_rl_cluster(request).await }) + }) + .map_err(|e| self.map_error_to_python(e))?; + serde_json::to_string(&response).map_err(|e| PyRuntimeError::new_err(e.to_string())) + } + + /// Get an RL cluster's status (phase, modelLoaded, activeJobName). + fn rl_get_cluster(&self, py: Python, name: String) -> PyResult { + let client = Arc::clone(&self.inner); + let response = py + .detach(|| { + self.runtime + .block_on(async move { client.get_rl_cluster(&name).await }) + }) + .map_err(|e| self.map_error_to_python(e))?; + serde_json::to_string(&response).map_err(|e| PyRuntimeError::new_err(e.to_string())) + } + + /// Submit a GRPO training job. `request_json` must match + /// basilica_sdk::rl::CreateRlJobRequest. + fn rl_create_job(&self, py: Python, request_json: String) -> PyResult { + let request: basilica_sdk::rl::CreateRlJobRequest = serde_json::from_str(&request_json) + .map_err(|e| PyValueError::new_err(format!("invalid RL job request: {e}")))?; + let client = Arc::clone(&self.inner); + let response = py + .detach(|| { + self.runtime + .block_on(async move { client.create_rl_job(request).await }) + }) + .map_err(|e| self.map_error_to_python(e))?; + serde_json::to_string(&response).map_err(|e| PyRuntimeError::new_err(e.to_string())) + } + + /// Get an RL job's status (phase, step, metrics, artifactURI). + fn rl_get_job(&self, py: Python, name: String) -> PyResult { + let client = Arc::clone(&self.inner); + let response = py + .detach(|| { + self.runtime + .block_on(async move { client.get_rl_job(&name).await }) + }) + .map_err(|e| self.map_error_to_python(e))?; + serde_json::to_string(&response).map_err(|e| PyRuntimeError::new_err(e.to_string())) + } + + /// Submit a declarative RL manifest (one document -> cluster and/or job). + fn rl_submit_manifest(&self, py: Python, manifest_json: String) -> PyResult { + let manifest: serde_json::Value = serde_json::from_str(&manifest_json) + .map_err(|e| PyValueError::new_err(format!("invalid RL manifest: {e}")))?; + let client = Arc::clone(&self.inner); + let response = py + .detach(|| { + self.runtime + .block_on(async move { client.submit_rl_manifest(manifest).await }) + }) + .map_err(|e| self.map_error_to_python(e))?; + serde_json::to_string(&response).map_err(|e| PyRuntimeError::new_err(e.to_string())) + } + /// List available nodes /// /// Args: diff --git a/crates/basilica-sdk-python/tests/test_rl_client.py b/crates/basilica-sdk-python/tests/test_rl_client.py new file mode 100644 index 000000000..f168c5d56 --- /dev/null +++ b/crates/basilica-sdk-python/tests/test_rl_client.py @@ -0,0 +1,208 @@ +"""basilica.rl contract tests, run against the COMPILED core transport: a +real stdlib HTTP server receives what the Rust client actually sends, +asserting the exact wire shapes the RL API's deny-unknown-fields DTOs +enforce — key casing, the nested `ref` renames, auth header, escape-hatch +passthrough, poll loops, and the core's error mapping. + +Requires the built extension (maturin develop / the installed wheel); +skipped cleanly where only the pure-python tree is on the path. +""" + +import json +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import pytest + +basilica = pytest.importorskip("basilica") +pytest.importorskip("basilica._basilica") + + +class _Recorder(BaseHTTPRequestHandler): + requests: list = [] + responses: list = [] # (status, body-dict) popped per request + + def _handle(self): + length = int(self.headers.get("Content-Length") or 0) + body = json.loads(self.rfile.read(length)) if length else None + _Recorder.requests.append( + { + "method": self.command, + "path": self.path, + "auth": self.headers.get("Authorization"), + "body": body, + } + ) + status, resp = ( + _Recorder.responses.pop(0) if _Recorder.responses else (200, {}) + ) + payload = json.dumps(resp).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + do_GET = do_POST = _handle + + def log_message(self, *_): + pass + + +def _api_error(message, code="BASILICA_API_BAD_REQUEST"): + """The basilica-api error envelope (error.rs into_response).""" + return {"error": {"code": code, "message": message, + "timestamp": "2026-08-24T00:00:00Z", "retryable": False}} + + +@pytest.fixture() +def server(): + _Recorder.requests = [] + _Recorder.responses = [] + httpd = ThreadingHTTPServer(("127.0.0.1", 0), _Recorder) + t = threading.Thread(target=httpd.serve_forever, daemon=True) + t.start() + yield f"http://127.0.0.1:{httpd.server_address[1]}", _Recorder + httpd.shutdown() + httpd.server_close() + + +def rl(base): + return basilica.BasilicaClient(base_url=base, api_key="test-key").rl + + +def test_create_cluster_wire_shape(server): + base, rec = server + rec.responses = [(200, {"name": "my-pool", "uid": "u1", "phase": "Provisioning"})] + rl(base).create_cluster( + name="my-pool", + base_model="Qwen/Qwen2.5-7B-Instruct", + gpu_model="H100", + trainer_gpus=4, + rollout_gpus=4, + idle_ttl="30m", + ) + (r,) = rec.requests + assert (r["method"], r["path"]) == ("POST", "/rl/clusters") + assert r["auth"] == "Bearer test-key" + assert r["body"] == { + "name": "my-pool", + "baseModel": "Qwen/Qwen2.5-7B-Instruct", + "trainer": {"replicas": 1, "gpu": {"model": "H100", "count": 4}}, + "rollout": {"replicas": 1, "gpu": {"model": "H100", "count": 4}}, + "idleTtl": "30m", + } + + +def test_create_job_full_wire_shape(server): + base, rec = server + rec.responses = [(200, {"name": "j1", "uid": "u2", "phase": "Pending"})] + rl(base).create_job( + cluster="my-pool", + max_steps=50, + reward_name="my-reward", + reward_source="def reward(prompt, completion, **ctx):\n return 1.0\n", + judge=True, + dataset_name="my-data", + dataset_repo="openai/gsm8k", + dataset_config="main", + dataset_split="train", + prompt_column="question", + answer_column="answer", + ) + (r,) = rec.requests + assert (r["method"], r["path"]) == ("POST", "/rl/jobs") + body = r["body"] + assert body["clusterRef"] == "my-pool" + assert body["maxSteps"] == 50 + assert body["algorithm"] == "grpo" + # the serde renames: nested identity fields are `ref` + assert body["reward"]["ref"] == "user:my-reward" + assert body["reward"]["judge"] == {} + assert body["dataset"]["ref"] == "user:my-data" + assert body["dataset"]["hf"] == { + "repo": "openai/gsm8k", + "config": "main", + "split": "train", + "promptColumn": "question", + "answerColumn": "answer", + } + # None-valued optionals must be ABSENT (deny_unknown_fields tolerates + # absence; null in a non-Option server slot would 400) + assert "name" not in body and "lr" not in body + + +def test_judge_without_custom_reward_is_a_client_error(server): + base, _ = server + with pytest.raises(ValueError, match="judge requires a custom reward"): + rl(base).create_job(cluster="c", max_steps=3, judge=True) + + +def test_builtin_job_minimal_body(server): + base, rec = server + rec.responses = [(200, {"name": "j1", "uid": "u2", "phase": "Pending"})] + rl(base).create_job(cluster="my-pool", max_steps=3) + (r,) = rec.requests + assert r["body"] == {"clusterRef": "my-pool", "algorithm": "grpo", "maxSteps": 3} + + +def test_raw_body_escape_hatch_preserves_unknown_fields(server): + # serde-flatten catch-alls in the core DTOs: a field this SDK version + # doesn't know must SURVIVE the typed round-trip verbatim. + base, rec = server + raw = {"clusterRef": "x", "algorithm": "grpo", "maxSteps": 1, "futureField": True} + rec.responses = [(200, {"name": "j1", "uid": "u2", "phase": "Pending"})] + rl(base).create_job(cluster="ignored", max_steps=99, body=raw) + (r,) = rec.requests + assert r["body"]["futureField"] is True + assert r["body"]["maxSteps"] == 1 + assert r["body"]["clusterRef"] == "x" + + +def test_wait_job_polls_to_terminal(server): + base, rec = server + rec.responses = [ + (200, {"phase": "Running"}), + (200, {"phase": "Running"}), + (200, {"phase": "Succeeded", "artifactURI": "s3://x/uid"}), + ] + final = rl(base).wait_job("j1", timeout_s=30, poll_s=0.01) + assert final["phase"] == "Succeeded" + assert final["artifactURI"] == "s3://x/uid" + assert len(rec.requests) == 3 + assert all(r["path"] == "/rl/jobs/j1" for r in rec.requests) + + +def test_wait_job_returns_failed_rather_than_raising(server): + base, rec = server + rec.responses = [(200, {"phase": "Failed"})] + final = rl(base).wait_job("j2", timeout_s=5, poll_s=0.01) + assert final["phase"] == "Failed" + + +def test_api_error_surfaces_server_message(server): + # 400 + the basilica-api envelope -> the core maps BadRequest -> + # PyValueError carrying the server's message verbatim. + base, rec = server + rec.responses = [(400, _api_error("trainer fleet totals 7 GPUs; the GRPO train batch ..."))] + with pytest.raises(ValueError, match="totals 7 GPUs"): + rl(base).create_cluster( + base_model="Qwen/Qwen2.5-7B-Instruct", gpu_model="H100", trainer_gpus=7 + ) + + +def test_invalid_request_json_rejected_client_side(server): + # the binding serde-validates BEFORE any HTTP: a body that cannot parse + # into the typed DTO raises without touching the network. + base, rec = server + with pytest.raises(ValueError, match="invalid RL job request"): + rl(base).create_job(cluster="c", max_steps=1, body={"maxSteps": "not-a-number"}) + assert rec.requests == [] + + +def test_manifest_posts_verbatim(server): + base, rec = server + doc = {"cluster": {"baseModel": "m"}, "job": {"maxSteps": 3}} + rl(base).submit_manifest(doc) + (r,) = rec.requests + assert (r["method"], r["path"], r["body"]) == ("POST", "/rl/manifest", doc) diff --git a/crates/basilica-sdk/src/client.rs b/crates/basilica-sdk/src/client.rs index f7fd9bc57..3ddc702a7 100644 --- a/crates/basilica-sdk/src/client.rs +++ b/crates/basilica-sdk/src/client.rs @@ -42,6 +42,10 @@ use crate::{ CreateJobRequest, CreateJobResponse, DeleteJobResponse, JobLogsResponse, JobStatusResponse, ReadFileRequest, ReadFileResponse, ResumeJobResponse, SuspendJobResponse, }, + rl::{ + CreateRlClusterRequest, CreateRlClusterResponse, CreateRlJobRequest, CreateRlJobResponse, + RlClusterStatusResponse, RlJobStatusResponse, RlManifestRequest, RlManifestResponse, + }, types::{ ApiKeyInfo, ApiKeyResponse, ApiListRentalsResponse, BalanceResponse, CardPurchaseResponse, CardPurchaseSummary, CreateApiKeyRequest, CreateCardPurchaseRequest, @@ -327,6 +331,40 @@ impl BasilicaClient { self.post(&path, &serde_json::json!({})).await } + // ----- RL Training API (GRPO post-training) -------------------------- + + /// Create a warm RL cluster. Poll [`get_rl_cluster`](Self::get_rl_cluster) + /// until its phase is `Ready`. + pub async fn create_rl_cluster( + &self, + request: CreateRlClusterRequest, + ) -> Result { + self.post("/rl/clusters", &request).await + } + + /// Get a cluster's status. + pub async fn get_rl_cluster(&self, name: &str) -> Result { + self.get(&format!("/rl/clusters/{}", name)).await + } + + /// Submit a GRPO training job to a Ready cluster. + pub async fn create_rl_job(&self, request: CreateRlJobRequest) -> Result { + self.post("/rl/jobs", &request).await + } + + /// Get a job's status (phase, step, metrics, `artifactURI`). + pub async fn get_rl_job(&self, name: &str) -> Result { + self.get(&format!("/rl/jobs/{}", name)).await + } + + /// Submit a declarative manifest (renders a cluster and/or a job). + pub async fn submit_rl_manifest( + &self, + manifest: RlManifestRequest, + ) -> Result { + self.post("/rl/manifest", &manifest).await + } + /// Resume a suspended job pub async fn resume_job(&self, job_id: &str) -> Result { let path = format!("/v2/jobs/{}/resume", job_id); diff --git a/crates/basilica-sdk/src/lib.rs b/crates/basilica-sdk/src/lib.rs index 035069a7d..7b764984f 100644 --- a/crates/basilica-sdk/src/lib.rs +++ b/crates/basilica-sdk/src/lib.rs @@ -9,12 +9,14 @@ pub mod auth; pub mod client; pub mod error; pub mod jobs; +pub mod rl; pub mod types; // Re-export main types pub use client::{BasilicaClient, ClientBuilder}; pub use error::{ApiError, ErrorResponse, Result}; pub use jobs::*; +pub use rl::*; pub use types::*; /// SDK version diff --git a/crates/basilica-sdk/src/rl.rs b/crates/basilica-sdk/src/rl.rs new file mode 100644 index 000000000..bdd616db3 --- /dev/null +++ b/crates/basilica-sdk/src/rl.rs @@ -0,0 +1,338 @@ +//! RL Training API client: GRPO post-training on the Basilica platform. +//! +//! Warm GPU clusters, GRPO jobs with custom rewards / datasets / in-cluster +//! LLM-judge, and the declarative manifest surface. The request DTOs here +//! mirror `basilica-api`'s `/rl/*` route DTOs field-for-field (the server +//! uses `deny_unknown_fields`, so the wire shape is the contract); keeping +//! them as shared serde types is what lets the compiler — not a runtime +//! test — catch drift from the server. + +use serde::{Deserialize, Serialize}; + +// Client methods live in `client.rs` (the crate convention: DTOs here, the +// `impl BasilicaClient` beside the private get/post transport helpers). + +// --------------------------------------------------------------------------- +// Cluster request DTOs +// --------------------------------------------------------------------------- + +/// A trainer/rollout fleet's per-pod GPU shape. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RlGpuRequest { + /// GPU model name, e.g. `H100`, `H200`. + pub model: String, + /// GPUs per pod. The trainer total (`count * replicas`) must divide the + /// GRPO train batch — admission rejects shapes that cannot (valid: 1/2/4/8). + pub count: u32, + /// Optional minimum VRAM per GPU in GB. >=140 marks an H200-class fleet, + /// which the >=16B recipe requires. + #[serde(skip_serializing_if = "Option::is_none")] + pub min_memory_gb: Option, +} + +/// One fleet (trainer or rollout). +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RlFleetRequest { + /// Pod replicas (v0 admits exactly 1 for the trainer). + pub replicas: u32, + /// Per-pod GPU shape. + pub gpu: RlGpuRequest, +} + +/// Create a warm RL cluster (`POST /rl/clusters`). +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CreateRlClusterRequest { + /// Optional cluster name (DNS-1035; generated when omitted). + #[serde(skip_serializing_if = "Option::is_none")] + pub name: Option, + /// Base model to pin, e.g. `Qwen/Qwen2.5-7B-Instruct`. Must be on the + /// platform allowlist; >=16B models require H200-class trainers. + pub base_model: String, + /// Trainer fleet (FSDP). + pub trainer: RlFleetRequest, + /// Rollout fleet (vLLM). + pub rollout: RlFleetRequest, + /// Optional idle-TTL (e.g. `30m`) after which an idle cluster reaps itself. + #[serde(skip_serializing_if = "Option::is_none")] + pub idle_ttl: Option, + /// Forward-compat catch-all: fields this SDK version doesn't know are + /// preserved verbatim on the wire (the `body=` escape hatch depends on + /// this — server-side schema additions must never be silently dropped). + #[serde(flatten)] + pub extra: serde_json::Map, +} + +// --------------------------------------------------------------------------- +// Job request DTOs +// --------------------------------------------------------------------------- + +/// In-cluster LLM-judge opt-in (WS-G v0.5). The judge is a platform-owned +/// vLLM pod serving an allowlisted open model; the reward reaches it via +/// `ctx["judge"](prompt, ...) -> str`. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RlJudgeRequest { + /// Judge model id; omit for the platform default (must be on the judge + /// allowlist). + #[serde(skip_serializing_if = "Option::is_none")] + pub model: Option, +} + +/// Custom reward: a user Python scoring function run in the isolated, +/// credential-free, zero-egress executor pod. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RlRewardRequest { + /// Reward ref, `user:`. + #[serde(rename = "ref")] + pub reward_ref: String, + /// Reward source: stdlib Python defining + /// `reward(prompt, completion, **ctx) -> float` (<=64 KiB). + pub source: String, + /// Optional in-cluster LLM-judge. + #[serde(skip_serializing_if = "Option::is_none")] + pub judge: Option, +} + +/// A public Hugging Face dataset source + column mapping. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RlHfDatasetSource { + /// HF dataset id, `org/name` (public datasets only). + pub repo: String, + /// Optional HF config (subset) name. + #[serde(skip_serializing_if = "Option::is_none")] + pub config: Option, + /// Split to load, e.g. `train`. + pub split: String, + /// Column holding the prompt text. + pub prompt_column: String, + /// Column handed to the reward as `ground_truth`. + pub answer_column: String, +} + +/// Custom dataset: platform code fetches and renders it (dataset-as-data; +/// no user code touches the data path). +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RlDatasetRequest { + /// Dataset ref, `user:`. + #[serde(rename = "ref")] + pub dataset_ref: String, + /// The public HF source + column mapping. + pub hf: RlHfDatasetSource, +} + +/// Create a GRPO training job on a Ready cluster (`POST /rl/jobs`). +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CreateRlJobRequest { + /// Name of the warm cluster to bind to. + pub cluster_ref: String, + /// Optional job name (generated when omitted). + #[serde(skip_serializing_if = "Option::is_none")] + pub name: Option, + /// Algorithm, e.g. `grpo`. + pub algorithm: String, + /// Custom reward; omit for the builtin reward. + #[serde(skip_serializing_if = "Option::is_none")] + pub reward: Option, + /// Custom dataset; omit for the builtin GSM8K dataset. + #[serde(skip_serializing_if = "Option::is_none")] + pub dataset: Option, + /// Training steps (<=3000 on custom datasets). + pub max_steps: u32, + /// Optional learning-rate override (string, e.g. `3.0e-6`). + #[serde(skip_serializing_if = "Option::is_none")] + pub lr: Option, + /// Forward-compat catch-all: fields this SDK version doesn't know are + /// preserved verbatim on the wire (the `body=` escape hatch depends on + /// this — server-side schema additions must never be silently dropped). + #[serde(flatten)] + pub extra: serde_json::Map, +} + +/// Declarative manifest: one document that renders a cluster and/or a job +/// (`POST /rl/manifest`). Freeform to mirror the server's document surface. +pub type RlManifestRequest = serde_json::Value; + +// --------------------------------------------------------------------------- +// Response DTOs +// --------------------------------------------------------------------------- + +/// Response after creating a cluster. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CreateRlClusterResponse { + /// The cluster's name (its `clusterRef`). + pub name: String, + /// Unique identifier. + pub uid: String, + /// Lifecycle phase at creation (`Provisioning`). + pub phase: String, +} + +/// Cluster status (`GET /rl/clusters/{name}`). +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RlClusterStatusResponse { + /// `Provisioning` | `Warming` | `Ready` | `Degraded` | `Terminating`. + pub phase: String, + /// Whether every fleet pod verified the pinned base model. + #[serde(skip_serializing_if = "Option::is_none")] + pub model_loaded: Option, + /// The bound job's name, if any. + #[serde(skip_serializing_if = "Option::is_none")] + pub active_job_name: Option, +} + +/// Response after creating a job. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CreateRlJobResponse { + /// The job's name (its identifier for status). + pub name: String, + /// Unique identifier. + pub uid: String, + /// Lifecycle phase at creation (`Pending`). + pub phase: String, + /// sha256 of the admitted reward source (absent for builtin-reward jobs). + #[serde(skip_serializing_if = "Option::is_none")] + pub reward_sha256: Option, +} + +/// Latest training metrics. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RlJobMetrics { + /// Latest training loss. + #[serde(skip_serializing_if = "Option::is_none")] + pub loss: Option, + /// Latest mean reward. + #[serde(skip_serializing_if = "Option::is_none")] + pub reward_mean: Option, + /// Latest KL divergence. + #[serde(skip_serializing_if = "Option::is_none")] + pub kl: Option, +} + +/// Job status (`GET /rl/jobs/{name}`). +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RlJobStatusResponse { + /// `Pending` | `Binding` | `Running` | `Succeeded` | `Failed` | `TimedOut`. + pub phase: String, + /// Best-effort latest training step. + #[serde(skip_serializing_if = "Option::is_none")] + pub step: Option, + /// Latest training metrics. + #[serde(skip_serializing_if = "Option::is_none")] + pub metrics: Option, + /// Object-store location of the trained model (null until bound). + #[serde(rename = "artifactURI", skip_serializing_if = "Option::is_none")] + pub artifact_uri: Option, +} + +/// Manifest submission result. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RlManifestResponse { + /// The created cluster, when the manifest carried a `cluster` block. + #[serde(skip_serializing_if = "Option::is_none")] + pub cluster: Option, + /// The created job, when the manifest carried a `job` block. + #[serde(skip_serializing_if = "Option::is_none")] + pub job: Option, +} + +#[cfg(test)] +mod tests { + use super::*; + + // The SDK request DTOs must serialize to the exact wire shape the server's + // deny_unknown_fields DTOs accept. These pin the field renames and casing + // that the server contract depends on. + #[test] + fn job_request_wire_shape() { + let req = CreateRlJobRequest { + cluster_ref: "my-pool".into(), + name: None, + algorithm: "grpo".into(), + reward: Some(RlRewardRequest { + reward_ref: "user:my-reward".into(), + source: "def reward(p, c, **k):\n return 1.0\n".into(), + judge: Some(RlJudgeRequest { model: None }), + }), + dataset: Some(RlDatasetRequest { + dataset_ref: "user:my-data".into(), + hf: RlHfDatasetSource { + repo: "openai/gsm8k".into(), + config: Some("main".into()), + split: "train".into(), + prompt_column: "question".into(), + answer_column: "answer".into(), + }, + }), + max_steps: 50, + lr: None, + extra: Default::default(), + }; + let v = serde_json::to_value(&req).unwrap(); + assert_eq!(v["clusterRef"], "my-pool"); + assert_eq!(v["maxSteps"], 50); + // the serde renames: nested identity fields are `ref` + assert_eq!(v["reward"]["ref"], "user:my-reward"); + assert_eq!(v["dataset"]["ref"], "user:my-data"); + assert_eq!(v["dataset"]["hf"]["promptColumn"], "question"); + // judge with no model serializes as an empty object (not null) + assert!(v["reward"]["judge"].is_object()); + assert_eq!(v["reward"]["judge"].as_object().unwrap().len(), 0); + // None optionals are ABSENT (deny_unknown_fields tolerates absence, + // but a null in a non-Option server slot would 400) + assert!(v.get("name").is_none()); + assert!(v.get("lr").is_none()); + } + + #[test] + fn cluster_request_wire_shape() { + let req = CreateRlClusterRequest { + name: Some("my-pool".into()), + base_model: "Qwen/Qwen2.5-7B-Instruct".into(), + trainer: RlFleetRequest { + replicas: 1, + gpu: RlGpuRequest { + model: "H100".into(), + count: 4, + min_memory_gb: None, + }, + }, + rollout: RlFleetRequest { + replicas: 1, + gpu: RlGpuRequest { + model: "H100".into(), + count: 4, + min_memory_gb: None, + }, + }, + idle_ttl: Some("30m".into()), + extra: Default::default(), + }; + let v = serde_json::to_value(&req).unwrap(); + assert_eq!(v["baseModel"], "Qwen/Qwen2.5-7B-Instruct"); + assert_eq!(v["trainer"]["gpu"]["count"], 4); + assert_eq!(v["idleTtl"], "30m"); + assert!(v["trainer"]["gpu"].get("minMemoryGb").is_none()); + } + + #[test] + fn job_status_parses_artifact_uri_rename() { + let body = r#"{"phase":"Succeeded","step":50,"artifactURI":"s3://x/uid"}"#; + let s: RlJobStatusResponse = serde_json::from_str(body).unwrap(); + assert_eq!(s.phase, "Succeeded"); + assert_eq!(s.artifact_uri.as_deref(), Some("s3://x/uid")); + assert_eq!(s.step, Some(50)); + } +} From c7594115eca0cd2eb9b103893b492e9d08cde835 Mon Sep 17 00:00:00 2001 From: Alexander Paskov Date: Mon, 24 Aug 2026 18:48:01 +0300 Subject: [PATCH 2/2] =?UTF-8?q?fix(sdk):=20address=20#557=20review=20?= =?UTF-8?q?=E2=80=94=20importability,=20orphan=20guards,=20wait=20semantic?= =?UTF-8?q?s,=20name=20validation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review findings from Talos + CodeRabbit, all verified before fixing: - `from basilica import RlNamespace` actually works now (it was in __all__ but only imported inside the .rl property): PEP 562 module __getattr__ keeps the import lazy. [confirmed broken by test] - create_job raises on orphan kwargs (reward_source without reward_name, dataset fields without dataset_name) instead of silently running the BUILTIN reward/dataset on a paid GPU job; symmetric with the existing judge guard. - wait_cluster: raises immediately on Terminating (it can never become Ready; burning the full 1800s timeout would hide the failure). Degraded deliberately keeps polling — fleets recover from transient unhealth; documented in the code. - Both wait loops tolerate transient poll errors (ConnectionError / RuntimeError) up to 5 CONSECUTIVE failures, reset on success — a single LB blip must not abort a 6-hour wait. - body= no longer needs placeholder kwargs: the other required kwargs are checked only when building the request. - get_rl_cluster/get_rl_job validate DNS-1035 client-side in the core before the name reaches the URL path (clearer error + path hardening). - Dead `or {}` removed; module docstring rewritten timeless (no PR archaeology). Tests: 15/15 contract tests against the compiled extension (5 new: wait_cluster poll + Terminating, transient-error budget, orphan guards, client-side name rejection); core 92 green; clippy clean. --- .../python/basilica/__init__.py | 12 +++ .../basilica-sdk-python/python/basilica/rl.py | 102 ++++++++++++++---- .../tests/test_rl_client.py | 68 +++++++++++- crates/basilica-sdk/src/client.rs | 26 +++++ 4 files changed, 185 insertions(+), 23 deletions(-) diff --git a/crates/basilica-sdk-python/python/basilica/__init__.py b/crates/basilica-sdk-python/python/basilica/__init__.py index e0be2fb63..b86cc73bb 100644 --- a/crates/basilica-sdk-python/python/basilica/__init__.py +++ b/crates/basilica-sdk-python/python/basilica/__init__.py @@ -452,6 +452,18 @@ def _build_inference_health_check(port: int) -> HealthCheckConfig: __version__ = _pkg_version("basilica-sdk") except PackageNotFoundError: __version__ = "0.0.0+unknown" + + +def __getattr__(name): + # PEP 562: `from basilica import RlNamespace` must work because it is in + # __all__, but basilica.rl stays lazily imported (mirrors the .rl property) + if name == "RlNamespace": + from basilica.rl import RlNamespace + + return RlNamespace + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + __all__ = [ # Main client "BasilicaClient", diff --git a/crates/basilica-sdk-python/python/basilica/rl.py b/crates/basilica-sdk-python/python/basilica/rl.py index e96d7f475..eef75dbc2 100644 --- a/crates/basilica-sdk-python/python/basilica/rl.py +++ b/crates/basilica-sdk-python/python/basilica/rl.py @@ -17,16 +17,16 @@ ... ) >>> final = client.rl.wait_job(job["name"]) # {phase, step, metrics, artifactURI} -THIN WRAPPER over the compiled core (#1509 review round): this module builds -the ergonomic kwargs into wire dicts and hands them to the Rust binding's -``rl_*`` methods, which serde-validate against the core's typed DTOs +THIN WRAPPER over the compiled core: this module builds the ergonomic +kwargs into wire dicts and hands them to the Rust binding's ``rl_*`` +methods, which serde-validate against the core's typed DTOs (``basilica_sdk::rl`` — the compile-time-shared contract with the server) and send through the core transport. That inherits the full auth chain -(explicit key, BASILICA_API_TOKEN, and the CLI-login token fallback the -earlier pure-python transport could not reach) and the core's error -mapping: non-2xx surfaces as ValueError (bad request), PermissionError -(authz), ConnectionError (transport), FileNotFoundError (not found), or -RuntimeError (server error), each carrying the server's message verbatim. +(explicit key, BASILICA_API_TOKEN, and the CLI-login token fallback) and +the core's error mapping: non-2xx surfaces as ValueError (bad request), +PermissionError (authz), ConnectionError (transport), FileNotFoundError +(not found), or RuntimeError (server error), each carrying the server's +message verbatim. The ``body=`` escape hatch on every create call sends a raw dict; unknown fields survive the typed round-trip verbatim (serde-flatten catch-alls in @@ -41,6 +41,13 @@ from typing import Any, Optional _TERMINAL_JOB_PHASES = frozenset({"Succeeded", "Failed", "TimedOut"}) +# Degraded is deliberately NOT here: a cluster degrades on transient fleet +# unhealth (pod restart, node blip) and can recover to Ready; only +# Terminating can never become Ready again. +_DEAD_CLUSTER_PHASES = frozenset({"Terminating"}) +# A single LB 502 or connection reset must not abort a multi-hour wait; +# this many CONSECUTIVE poll failures (reset on any success) give up. +_POLL_FAILURE_BUDGET = 5 def _drop_none(d: dict) -> dict: @@ -59,8 +66,8 @@ def __init__(self, core: Any): def create_cluster( self, *, - base_model: str, - gpu_model: str, + base_model: Optional[str] = None, + gpu_model: Optional[str] = None, trainer_gpus: int = 4, rollout_gpus: int = 4, name: Optional[str] = None, @@ -72,9 +79,14 @@ def create_cluster( Certified shapes: 4+4 (H100 for <16B models, H200-class for >=16B — admission rejects bad pairings with an actionable message). - ``body`` overrides everything (raw wire dict, escape hatch). + ``body`` replaces the built request entirely (raw wire dict, escape + hatch) — no other kwargs are consulted when it is given. """ if body is None: + if base_model is None or gpu_model is None: + raise ValueError( + "base_model and gpu_model are required (unless a raw body= is given)" + ) def fleet(count: int) -> dict: return { @@ -101,16 +113,31 @@ def get_cluster(self, name: str) -> dict: def wait_cluster( self, name: str, timeout_s: float = 1800.0, poll_s: float = 15.0 ) -> dict: - """Poll until phase == Ready (raises TimeoutError otherwise).""" + """Poll until phase == Ready. Raises RuntimeError immediately on + Terminating (it can never become Ready — waiting out the full + timeout would hide the failure), TimeoutError on the deadline. + Degraded keeps polling: fleets recover from transient unhealth.""" deadline = time.monotonic() + timeout_s + failures = 0 while True: - cluster = self.get_cluster(name) - if cluster.get("phase") == "Ready": + try: + cluster = self.get_cluster(name) + except (ConnectionError, RuntimeError): + failures += 1 + if failures >= _POLL_FAILURE_BUDGET: + raise + time.sleep(poll_s) + continue + failures = 0 + phase = cluster.get("phase") + if phase == "Ready": return cluster + if phase in _DEAD_CLUSTER_PHASES: + raise RuntimeError(f"cluster {name!r} entered {phase}: {cluster}") if time.monotonic() >= deadline: raise TimeoutError( f"cluster {name!r} not Ready after {timeout_s}s " - f"(last phase: {cluster.get('phase')!r})" + f"(last phase: {phase!r})" ) time.sleep(poll_s) @@ -119,8 +146,8 @@ def wait_cluster( def create_job( self, *, - cluster: str, - max_steps: int, + cluster: Optional[str] = None, + max_steps: Optional[int] = None, name: Optional[str] = None, algorithm: str = "grpo", # custom reward (user: + inline source); omit for the builtin @@ -144,16 +171,27 @@ def create_job( ``reward(prompt, completion, **ctx) -> float``; it runs in an isolated credential-free pod. ``judge=True`` exposes ``ctx["judge"](prompt)`` backed by an in-cluster judge model - (requires a custom reward). ``body`` overrides everything. + (requires a custom reward). ``body`` replaces the built request + entirely — no other kwargs are consulted when it is given. + + Orphan kwargs raise: a ``reward_source`` without ``reward_name`` (or + dataset fields without ``dataset_name``) would otherwise be silently + dropped and the BUILTIN reward/dataset would run on a paid GPU job. """ if body is None: + if cluster is None or max_steps is None: + raise ValueError( + "cluster and max_steps are required (unless a raw body= is given)" + ) reward = None if reward_name is not None: if reward_source is None: raise ValueError("reward_source is required with reward_name") reward = {"ref": f"user:{reward_name}", "source": reward_source} if judge or judge_model: - reward["judge"] = _drop_none({"model": judge_model}) or {} + reward["judge"] = _drop_none({"model": judge_model}) + elif reward_source is not None: + raise ValueError("reward_name is required with reward_source") elif judge or judge_model: raise ValueError( "judge requires a custom reward (it is called from your reward code)" @@ -172,6 +210,17 @@ def create_job( } ), } + elif any( + v is not None + for v in ( + dataset_repo, + dataset_config, + dataset_split, + prompt_column, + answer_column, + ) + ): + raise ValueError("dataset_name is required with dataset fields") body = _drop_none( { "clusterRef": cluster, @@ -194,10 +243,21 @@ def wait_job( """Poll until the job is terminal (Succeeded/Failed/TimedOut) and return the final document either way — check ``phase`` yourself; raising on Failed would hide the failure detail behind an - exception.""" + exception. Transient poll errors are tolerated up to + ``_POLL_FAILURE_BUDGET`` consecutive failures — a single LB blip + must not abort a multi-hour wait.""" deadline = time.monotonic() + timeout_s + failures = 0 while True: - job = self.get_job(name) + try: + job = self.get_job(name) + except (ConnectionError, RuntimeError): + failures += 1 + if failures >= _POLL_FAILURE_BUDGET: + raise + time.sleep(poll_s) + continue + failures = 0 if job.get("phase") in _TERMINAL_JOB_PHASES: return job if time.monotonic() >= deadline: diff --git a/crates/basilica-sdk-python/tests/test_rl_client.py b/crates/basilica-sdk-python/tests/test_rl_client.py index f168c5d56..7f7c80018 100644 --- a/crates/basilica-sdk-python/tests/test_rl_client.py +++ b/crates/basilica-sdk-python/tests/test_rl_client.py @@ -148,17 +148,32 @@ def test_builtin_job_minimal_body(server): def test_raw_body_escape_hatch_preserves_unknown_fields(server): # serde-flatten catch-alls in the core DTOs: a field this SDK version - # doesn't know must SURVIVE the typed round-trip verbatim. + # doesn't know must SURVIVE the typed round-trip verbatim. body= needs + # no placeholder kwargs — it replaces the built request entirely. base, rec = server raw = {"clusterRef": "x", "algorithm": "grpo", "maxSteps": 1, "futureField": True} rec.responses = [(200, {"name": "j1", "uid": "u2", "phase": "Pending"})] - rl(base).create_job(cluster="ignored", max_steps=99, body=raw) + rl(base).create_job(body=raw) (r,) = rec.requests assert r["body"]["futureField"] is True assert r["body"]["maxSteps"] == 1 assert r["body"]["clusterRef"] == "x" +def test_orphan_kwargs_raise_instead_of_silently_dropping(server): + # reward_source without reward_name (or dataset fields without + # dataset_name) would otherwise silently run the BUILTIN reward/dataset + # on a paid GPU job. + base, rec = server + with pytest.raises(ValueError, match="reward_name is required"): + rl(base).create_job(cluster="c", max_steps=3, reward_source="def reward(): ...") + with pytest.raises(ValueError, match="dataset_name is required"): + rl(base).create_job(cluster="c", max_steps=3, dataset_repo="openai/gsm8k") + with pytest.raises(ValueError, match="cluster and max_steps are required"): + rl(base).create_job(reward_name="r", reward_source="...") + assert rec.requests == [] + + def test_wait_job_polls_to_terminal(server): base, rec = server rec.responses = [ @@ -206,3 +221,52 @@ def test_manifest_posts_verbatim(server): rl(base).submit_manifest(doc) (r,) = rec.requests assert (r["method"], r["path"], r["body"]) == ("POST", "/rl/manifest", doc) + + +def test_wait_cluster_polls_to_ready(server): + base, rec = server + rec.responses = [ + (200, {"name": "c1", "uid": "u1", "phase": "Provisioning"}), + (200, {"name": "c1", "uid": "u1", "phase": "Warming"}), + (200, {"name": "c1", "uid": "u1", "phase": "Ready"}), + ] + final = rl(base).wait_cluster("c1", timeout_s=30, poll_s=0.01) + assert final["phase"] == "Ready" + assert len(rec.requests) == 3 + assert all(r["path"] == "/rl/clusters/c1" for r in rec.requests) + + +def test_wait_cluster_raises_immediately_on_terminating(server): + # Terminating can never become Ready: waiting out the full timeout would + # hide the failure. (Degraded, by contrast, keeps polling — it recovers.) + base, rec = server + rec.responses = [(200, {"name": "c1", "uid": "u1", "phase": "Terminating"})] + with pytest.raises(RuntimeError, match="entered Terminating"): + rl(base).wait_cluster("c1", timeout_s=30, poll_s=0.01) + assert len(rec.requests) == 1 + + +def test_wait_job_survives_transient_poll_errors(server): + # A single LB 502 mid-wait must not abort a multi-hour poll; only + # CONSECUTIVE failures beyond the budget give up. + base, rec = server + rec.responses = [ + (200, {"phase": "Running"}), + (500, _api_error("upstream blip", code="BASILICA_API_INTERNAL_ERROR")), + (500, _api_error("upstream blip", code="BASILICA_API_INTERNAL_ERROR")), + (200, {"phase": "Succeeded"}), + ] + final = rl(base).wait_job("j1", timeout_s=30, poll_s=0.01) + assert final["phase"] == "Succeeded" + assert len(rec.requests) == 4 + + +def test_invalid_name_rejected_client_side(server): + # names become URL path segments + k8s object names: DNS-1035 is checked + # in the core before any HTTP + base, rec = server + with pytest.raises(ValueError, match="DNS-1035"): + rl(base).get_job("Bad/Name") + with pytest.raises(ValueError, match="DNS-1035"): + rl(base).get_cluster("-leading-dash") + assert rec.requests == [] diff --git a/crates/basilica-sdk/src/client.rs b/crates/basilica-sdk/src/client.rs index 3ddc702a7..9793afda2 100644 --- a/crates/basilica-sdk/src/client.rs +++ b/crates/basilica-sdk/src/client.rs @@ -333,6 +333,30 @@ impl BasilicaClient { // ----- RL Training API (GRPO post-training) -------------------------- + /// RL names become URL path segments and Kubernetes object names, so the + /// server enforces DNS-1035. Checking it client-side keeps a malformed + /// name from ever reaching the path (and yields a clearer error). + fn validate_rl_name(name: &str) -> Result<()> { + let dns1035 = !name.is_empty() + && name.len() <= 63 + && name.starts_with(|c: char| c.is_ascii_lowercase()) + && !name.ends_with('-') + && name + .chars() + .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-'); + if dns1035 { + Ok(()) + } else { + Err(ApiError::InvalidRequest { + message: format!( + "invalid RL resource name {name:?}: must be DNS-1035 \ + (lowercase alphanumeric or '-', start with a letter, \ + not end with '-', max 63 chars)" + ), + }) + } + } + /// Create a warm RL cluster. Poll [`get_rl_cluster`](Self::get_rl_cluster) /// until its phase is `Ready`. pub async fn create_rl_cluster( @@ -344,6 +368,7 @@ impl BasilicaClient { /// Get a cluster's status. pub async fn get_rl_cluster(&self, name: &str) -> Result { + Self::validate_rl_name(name)?; self.get(&format!("/rl/clusters/{}", name)).await } @@ -354,6 +379,7 @@ impl BasilicaClient { /// Get a job's status (phase, step, metrics, `artifactURI`). pub async fn get_rl_job(&self, name: &str) -> Result { + Self::validate_rl_name(name)?; self.get(&format!("/rl/jobs/{}", name)).await }