mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-23 20:23:00 +00:00
feat(hanzo-train): add Tinker-shaped training client (0.1.0)
Hand-written package for the engine /v1/training API: ServiceClient + TrainingClient (create/list/get/delete/forward_backward/optim_step/ sample/save_weights), dataclasses, sync httpx, completed-future .result() wrapper so tinker fut.result() code ports 1:1. 22 pytest tests over httpx.MockTransport, no network. Wired like siblings: publish-pypi hanzo-train-* tag case, rye+uv workspace member, uv source, pyright exclude. uv.lock relock also catches the lockfile up to on-disk truth (hanzoai 3.1.1, hanzo-tools-browser 0.5.10) -- no package versions changed.
This commit is contained in:
@@ -131,6 +131,8 @@ jobs:
|
||||
PACKAGES="hanzo-aci"
|
||||
elif [[ $TAG == hanzo-dev-* ]]; then
|
||||
PACKAGES="hanzo-dev"
|
||||
elif [[ $TAG == hanzo-train-* ]]; then
|
||||
PACKAGES="hanzo-train"
|
||||
elif [[ $TAG == hanzoai-* ]]; then
|
||||
PACKAGES="hanzoai"
|
||||
elif [[ $TAG == hanzo-* ]]; then
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
# hanzo-train
|
||||
|
||||
Tinker-shaped Python client for the Hanzo Engine training API.
|
||||
|
||||
It mirrors the shapes of Thinking Machines' [`tinker`](https://github.com/thinking-machines-lab/tinker)
|
||||
SDK, so training loops written against tinker port across unchanged, while using
|
||||
Hanzo-canonical field names. The transport is synchronous `httpx` with no retries;
|
||||
the four training ops return a completed future exposing `.result(timeout=None)`
|
||||
so `fut.result()` code ports 1:1.
|
||||
|
||||
## Install
|
||||
|
||||
```bash
|
||||
pip install hanzo-train
|
||||
```
|
||||
|
||||
## Quickstart
|
||||
|
||||
```python
|
||||
from hanzo_train import ServiceClient, LoraConfig, AdamParams, SamplingParams, Datum, ModelInput
|
||||
|
||||
sc = ServiceClient(base_url="http://localhost:1234", api_key=None) # api_key -> Authorization: Bearer
|
||||
|
||||
tc = sc.create_lora_training_client(
|
||||
"HuggingFaceTB/SmolLM2-135M", lora_config=LoraConfig(rank=16), wait=True, timeout=600.0
|
||||
)
|
||||
# wait=True polls the client until status is ready (raises on failed); wait=False returns immediately.
|
||||
|
||||
out = tc.forward_backward([
|
||||
{"prompt": "2+2=", "completion": "4"},
|
||||
Datum(model_input=ModelInput(tokens=[1, 2, 3]), target_tokens=[2, 3, 4], weights=[0.0, 1.0, 1.0]),
|
||||
]).result()
|
||||
# out.loss, out.num_tokens, out.metrics
|
||||
|
||||
tc.optim_step(AdamParams(lr=1e-4)).result()
|
||||
|
||||
resp = tc.sample(
|
||||
prompt="2+2=", sampling_params=SamplingParams(max_tokens=8, temperature=0.0), num_samples=1
|
||||
).result()
|
||||
# resp.sequences[0].tokens / .text
|
||||
|
||||
saved = tc.save_weights_and_get_sampling_client(name="my-adapter").result() # .path, .format == "peft"
|
||||
# alias: tc.save_weights(name="my-adapter")
|
||||
|
||||
info = tc.get_info() # TrainingClientInfo incl. loss_history
|
||||
sc.list_training_clients() # list[TrainingClientInfo]
|
||||
tc.delete()
|
||||
```
|
||||
|
||||
## Errors
|
||||
|
||||
Every non-2xx response raises `HanzoTrainError(status, message)`:
|
||||
|
||||
- `400` — bad input
|
||||
- `404` — unknown training client id
|
||||
- `409` — client still loading, or failed to load
|
||||
|
||||
## API
|
||||
|
||||
`ServiceClient(base_url, api_key=None)`
|
||||
|
||||
- `create_lora_training_client(base_model, lora_config=None, wait=True, timeout=600.0, poll_interval=1.0) -> TrainingClient`
|
||||
- `list_training_clients() -> list[TrainingClientInfo]`
|
||||
- `close()`
|
||||
|
||||
`TrainingClient`
|
||||
|
||||
- `forward_backward(data) -> Future[ForwardBackwardResult]` — `data` items are either
|
||||
`{"prompt": str, "completion": str}` dicts or `Datum` values.
|
||||
- `optim_step(adam_params=None) -> Future[OptimResult]`
|
||||
- `sample(prompt=None, tokens=None, sampling_params=None, num_samples=1) -> Future[SampleResult]`
|
||||
— pass exactly one of `prompt` or `tokens`.
|
||||
- `save_weights_and_get_sampling_client(name, dir=None) -> Future[SaveResult]` (alias: `save_weights`)
|
||||
- `get_info() -> TrainingClientInfo`
|
||||
- `delete()`
|
||||
@@ -0,0 +1,52 @@
|
||||
"""Hanzo Train -- Tinker-shaped client for the Hanzo Engine training API.
|
||||
|
||||
Mirrors the shapes of Thinking Machines' ``tinker`` SDK so training loops port
|
||||
across unchanged. Synchronous httpx under the hood; the four training ops return
|
||||
a completed future exposing ``.result(timeout=None)``.
|
||||
|
||||
Usage::
|
||||
|
||||
from hanzo_train import ServiceClient, LoraConfig, AdamParams, SamplingParams
|
||||
|
||||
sc = ServiceClient("http://localhost:1234")
|
||||
tc = sc.create_lora_training_client("HuggingFaceTB/SmolLM2-135M", LoraConfig(rank=16))
|
||||
out = tc.forward_backward([{"prompt": "2+2=", "completion": "4"}]).result()
|
||||
tc.optim_step(AdamParams(lr=1e-4)).result()
|
||||
resp = tc.sample(prompt="2+2=", sampling_params=SamplingParams(max_tokens=8)).result()
|
||||
"""
|
||||
|
||||
from hanzo_train.client import (
|
||||
AdamParams,
|
||||
Datum,
|
||||
ForwardBackwardResult,
|
||||
HanzoTrainError,
|
||||
LoraConfig,
|
||||
ModelInput,
|
||||
OptimResult,
|
||||
SampleResult,
|
||||
SamplingParams,
|
||||
SaveResult,
|
||||
Sequence,
|
||||
ServiceClient,
|
||||
TrainingClient,
|
||||
TrainingClientInfo,
|
||||
)
|
||||
|
||||
__version__ = "0.1.0"
|
||||
|
||||
__all__ = [
|
||||
"ServiceClient",
|
||||
"TrainingClient",
|
||||
"TrainingClientInfo",
|
||||
"LoraConfig",
|
||||
"AdamParams",
|
||||
"SamplingParams",
|
||||
"Datum",
|
||||
"ModelInput",
|
||||
"ForwardBackwardResult",
|
||||
"OptimResult",
|
||||
"SampleResult",
|
||||
"Sequence",
|
||||
"SaveResult",
|
||||
"HanzoTrainError",
|
||||
]
|
||||
@@ -0,0 +1,335 @@
|
||||
"""Tinker-shaped client for the Hanzo Engine training API.
|
||||
|
||||
The engine exposes a small LoRA training surface under ``/v1/training``. This
|
||||
client mirrors the shapes of Thinking Machines' ``tinker`` SDK -- so training
|
||||
loops written against tinker port across unchanged -- using Hanzo-canonical
|
||||
field names and a single synchronous HTTP path (httpx, no retries).
|
||||
|
||||
Usage::
|
||||
|
||||
from hanzo_train import ServiceClient, LoraConfig, AdamParams, SamplingParams
|
||||
|
||||
sc = ServiceClient("http://localhost:1234")
|
||||
tc = sc.create_lora_training_client("HuggingFaceTB/SmolLM2-135M", LoraConfig(rank=16))
|
||||
tc.forward_backward([{"prompt": "2+2=", "completion": "4"}]).result()
|
||||
tc.optim_step(AdamParams(lr=1e-4)).result()
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable
|
||||
from concurrent.futures import Future
|
||||
from dataclasses import dataclass, field
|
||||
from time import monotonic, sleep
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
_TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
|
||||
|
||||
|
||||
class HanzoTrainError(Exception):
|
||||
"""A non-2xx response from the training API, carrying its status and message."""
|
||||
|
||||
def __init__(self, status: int, message: str) -> None:
|
||||
super().__init__(f"[{status}] {message}")
|
||||
self.status = status
|
||||
self.message = message
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelInput:
|
||||
"""Pre-tokenized model input."""
|
||||
|
||||
tokens: list[int]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Datum:
|
||||
"""One tokenized example: input tokens, target tokens, and per-token weights."""
|
||||
|
||||
model_input: ModelInput
|
||||
target_tokens: list[int]
|
||||
weights: list[float]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LoraConfig:
|
||||
"""LoRA adapter configuration; defaults match the engine wire defaults."""
|
||||
|
||||
rank: int = 16
|
||||
alpha: float = 32.0
|
||||
target_modules: list[str] = field(default_factory=lambda: list(_TARGET_MODULES))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AdamParams:
|
||||
"""Adam optimizer parameters for one step; defaults match the wire defaults."""
|
||||
|
||||
lr: float = 1e-4
|
||||
beta1: float = 0.9
|
||||
beta2: float = 0.95
|
||||
eps: float = 1e-8
|
||||
weight_decay: float = 0.0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SamplingParams:
|
||||
"""Sampling parameters; defaults match the engine wire defaults."""
|
||||
|
||||
max_tokens: int = 64
|
||||
temperature: float = 1.0
|
||||
top_k: int = -1
|
||||
top_p: float = 1.0
|
||||
seed: int = 0
|
||||
stop_tokens: list[int] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ForwardBackwardResult:
|
||||
"""Loss and token count returned by a forward-backward pass."""
|
||||
|
||||
loss: float
|
||||
num_tokens: int
|
||||
metrics: dict[str, float]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OptimResult:
|
||||
"""Cumulative optimizer step count after an optim step."""
|
||||
|
||||
optim_steps: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Sequence:
|
||||
"""One sampled sequence: token ids and decoded text."""
|
||||
|
||||
tokens: list[int]
|
||||
text: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SampleResult:
|
||||
"""Sequences produced by a sample call."""
|
||||
|
||||
sequences: list[Sequence]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SaveResult:
|
||||
"""Where a saved adapter lives and in what format."""
|
||||
|
||||
path: str
|
||||
format: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TrainingClientInfo:
|
||||
"""Server-side state of a training client."""
|
||||
|
||||
id: str
|
||||
base_model: str
|
||||
status: str
|
||||
lora_config: LoraConfig
|
||||
forward_backward_calls: int
|
||||
optim_steps: int
|
||||
error: str | None = None
|
||||
trainable_params: int | None = None
|
||||
last_loss: float | None = None
|
||||
loss_history: list[float] = field(default_factory=list)
|
||||
|
||||
|
||||
def _future[T](value: T) -> Future[T]:
|
||||
fut: Future[T] = Future()
|
||||
fut.set_result(value)
|
||||
return fut
|
||||
|
||||
|
||||
def _datum_wire(datum: Datum | dict[str, Any]) -> dict[str, Any]:
|
||||
if isinstance(datum, Datum):
|
||||
return {
|
||||
"model_input": {"tokens": list(datum.model_input.tokens)},
|
||||
"target_tokens": list(datum.target_tokens),
|
||||
"weights": list(datum.weights),
|
||||
}
|
||||
return datum
|
||||
|
||||
|
||||
def _info(data: dict[str, Any]) -> TrainingClientInfo:
|
||||
cfg = data["lora_config"]
|
||||
return TrainingClientInfo(
|
||||
id=data["id"],
|
||||
base_model=data["base_model"],
|
||||
status=data["status"],
|
||||
lora_config=LoraConfig(
|
||||
rank=cfg["rank"], alpha=cfg["alpha"], target_modules=list(cfg["target_modules"])
|
||||
),
|
||||
forward_backward_calls=data["forward_backward_calls"],
|
||||
optim_steps=data["optim_steps"],
|
||||
error=data.get("error"),
|
||||
trainable_params=data.get("trainable_params"),
|
||||
last_loss=data.get("last_loss"),
|
||||
loss_history=list(data.get("loss_history", [])),
|
||||
)
|
||||
|
||||
|
||||
class ServiceClient:
|
||||
"""Entry point to the Hanzo Engine training API."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
api_key: str | None = None,
|
||||
transport: httpx.BaseTransport | None = None,
|
||||
) -> None:
|
||||
headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
|
||||
self._http = httpx.Client(
|
||||
base_url=base_url.rstrip("/"),
|
||||
headers=headers,
|
||||
transport=transport,
|
||||
timeout=httpx.Timeout(None, connect=10.0),
|
||||
)
|
||||
|
||||
def _request(self, method: str, path: str, body: Any = None) -> Any:
|
||||
resp = self._http.request(method, path, json=body)
|
||||
if resp.status_code >= 400:
|
||||
raise HanzoTrainError(resp.status_code, resp.text)
|
||||
return resp.json()
|
||||
|
||||
def create_lora_training_client(
|
||||
self,
|
||||
base_model: str,
|
||||
lora_config: LoraConfig | None = None,
|
||||
wait: bool = True,
|
||||
timeout: float = 600.0,
|
||||
poll_interval: float = 1.0,
|
||||
) -> TrainingClient:
|
||||
cfg = lora_config or LoraConfig()
|
||||
data = self._request(
|
||||
"POST",
|
||||
"/v1/training/clients",
|
||||
{
|
||||
"base_model": base_model,
|
||||
"lora_config": {
|
||||
"rank": cfg.rank,
|
||||
"alpha": cfg.alpha,
|
||||
"target_modules": list(cfg.target_modules),
|
||||
},
|
||||
},
|
||||
)
|
||||
tc = TrainingClient(self, data["id"], poll_interval)
|
||||
if wait:
|
||||
tc._wait_ready(timeout)
|
||||
return tc
|
||||
|
||||
def list_training_clients(self) -> list[TrainingClientInfo]:
|
||||
data = self._request("GET", "/v1/training/clients")
|
||||
return [_info(c) for c in data["clients"]]
|
||||
|
||||
def close(self) -> None:
|
||||
self._http.close()
|
||||
|
||||
|
||||
class TrainingClient:
|
||||
"""A handle to one LoRA training client on the engine."""
|
||||
|
||||
def __init__(self, service: ServiceClient, id: str, poll_interval: float = 1.0) -> None:
|
||||
self.id = id
|
||||
self._service = service
|
||||
self._poll_interval = poll_interval
|
||||
|
||||
def _path(self, suffix: str = "") -> str:
|
||||
return f"/v1/training/clients/{self.id}{suffix}"
|
||||
|
||||
def forward_backward(
|
||||
self, data: Iterable[Datum | dict[str, Any]]
|
||||
) -> Future[ForwardBackwardResult]:
|
||||
result = self._service._request(
|
||||
"POST", self._path("/forward_backward"), {"data": [_datum_wire(d) for d in data]}
|
||||
)
|
||||
return _future(
|
||||
ForwardBackwardResult(
|
||||
loss=result["loss"], num_tokens=result["num_tokens"], metrics=result["metrics"]
|
||||
)
|
||||
)
|
||||
|
||||
def optim_step(self, adam_params: AdamParams | None = None) -> Future[OptimResult]:
|
||||
adam = adam_params or AdamParams()
|
||||
result = self._service._request(
|
||||
"POST",
|
||||
self._path("/optim_step"),
|
||||
{
|
||||
"adam_params": {
|
||||
"lr": adam.lr,
|
||||
"beta1": adam.beta1,
|
||||
"beta2": adam.beta2,
|
||||
"eps": adam.eps,
|
||||
"weight_decay": adam.weight_decay,
|
||||
}
|
||||
},
|
||||
)
|
||||
return _future(OptimResult(optim_steps=result["optim_steps"]))
|
||||
|
||||
def sample(
|
||||
self,
|
||||
prompt: str | None = None,
|
||||
tokens: list[int] | None = None,
|
||||
sampling_params: SamplingParams | None = None,
|
||||
num_samples: int = 1,
|
||||
) -> Future[SampleResult]:
|
||||
if (prompt is None) == (tokens is None):
|
||||
raise ValueError("pass exactly one of prompt or tokens")
|
||||
sp = sampling_params or SamplingParams()
|
||||
body: dict[str, Any] = {
|
||||
"sampling_params": {
|
||||
"max_tokens": sp.max_tokens,
|
||||
"temperature": sp.temperature,
|
||||
"top_k": sp.top_k,
|
||||
"top_p": sp.top_p,
|
||||
"seed": sp.seed,
|
||||
"stop_tokens": list(sp.stop_tokens),
|
||||
},
|
||||
"num_samples": num_samples,
|
||||
}
|
||||
if prompt is not None:
|
||||
body["prompt"] = prompt
|
||||
else:
|
||||
body["tokens"] = tokens
|
||||
result = self._service._request("POST", self._path("/sample"), body)
|
||||
return _future(
|
||||
SampleResult(
|
||||
sequences=[
|
||||
Sequence(tokens=s["tokens"], text=s["text"]) for s in result["sequences"]
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
def save_weights_and_get_sampling_client(
|
||||
self, name: str, dir: str | None = None
|
||||
) -> Future[SaveResult]:
|
||||
body: dict[str, Any] = {"name": name}
|
||||
if dir is not None:
|
||||
body["dir"] = dir
|
||||
result = self._service._request("POST", self._path("/save_weights"), body)
|
||||
return _future(SaveResult(path=result["path"], format=result["format"]))
|
||||
|
||||
save_weights = save_weights_and_get_sampling_client
|
||||
|
||||
def get_info(self) -> TrainingClientInfo:
|
||||
return _info(self._service._request("GET", self._path()))
|
||||
|
||||
def delete(self) -> None:
|
||||
self._service._request("DELETE", self._path())
|
||||
|
||||
def _wait_ready(self, timeout: float) -> None:
|
||||
deadline = monotonic() + timeout
|
||||
while True:
|
||||
info = self.get_info()
|
||||
if info.status == "ready":
|
||||
return
|
||||
if info.status == "failed":
|
||||
raise HanzoTrainError(409, info.error or "training client failed to load")
|
||||
if monotonic() >= deadline:
|
||||
raise HanzoTrainError(408, f"training client {self.id} not ready after {timeout}s")
|
||||
sleep(self._poll_interval)
|
||||
@@ -0,0 +1,59 @@
|
||||
[project]
|
||||
name = "hanzo-train"
|
||||
version = "0.1.0"
|
||||
description = "Hanzo Train — Tinker-shaped client for the Hanzo Engine training API"
|
||||
readme = "README.md"
|
||||
license = { text = "Apache-2.0" }
|
||||
requires-python = ">=3.12"
|
||||
authors = [{ name = "Hanzo AI", email = "[email protected]" }]
|
||||
keywords = ["hanzo", "training", "lora", "fine-tuning", "tinker", "llm"]
|
||||
classifiers = [
|
||||
"Development Status :: 4 - Beta",
|
||||
"Intended Audience :: Developers",
|
||||
"License :: OSI Approved :: Apache Software License",
|
||||
"Operating System :: OS Independent",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
"Programming Language :: Python :: 3.13",
|
||||
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
||||
"Topic :: Software Development :: Libraries :: Python Modules",
|
||||
"Typing :: Typed",
|
||||
]
|
||||
|
||||
dependencies = [
|
||||
"httpx>=0.25.0",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
"pytest>=8.0.0",
|
||||
"ruff>=0.5.0",
|
||||
"mypy>=1.10.0",
|
||||
]
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://hanzo.ai"
|
||||
Documentation = "https://docs.hanzo.ai/train"
|
||||
Backend = "https://github.com/hanzoai/engine"
|
||||
Source = "https://github.com/hanzoai/python-sdk"
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["hanzo_train"]
|
||||
|
||||
[tool.ruff]
|
||||
target-version = "py312"
|
||||
line-length = 100
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = ["E", "W", "F", "I", "B", "C4", "UP"]
|
||||
ignore = ["E501"]
|
||||
|
||||
[tool.mypy]
|
||||
python_version = "3.12"
|
||||
strict = true
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
testpaths = ["tests"]
|
||||
@@ -0,0 +1,421 @@
|
||||
"""Unit tests for hanzo_train over httpx.MockTransport -- no network."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from hanzo_train import (
|
||||
AdamParams,
|
||||
Datum,
|
||||
HanzoTrainError,
|
||||
LoraConfig,
|
||||
ModelInput,
|
||||
SamplingParams,
|
||||
ServiceClient,
|
||||
TrainingClient,
|
||||
)
|
||||
|
||||
BASE = "http://engine"
|
||||
Handler = Callable[[httpx.Request], httpx.Response]
|
||||
|
||||
|
||||
def service(handler: Handler, api_key: str | None = None) -> ServiceClient:
|
||||
return ServiceClient(BASE, api_key=api_key, transport=httpx.MockTransport(handler))
|
||||
|
||||
|
||||
def sent(request: httpx.Request) -> dict[str, Any]:
|
||||
return json.loads(request.content)
|
||||
|
||||
|
||||
def info_json(**over: Any) -> dict[str, Any]:
|
||||
data: dict[str, Any] = {
|
||||
"id": "tc-abc",
|
||||
"base_model": "m",
|
||||
"status": "ready",
|
||||
"lora_config": {"rank": 16, "alpha": 32.0, "target_modules": ["q_proj"]},
|
||||
"forward_backward_calls": 0,
|
||||
"optim_steps": 0,
|
||||
}
|
||||
data.update(over)
|
||||
return data
|
||||
|
||||
|
||||
def test_create_defaults_and_no_wait() -> None:
|
||||
seen: dict[str, Any] = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
seen["method"] = request.method
|
||||
seen["path"] = request.url.path
|
||||
seen["body"] = sent(request)
|
||||
return httpx.Response(200, json=info_json(status="loading"))
|
||||
|
||||
tc = service(handler).create_lora_training_client("m", wait=False)
|
||||
|
||||
assert tc.id == "tc-abc"
|
||||
assert seen["method"] == "POST"
|
||||
assert seen["path"] == "/v1/training/clients"
|
||||
assert seen["body"]["base_model"] == "m"
|
||||
lora = seen["body"]["lora_config"]
|
||||
assert lora["rank"] == 16
|
||||
assert lora["alpha"] == 32.0
|
||||
assert lora["target_modules"] == [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
"o_proj",
|
||||
"gate_proj",
|
||||
"up_proj",
|
||||
"down_proj",
|
||||
]
|
||||
|
||||
|
||||
def test_create_custom_lora_config() -> None:
|
||||
seen: dict[str, Any] = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
seen["body"] = sent(request)
|
||||
return httpx.Response(200, json=info_json(status="loading"))
|
||||
|
||||
service(handler).create_lora_training_client(
|
||||
"m", lora_config=LoraConfig(rank=8, alpha=16.0, target_modules=["q_proj"]), wait=False
|
||||
)
|
||||
|
||||
assert seen["body"]["lora_config"] == {
|
||||
"rank": 8,
|
||||
"alpha": 16.0,
|
||||
"target_modules": ["q_proj"],
|
||||
}
|
||||
|
||||
|
||||
def test_create_wait_until_ready() -> None:
|
||||
calls = {"get": 0}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
if request.method == "POST":
|
||||
return httpx.Response(200, json=info_json(status="loading"))
|
||||
calls["get"] += 1
|
||||
status = "ready" if calls["get"] >= 2 else "loading"
|
||||
return httpx.Response(200, json=info_json(status=status))
|
||||
|
||||
tc = service(handler).create_lora_training_client("m", wait=True, poll_interval=0.0)
|
||||
|
||||
assert tc.id == "tc-abc"
|
||||
assert calls["get"] == 2
|
||||
|
||||
|
||||
def test_create_wait_failed_raises() -> None:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
if request.method == "POST":
|
||||
return httpx.Response(200, json=info_json(status="loading"))
|
||||
return httpx.Response(200, json=info_json(status="failed", error="oom"))
|
||||
|
||||
with pytest.raises(HanzoTrainError) as excinfo:
|
||||
service(handler).create_lora_training_client("m", wait=True, poll_interval=0.0)
|
||||
|
||||
assert excinfo.value.status == 409
|
||||
assert "oom" in excinfo.value.message
|
||||
|
||||
|
||||
def test_forward_backward_prompt_form() -> None:
|
||||
seen: dict[str, Any] = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
seen["path"] = request.url.path
|
||||
seen["body"] = sent(request)
|
||||
return httpx.Response(200, json={"loss": 1.5, "num_tokens": 3, "metrics": {"loss": 1.5}})
|
||||
|
||||
tc = TrainingClient(service(handler), "tc-abc")
|
||||
out = tc.forward_backward([{"prompt": "2+2=", "completion": "4"}]).result()
|
||||
|
||||
assert seen["path"] == "/v1/training/clients/tc-abc/forward_backward"
|
||||
assert seen["body"]["data"] == [{"prompt": "2+2=", "completion": "4"}]
|
||||
assert out.loss == 1.5
|
||||
assert out.num_tokens == 3
|
||||
assert out.metrics == {"loss": 1.5}
|
||||
|
||||
|
||||
def test_forward_backward_datum_form() -> None:
|
||||
seen: dict[str, Any] = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
seen["body"] = sent(request)
|
||||
return httpx.Response(200, json={"loss": 0.5, "num_tokens": 3, "metrics": {"loss": 0.5}})
|
||||
|
||||
tc = TrainingClient(service(handler), "tc-abc")
|
||||
datum = Datum(
|
||||
model_input=ModelInput(tokens=[1, 2, 3]), target_tokens=[2, 3, 4], weights=[0.0, 1.0, 1.0]
|
||||
)
|
||||
out = tc.forward_backward([datum]).result()
|
||||
|
||||
assert seen["body"]["data"] == [
|
||||
{
|
||||
"model_input": {"tokens": [1, 2, 3]},
|
||||
"target_tokens": [2, 3, 4],
|
||||
"weights": [0.0, 1.0, 1.0],
|
||||
}
|
||||
]
|
||||
assert out.loss == 0.5
|
||||
|
||||
|
||||
def test_forward_backward_mixed_forms() -> None:
|
||||
seen: dict[str, Any] = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
seen["body"] = sent(request)
|
||||
return httpx.Response(200, json={"loss": 1.0, "num_tokens": 6, "metrics": {"loss": 1.0}})
|
||||
|
||||
tc = TrainingClient(service(handler), "tc-abc")
|
||||
tc.forward_backward(
|
||||
[
|
||||
{"prompt": "2+2=", "completion": "4"},
|
||||
Datum(model_input=ModelInput(tokens=[1]), target_tokens=[2], weights=[1.0]),
|
||||
]
|
||||
).result()
|
||||
|
||||
assert seen["body"]["data"] == [
|
||||
{"prompt": "2+2=", "completion": "4"},
|
||||
{"model_input": {"tokens": [1]}, "target_tokens": [2], "weights": [1.0]},
|
||||
]
|
||||
|
||||
|
||||
def test_optim_step_explicit_params() -> None:
|
||||
seen: dict[str, Any] = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
seen["path"] = request.url.path
|
||||
seen["body"] = sent(request)
|
||||
return httpx.Response(200, json={"optim_steps": 1})
|
||||
|
||||
tc = TrainingClient(service(handler), "tc-abc")
|
||||
res = tc.optim_step(AdamParams(lr=1e-4)).result()
|
||||
|
||||
assert seen["path"] == "/v1/training/clients/tc-abc/optim_step"
|
||||
assert seen["body"]["adam_params"] == {
|
||||
"lr": 1e-4,
|
||||
"beta1": 0.9,
|
||||
"beta2": 0.95,
|
||||
"eps": 1e-8,
|
||||
"weight_decay": 0.0,
|
||||
}
|
||||
assert res.optim_steps == 1
|
||||
|
||||
|
||||
def test_optim_step_defaults() -> None:
|
||||
seen: dict[str, Any] = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
seen["body"] = sent(request)
|
||||
return httpx.Response(200, json={"optim_steps": 7})
|
||||
|
||||
tc = TrainingClient(service(handler), "tc-abc")
|
||||
res = tc.optim_step().result()
|
||||
|
||||
assert seen["body"]["adam_params"]["lr"] == 1e-4
|
||||
assert res.optim_steps == 7
|
||||
|
||||
|
||||
def test_sample_prompt_form() -> None:
|
||||
seen: dict[str, Any] = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
seen["path"] = request.url.path
|
||||
seen["body"] = sent(request)
|
||||
return httpx.Response(200, json={"sequences": [{"tokens": [4], "text": "4"}]})
|
||||
|
||||
tc = TrainingClient(service(handler), "tc-abc")
|
||||
resp = tc.sample(
|
||||
prompt="2+2=", sampling_params=SamplingParams(max_tokens=8, temperature=0.0), num_samples=1
|
||||
).result()
|
||||
|
||||
assert seen["path"] == "/v1/training/clients/tc-abc/sample"
|
||||
assert seen["body"]["prompt"] == "2+2="
|
||||
assert "tokens" not in seen["body"]
|
||||
assert seen["body"]["num_samples"] == 1
|
||||
assert seen["body"]["sampling_params"] == {
|
||||
"max_tokens": 8,
|
||||
"temperature": 0.0,
|
||||
"top_k": -1,
|
||||
"top_p": 1.0,
|
||||
"seed": 0,
|
||||
"stop_tokens": [],
|
||||
}
|
||||
assert resp.sequences[0].tokens == [4]
|
||||
assert resp.sequences[0].text == "4"
|
||||
|
||||
|
||||
def test_sample_tokens_form_and_defaults() -> None:
|
||||
seen: dict[str, Any] = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
seen["body"] = sent(request)
|
||||
return httpx.Response(200, json={"sequences": [{"tokens": [9], "text": "x"}]})
|
||||
|
||||
tc = TrainingClient(service(handler), "tc-abc")
|
||||
tc.sample(tokens=[1, 2, 3]).result()
|
||||
|
||||
assert seen["body"]["tokens"] == [1, 2, 3]
|
||||
assert "prompt" not in seen["body"]
|
||||
assert seen["body"]["sampling_params"]["max_tokens"] == 64
|
||||
assert seen["body"]["sampling_params"]["temperature"] == 1.0
|
||||
|
||||
|
||||
def test_sample_requires_exactly_one_input() -> None:
|
||||
tc = TrainingClient(service(lambda r: httpx.Response(200, json={"sequences": []})), "tc-abc")
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
tc.sample()
|
||||
with pytest.raises(ValueError):
|
||||
tc.sample(prompt="a", tokens=[1])
|
||||
|
||||
|
||||
def test_save_weights_and_alias() -> None:
|
||||
seen: dict[str, Any] = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
seen["path"] = request.url.path
|
||||
seen["body"] = sent(request)
|
||||
return httpx.Response(200, json={"path": "/data/adapters/my-adapter", "format": "peft"})
|
||||
|
||||
tc = TrainingClient(service(handler), "tc-abc")
|
||||
saved = tc.save_weights_and_get_sampling_client(name="my-adapter").result()
|
||||
|
||||
assert seen["path"] == "/v1/training/clients/tc-abc/save_weights"
|
||||
assert seen["body"] == {"name": "my-adapter"}
|
||||
assert saved.path == "/data/adapters/my-adapter"
|
||||
assert saved.format == "peft"
|
||||
|
||||
saved2 = tc.save_weights(name="a", dir="/tmp/x").result()
|
||||
assert seen["body"] == {"name": "a", "dir": "/tmp/x"}
|
||||
assert saved2.format == "peft"
|
||||
|
||||
|
||||
def test_get_info() -> None:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert request.method == "GET"
|
||||
assert request.url.path == "/v1/training/clients/tc-abc"
|
||||
return httpx.Response(
|
||||
200,
|
||||
json=info_json(
|
||||
status="ready",
|
||||
forward_backward_calls=3,
|
||||
optim_steps=2,
|
||||
trainable_params=1234,
|
||||
last_loss=0.25,
|
||||
loss_history=[1.0, 0.5, 0.25],
|
||||
),
|
||||
)
|
||||
|
||||
info = TrainingClient(service(handler), "tc-abc").get_info()
|
||||
|
||||
assert info.status == "ready"
|
||||
assert info.forward_backward_calls == 3
|
||||
assert info.optim_steps == 2
|
||||
assert info.trainable_params == 1234
|
||||
assert info.last_loss == 0.25
|
||||
assert info.loss_history == [1.0, 0.5, 0.25]
|
||||
assert info.lora_config.rank == 16
|
||||
assert info.lora_config.alpha == 32.0
|
||||
|
||||
|
||||
def test_list_training_clients() -> None:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert request.method == "GET"
|
||||
assert request.url.path == "/v1/training/clients"
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"clients": [
|
||||
info_json(id="tc-1", status="ready"),
|
||||
info_json(id="tc-2", status="loading"),
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
infos = service(handler).list_training_clients()
|
||||
|
||||
assert [i.id for i in infos] == ["tc-1", "tc-2"]
|
||||
assert infos[0].status == "ready"
|
||||
assert isinstance(infos[0].lora_config, LoraConfig)
|
||||
|
||||
|
||||
def test_delete() -> None:
|
||||
seen: dict[str, Any] = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
seen["method"] = request.method
|
||||
seen["path"] = request.url.path
|
||||
return httpx.Response(200, json={"id": "tc-abc", "deleted": True})
|
||||
|
||||
result = TrainingClient(service(handler), "tc-abc").delete()
|
||||
|
||||
assert result is None
|
||||
assert seen["method"] == "DELETE"
|
||||
assert seen["path"] == "/v1/training/clients/tc-abc"
|
||||
|
||||
|
||||
def test_error_404() -> None:
|
||||
tc = TrainingClient(service(lambda r: httpx.Response(404, text="unknown id")), "tc-missing")
|
||||
|
||||
with pytest.raises(HanzoTrainError) as excinfo:
|
||||
tc.get_info()
|
||||
|
||||
assert excinfo.value.status == 404
|
||||
assert excinfo.value.message == "unknown id"
|
||||
|
||||
|
||||
def test_error_409_forward_backward() -> None:
|
||||
tc = TrainingClient(
|
||||
service(lambda r: httpx.Response(409, text="client still loading")), "tc-abc"
|
||||
)
|
||||
|
||||
with pytest.raises(HanzoTrainError) as excinfo:
|
||||
tc.forward_backward([{"prompt": "a", "completion": "b"}]).result()
|
||||
|
||||
assert excinfo.value.status == 409
|
||||
assert "loading" in excinfo.value.message
|
||||
|
||||
|
||||
def test_error_400_optim_step() -> None:
|
||||
tc = TrainingClient(service(lambda r: httpx.Response(400, text="bad data")), "tc-abc")
|
||||
|
||||
with pytest.raises(HanzoTrainError) as excinfo:
|
||||
tc.optim_step().result()
|
||||
|
||||
assert excinfo.value.status == 400
|
||||
|
||||
|
||||
def test_auth_header_bearer() -> None:
|
||||
seen: dict[str, Any] = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
seen["auth"] = request.headers.get("authorization")
|
||||
return httpx.Response(200, json=info_json())
|
||||
|
||||
TrainingClient(service(handler, api_key="secret-token"), "tc-abc").get_info()
|
||||
|
||||
assert seen["auth"] == "Bearer secret-token"
|
||||
|
||||
|
||||
def test_no_auth_header_without_key() -> None:
|
||||
seen: dict[str, Any] = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
seen["auth"] = request.headers.get("authorization")
|
||||
return httpx.Response(200, json=info_json())
|
||||
|
||||
TrainingClient(service(handler), "tc-abc").get_info()
|
||||
|
||||
assert seen["auth"] is None
|
||||
|
||||
|
||||
def test_future_result_returns_value() -> None:
|
||||
tc = TrainingClient(service(lambda r: httpx.Response(200, json={"optim_steps": 5})), "tc-abc")
|
||||
|
||||
fut = tc.optim_step()
|
||||
|
||||
assert fut.result(timeout=1.0).optim_steps == 5
|
||||
@@ -72,6 +72,7 @@ members = [
|
||||
"pkg/hanzo-network",
|
||||
"pkg/hanzo-dev-py",
|
||||
"pkg/hanzo-s3",
|
||||
"pkg/hanzo-train",
|
||||
"pkg/hanzo-tools-ui",
|
||||
"pkg/hanzo-zap",
|
||||
"pkg/hanzoai",
|
||||
@@ -209,6 +210,7 @@ exclude = [
|
||||
"pkg/hanzo-network",
|
||||
"pkg/hanzo-node",
|
||||
"pkg/hanzo-dev-py",
|
||||
"pkg/hanzo-train",
|
||||
"pkg/hanzo-tools",
|
||||
"pkg/hanzo-tools-agent",
|
||||
"pkg/hanzo-tools-api",
|
||||
@@ -482,6 +484,7 @@ exclude = [
|
||||
members = [
|
||||
"pkg/hanzo-memory",
|
||||
"pkg/hanzo-s3",
|
||||
"pkg/hanzo-train",
|
||||
"pkg/hanzo-tools",
|
||||
"pkg/hanzo-tools-core",
|
||||
# Required tool packages — the hanzo-tools-core integration suite asserts each
|
||||
@@ -533,6 +536,7 @@ hanzo-cli = { workspace = true }
|
||||
hanzo-iam = { workspace = true }
|
||||
hanzo-kms = { workspace = true }
|
||||
hanzo-s3 = { workspace = true }
|
||||
hanzo-train = { workspace = true }
|
||||
hanzo-tools-s3 = { workspace = true }
|
||||
|
||||
[tool.ruff.lint.extend-per-file-ignores]
|
||||
|
||||
@@ -28,6 +28,7 @@ members = [
|
||||
"hanzo-tools-shell",
|
||||
"hanzo-tools-todo",
|
||||
"hanzo-tools-ui",
|
||||
"hanzo-train",
|
||||
"hanzoai",
|
||||
]
|
||||
|
||||
@@ -1766,7 +1767,7 @@ provides-extras = ["consensus", "api", "full", "dev"]
|
||||
|
||||
[[package]]
|
||||
name = "hanzo-tools-browser"
|
||||
version = "0.5.9"
|
||||
version = "0.5.10"
|
||||
source = { editable = "pkg/hanzo-tools-browser" }
|
||||
dependencies = [
|
||||
{ name = "hanzo-tools" },
|
||||
@@ -2210,9 +2211,46 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/b3/ee/a6d265860d51a52a8ac931579f722a15173ca480bbcce233fd78505b491b/hanzo_tools_vector-0.2.0-py3-none-any.whl", hash = "sha256:2c5dca66e657bf1098d9ac4f8aff625addf4cb623387a1d2cb234a935ab4e7fc", size = 32480, upload-time = "2025-12-26T18:29:58.591Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hanzo-train"
|
||||
version = "0.1.0"
|
||||
source = { editable = "pkg/hanzo-train" }
|
||||
dependencies = [
|
||||
{ name = "httpx" },
|
||||
]
|
||||
|
||||
[package.optional-dependencies]
|
||||
dev = [
|
||||
{ name = "mypy" },
|
||||
{ name = "pytest" },
|
||||
{ name = "ruff" },
|
||||
]
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "httpx", specifier = ">=0.25.0" },
|
||||
{ name = "mypy", marker = "extra == 'dev'", specifier = ">=1.10.0" },
|
||||
{ name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0.0" },
|
||||
{ name = "ruff", marker = "extra == 'dev'", specifier = ">=0.5.0" },
|
||||
]
|
||||
provides-extras = ["dev"]
|
||||
|
||||
[[package]]
|
||||
name = "hanzo-zap"
|
||||
version = "0.7.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "httpx" },
|
||||
{ name = "zap-proto" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/72/fc/98b9fcbc0b104a25cc019d1ac5fe07270e16d4e8ea454a97ad86ab0ff903/hanzo_zap-0.7.0.tar.gz", hash = "sha256:74c5ce45b69c3fb7cb72a095f83434c0fa7706484e2fc2d96417ebc4300821e0", size = 19759, upload-time = "2026-06-30T20:57:16.522Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/24/b5/70a04edf50d80fc13c090da454bce7405e5563c0b640f6eb25a7ae20d51c/hanzo_zap-0.7.0-py3-none-any.whl", hash = "sha256:7590fc2785db10c2b727a14b8c8556c33e30c4d06af5bbb95cb43748395790fc", size = 14490, upload-time = "2026-06-30T20:57:15.629Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hanzoai"
|
||||
version = "3.0.1"
|
||||
version = "3.1.1"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "pydantic" },
|
||||
@@ -2225,6 +2263,10 @@ dependencies = [
|
||||
llm = [
|
||||
{ name = "hanzo-llm" },
|
||||
]
|
||||
zap = [
|
||||
{ name = "hanzo-zap" },
|
||||
{ name = "httpx" },
|
||||
]
|
||||
|
||||
[package.dev-dependencies]
|
||||
dev = [
|
||||
@@ -2251,12 +2293,14 @@ dev = [
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "hanzo-llm", marker = "extra == 'llm'", specifier = ">=1.0.0" },
|
||||
{ name = "hanzo-zap", marker = "extra == 'zap'", specifier = ">=0.7.0" },
|
||||
{ name = "httpx", marker = "extra == 'zap'", specifier = ">=0.27.0" },
|
||||
{ name = "pydantic", specifier = ">=2.10,<3" },
|
||||
{ name = "python-dateutil", specifier = ">=2.8.2" },
|
||||
{ name = "typing-extensions", specifier = ">=4.10,<5" },
|
||||
{ name = "urllib3", specifier = ">=2.1.0,<3" },
|
||||
]
|
||||
provides-extras = ["llm"]
|
||||
provides-extras = ["llm", "zap"]
|
||||
|
||||
[package.metadata.requires-dev]
|
||||
dev = [
|
||||
|
||||
Reference in New Issue
Block a user