diff --git a/crates/basilica-sdk-python/python/basilica/__init__.py b/crates/basilica-sdk-python/python/basilica/__init__.py index 02bff4c33..b86cc73bb 100644 --- a/crates/basilica-sdk-python/python/basilica/__init__.py +++ b/crates/basilica-sdk-python/python/basilica/__init__.py @@ -452,9 +452,23 @@ 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", + # RL training namespace (client.rl; module: basilica.rl) + "RlNamespace", # Decorator API "deployment", "DeployedFunction", @@ -600,12 +614,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..eef75dbc2 --- /dev/null +++ b/crates/basilica-sdk-python/python/basilica/rl.py @@ -0,0 +1,273 @@ +"""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: 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) 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"}) +# 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: + 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: Optional[str] = None, + gpu_model: Optional[str] = None, + 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`` 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 { + "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 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: + 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: {phase!r})" + ) + time.sleep(poll_s) + + # -- jobs -------------------------------------------------------------- + + def create_job( + self, + *, + 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 + 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`` 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}) + 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)" + ) + 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, + } + ), + } + 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, + "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. 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: + 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: + 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..7f7c80018 --- /dev/null +++ b/crates/basilica-sdk-python/tests/test_rl_client.py @@ -0,0 +1,272 @@ +"""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. 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(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 = [ + (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) + + +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 f7fd9bc57..9793afda2 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,66 @@ impl BasilicaClient { self.post(&path, &serde_json::json!({})).await } + // ----- 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( + &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::validate_rl_name(name)?; + 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::validate_rl_name(name)?; + 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)); + } +}