ariapin 1.7.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- ariapin-1.7.0/PKG-INFO +15 -0
- ariapin-1.7.0/ariapin/__init__.py +29 -0
- ariapin-1.7.0/ariapin/_http.py +169 -0
- ariapin-1.7.0/ariapin/client.py +282 -0
- ariapin-1.7.0/ariapin/exceptions.py +25 -0
- ariapin-1.7.0/ariapin/models.py +334 -0
- ariapin-1.7.0/ariapin.egg-info/PKG-INFO +15 -0
- ariapin-1.7.0/ariapin.egg-info/SOURCES.txt +12 -0
- ariapin-1.7.0/ariapin.egg-info/dependency_links.txt +1 -0
- ariapin-1.7.0/ariapin.egg-info/requires.txt +3 -0
- ariapin-1.7.0/ariapin.egg-info/top_level.txt +1 -0
- ariapin-1.7.0/pyproject.toml +29 -0
- ariapin-1.7.0/setup.cfg +4 -0
- ariapin-1.7.0/tests/test_client.py +252 -0
ariapin-1.7.0/PKG-INFO
ADDED
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: ariapin
|
|
3
|
+
Version: 1.7.0
|
|
4
|
+
Summary: Python client for the ariapin end-side post-training API (LoRA/QLoRA + RL with SFT/DPO/GRPO)
|
|
5
|
+
Author: aria compute
|
|
6
|
+
License: Apache-2.0
|
|
7
|
+
Project-URL: Homepage, https://github.com/ariacompute/harness
|
|
8
|
+
Keywords: llm,post-training,lora,qlora,dora,grpo,dpo,sft,rl
|
|
9
|
+
Classifier: Programming Language :: Python :: 3
|
|
10
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
11
|
+
Classifier: License :: OSI Approved :: Apache Software License
|
|
12
|
+
Requires-Python: >=3.10
|
|
13
|
+
Description-Content-Type: text/markdown
|
|
14
|
+
Provides-Extra: dev
|
|
15
|
+
Requires-Dist: pytest>=7; extra == "dev"
|
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
"""ariapin — Python client for the end-side post-training API."""
|
|
2
|
+
|
|
3
|
+
from .client import AriaPinClient, DEFAULT_BASE_URL
|
|
4
|
+
from .exceptions import AriaPinConnectionError, AriaPinError
|
|
5
|
+
from .models import (
|
|
6
|
+
Agent,
|
|
7
|
+
AutoIntervention,
|
|
8
|
+
AutoInterventionRule,
|
|
9
|
+
GPUNode,
|
|
10
|
+
Hyperparams,
|
|
11
|
+
Job,
|
|
12
|
+
LoraConfig,
|
|
13
|
+
)
|
|
14
|
+
|
|
15
|
+
__all__ = [
|
|
16
|
+
"AriaPinClient",
|
|
17
|
+
"DEFAULT_BASE_URL",
|
|
18
|
+
"AriaPinError",
|
|
19
|
+
"AriaPinConnectionError",
|
|
20
|
+
"Job",
|
|
21
|
+
"LoraConfig",
|
|
22
|
+
"Hyperparams",
|
|
23
|
+
"AutoIntervention",
|
|
24
|
+
"AutoInterventionRule",
|
|
25
|
+
"GPUNode",
|
|
26
|
+
"Agent",
|
|
27
|
+
]
|
|
28
|
+
|
|
29
|
+
__version__ = "1.7.0"
|
|
@@ -0,0 +1,169 @@
|
|
|
1
|
+
"""Internal HTTP transport for the ariapin client.
|
|
2
|
+
|
|
3
|
+
Zero runtime dependencies: uses only the standard library. The envelope is
|
|
4
|
+
the unified ``{"code": int, "data": ..., "message": str}`` shape returned by
|
|
5
|
+
the backend. A non-zero ``code`` (or any non-2xx HTTP status) raises
|
|
6
|
+
``AriaPinError``. Only 5xx / network errors are retried with exponential
|
|
7
|
+
backoff; 4xx responses are surfaced immediately (retrying would waste quota).
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
from __future__ import annotations
|
|
11
|
+
|
|
12
|
+
import json
|
|
13
|
+
import time
|
|
14
|
+
import urllib.error
|
|
15
|
+
import urllib.request
|
|
16
|
+
from typing import Any
|
|
17
|
+
|
|
18
|
+
from .exceptions import AriaPinConnectionError, AriaPinError
|
|
19
|
+
|
|
20
|
+
_DEFAULT_TIMEOUT = 30
|
|
21
|
+
_MAX_RETRIES = 3
|
|
22
|
+
_BACKOFF_BASE = 0.5
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class _HTTPClient:
|
|
26
|
+
def __init__(self, api_key: str, base_url: str, timeout: int = _DEFAULT_TIMEOUT):
|
|
27
|
+
self._api_key = api_key
|
|
28
|
+
self._base_url = base_url.rstrip("/")
|
|
29
|
+
self._timeout = timeout
|
|
30
|
+
|
|
31
|
+
def _url(self, path: str) -> str:
|
|
32
|
+
if not path.startswith("/"):
|
|
33
|
+
path = "/" + path
|
|
34
|
+
return self._base_url + path
|
|
35
|
+
|
|
36
|
+
def request(
|
|
37
|
+
self,
|
|
38
|
+
method: str,
|
|
39
|
+
path: str,
|
|
40
|
+
*,
|
|
41
|
+
json_body: Any = None,
|
|
42
|
+
query: dict[str, Any] | None = None,
|
|
43
|
+
files: dict[str, tuple[str, bytes, str]] | None = None,
|
|
44
|
+
) -> Any:
|
|
45
|
+
url = self._url(path)
|
|
46
|
+
if query:
|
|
47
|
+
from urllib.parse import urlencode
|
|
48
|
+
|
|
49
|
+
url += "?" + urlencode({k: v for k, v in query.items() if v is not None})
|
|
50
|
+
|
|
51
|
+
last_err: Exception | None = None
|
|
52
|
+
for attempt in range(1, _MAX_RETRIES + 1):
|
|
53
|
+
try:
|
|
54
|
+
data, headers = self._prepare(json_body, files)
|
|
55
|
+
req = urllib.request.Request(url, data=data, method=method, headers=headers)
|
|
56
|
+
try:
|
|
57
|
+
with urllib.request.urlopen(req, timeout=self._timeout) as resp:
|
|
58
|
+
return self._parse(resp)
|
|
59
|
+
except urllib.error.HTTPError as e:
|
|
60
|
+
# Server responded with an error status.
|
|
61
|
+
body = self._read_bytes(e)
|
|
62
|
+
err = self._into_error(e.code, body)
|
|
63
|
+
if 400 <= e.code < 500:
|
|
64
|
+
# Do not retry client errors, but allow 429 to retry.
|
|
65
|
+
if e.code == 429:
|
|
66
|
+
last_err = err
|
|
67
|
+
time.sleep(_BACKOFF_BASE * attempt)
|
|
68
|
+
continue
|
|
69
|
+
raise err
|
|
70
|
+
last_err = err # 5xx -> retry
|
|
71
|
+
except urllib.error.URLError as e:
|
|
72
|
+
raise AriaPinConnectionError(
|
|
73
|
+
f"cannot reach {url}: {e.reason}"
|
|
74
|
+
) from e
|
|
75
|
+
except AriaPinConnectionError:
|
|
76
|
+
last_err = err if (err := last_err) else last_err # keep last
|
|
77
|
+
last_err = _wrap_conn_err(url)
|
|
78
|
+
if attempt < _MAX_RETRIES:
|
|
79
|
+
time.sleep(_BACKOFF_BASE * attempt)
|
|
80
|
+
continue
|
|
81
|
+
raise last_err
|
|
82
|
+
# successful return or 5xx retry
|
|
83
|
+
if last_err is not None and isinstance(last_err, AriaPinError) and last_err.status >= 500:
|
|
84
|
+
if attempt < _MAX_RETRIES:
|
|
85
|
+
time.sleep(_BACKOFF_BASE * attempt)
|
|
86
|
+
continue
|
|
87
|
+
raise last_err
|
|
88
|
+
if last_err is not None:
|
|
89
|
+
raise last_err
|
|
90
|
+
raise AriaPinError("unknown error")
|
|
91
|
+
|
|
92
|
+
# ---- helpers -------------------------------------------------------
|
|
93
|
+
def _prepare(self, json_body, files):
|
|
94
|
+
headers = {"Authorization": f"Bearer {self._api_key}"}
|
|
95
|
+
if json_body is not None:
|
|
96
|
+
headers["Content-Type"] = "application/json"
|
|
97
|
+
return json.dumps(json_body).encode("utf-8"), headers
|
|
98
|
+
if files is not None:
|
|
99
|
+
boundary = "----ariapin%s" % int(time.time() * 1000)
|
|
100
|
+
body = _encode_multipart(boundary, files)
|
|
101
|
+
headers["Content-Type"] = f"multipart/form-data; boundary={boundary}"
|
|
102
|
+
return body, headers
|
|
103
|
+
return None, headers
|
|
104
|
+
|
|
105
|
+
@staticmethod
|
|
106
|
+
def _read_bytes(e: urllib.error.HTTPError) -> bytes:
|
|
107
|
+
try:
|
|
108
|
+
return e.read()
|
|
109
|
+
except Exception: # pragma: no cover - defensive
|
|
110
|
+
return b""
|
|
111
|
+
|
|
112
|
+
def _into_error(self, status: int, raw: bytes) -> AriaPinError:
|
|
113
|
+
msg = ""
|
|
114
|
+
code = 0
|
|
115
|
+
data = None
|
|
116
|
+
if raw:
|
|
117
|
+
try:
|
|
118
|
+
payload = json.loads(raw.decode("utf-8", "replace"))
|
|
119
|
+
code = int(payload.get("code", 0))
|
|
120
|
+
msg = str(payload.get("message", ""))
|
|
121
|
+
data = payload.get("data")
|
|
122
|
+
except Exception:
|
|
123
|
+
msg = raw.decode("utf-8", "replace")
|
|
124
|
+
if not msg:
|
|
125
|
+
msg = f"HTTP {status}"
|
|
126
|
+
return AriaPinError(msg, status=status, code=code, data=data)
|
|
127
|
+
|
|
128
|
+
@staticmethod
|
|
129
|
+
def _parse(resp) -> Any:
|
|
130
|
+
raw = resp.read()
|
|
131
|
+
if not raw:
|
|
132
|
+
return None
|
|
133
|
+
try:
|
|
134
|
+
payload = json.loads(raw.decode("utf-8", "replace"))
|
|
135
|
+
except Exception:
|
|
136
|
+
return raw # unexpected binary (e.g. download stream) -> caller handles
|
|
137
|
+
if isinstance(payload, dict) and "code" in payload:
|
|
138
|
+
if payload.get("code") not in (0, None):
|
|
139
|
+
raise AriaPinError(
|
|
140
|
+
str(payload.get("message", "")),
|
|
141
|
+
status=200,
|
|
142
|
+
code=int(payload.get("code", 0)),
|
|
143
|
+
data=payload.get("data"),
|
|
144
|
+
)
|
|
145
|
+
return payload.get("data")
|
|
146
|
+
return payload
|
|
147
|
+
|
|
148
|
+
|
|
149
|
+
def _wrap_conn_err(url: str) -> AriaPinConnectionError:
|
|
150
|
+
return AriaPinConnectionError(f"cannot reach {url}")
|
|
151
|
+
|
|
152
|
+
|
|
153
|
+
def _encode_multipart(boundary: str, files: dict[str, tuple[str, bytes, str]]) -> bytes:
|
|
154
|
+
"""Minimal multipart encoder for the dataset upload endpoint.
|
|
155
|
+
|
|
156
|
+
``files`` maps field name -> (filename, content_bytes, content_type).
|
|
157
|
+
"""
|
|
158
|
+
crlf = b"\r\n"
|
|
159
|
+
parts = []
|
|
160
|
+
for name, (filename, content, ctype) in files.items():
|
|
161
|
+
parts.append(f"--{boundary}".encode())
|
|
162
|
+
disp = f'Content-Disposition: form-data; name="{name}"; filename="{filename}"'
|
|
163
|
+
parts.append(disp.encode())
|
|
164
|
+
parts.append(f"Content-Type: {ctype}".encode())
|
|
165
|
+
parts.append(b"")
|
|
166
|
+
parts.append(content if isinstance(content, bytes) else content.encode("utf-8"))
|
|
167
|
+
parts.append(f"--{boundary}--".encode())
|
|
168
|
+
parts.append(b"")
|
|
169
|
+
return crlf.join(parts)
|
|
@@ -0,0 +1,282 @@
|
|
|
1
|
+
"""ariapin Python client.
|
|
2
|
+
|
|
3
|
+
Drop-in training experience inspired by river-client: submit a LoRA/QLoRA +
|
|
4
|
+
RL job and poll it. Pure standard library, no third-party runtime deps.
|
|
5
|
+
|
|
6
|
+
Example
|
|
7
|
+
-------
|
|
8
|
+
from ariapin import AriaPinClient, LoraConfig, Hyperparams
|
|
9
|
+
|
|
10
|
+
client = AriaPinClient(api_key="sk-...", base_url="http://localhost:8001")
|
|
11
|
+
ds = client.create_dataset("my-sft", "sft", "data.jsonl")
|
|
12
|
+
job = client.train(
|
|
13
|
+
base_model="Qwen/Qwen2.5-0.5B-Instruct",
|
|
14
|
+
job_type="grpo",
|
|
15
|
+
lora=LoraConfig(method="lora", rank=16),
|
|
16
|
+
dataset_id=ds["dataset_id"],
|
|
17
|
+
hyperparams=Hyperparams(epochs=1, lr=4e-5, group_size=4, beta=0.04),
|
|
18
|
+
)
|
|
19
|
+
job.wait() # blocks until terminal
|
|
20
|
+
print(job.status, job.progress) # -> "succeeded" {...}
|
|
21
|
+
model = client.models()[-1]
|
|
22
|
+
client.export_model(model["model_id"], "gguf")
|
|
23
|
+
"""
|
|
24
|
+
|
|
25
|
+
from __future__ import annotations
|
|
26
|
+
|
|
27
|
+
from typing import Any
|
|
28
|
+
|
|
29
|
+
from ._http import _HTTPClient
|
|
30
|
+
from .exceptions import AriaPinError
|
|
31
|
+
from .models import AutoIntervention, Hyperparams, Job, LoraConfig
|
|
32
|
+
|
|
33
|
+
DEFAULT_BASE_URL = "http://localhost:8001"
|
|
34
|
+
|
|
35
|
+
# 支持的训练方法(对齐 TRL:SFT/DPO/KTO/ORPO/SimPO + RL GRPO/PPO)。
|
|
36
|
+
JOB_TYPES = ("sft", "dpo", "kto", "orpo", "simpo", "grpo", "ppo")
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
class AriaPinClient:
|
|
40
|
+
def __init__(self, api_key: str, base_url: str = DEFAULT_BASE_URL, timeout: int = 30):
|
|
41
|
+
if not api_key:
|
|
42
|
+
raise ValueError("api_key is required")
|
|
43
|
+
self._http = _HTTPClient(api_key, base_url, timeout=timeout)
|
|
44
|
+
|
|
45
|
+
# --- capabilities --------------------------------------------------
|
|
46
|
+
def capabilities(self) -> dict[str, Any]:
|
|
47
|
+
"""``GET /v1/capabilities`` — model/method/LoRA/quota catalog."""
|
|
48
|
+
return self._http.request("GET", "/v1/capabilities") or {}
|
|
49
|
+
|
|
50
|
+
# --- datasets ------------------------------------------------------
|
|
51
|
+
def create_dataset(self, name: str, fmt: str, file_path: str) -> dict[str, Any]:
|
|
52
|
+
"""``POST /v1/datasets`` multipart upload. Returns dataset metadata."""
|
|
53
|
+
with open(file_path, "rb") as f:
|
|
54
|
+
content = f.read()
|
|
55
|
+
files = {"name": (name, name.encode("utf-8"), "text/plain"),
|
|
56
|
+
"format": (fmt, fmt.encode("utf-8"), "text/plain"),
|
|
57
|
+
"file": (name, content, "application/octet-stream")}
|
|
58
|
+
return self._http.request("POST", "/v1/datasets", files=files) or {}
|
|
59
|
+
|
|
60
|
+
def list_datasets(self) -> list[dict[str, Any]]:
|
|
61
|
+
return self._http.request("GET", "/v1/datasets") or []
|
|
62
|
+
|
|
63
|
+
def get_dataset(self, dataset_id: str) -> dict[str, Any]:
|
|
64
|
+
return self._http.request("GET", f"/v1/datasets/{dataset_id}") or {}
|
|
65
|
+
|
|
66
|
+
# --- jobs ----------------------------------------------------------
|
|
67
|
+
def train(
|
|
68
|
+
self,
|
|
69
|
+
base_model: str,
|
|
70
|
+
job_type: str = "sft",
|
|
71
|
+
*,
|
|
72
|
+
lora: LoraConfig | dict[str, Any] | None = None,
|
|
73
|
+
dataset_id: str | None = None,
|
|
74
|
+
hyperparams: Hyperparams | dict[str, Any] | None = None,
|
|
75
|
+
auto_intervention: AutoIntervention | dict[str, Any] | None = None,
|
|
76
|
+
) -> Job:
|
|
77
|
+
"""``POST /v1/jobs`` — submit an async training job, return a :class:`Job`.
|
|
78
|
+
|
|
79
|
+
Args:
|
|
80
|
+
base_model: HuggingFace base model id (must be in allow-list).
|
|
81
|
+
job_type: sft/dpo/kto/orpo/simpo/grpo/ppo.
|
|
82
|
+
lora_config: LoRA/QLoRA/DoRA config (defaults to rank=16, alpha=32).
|
|
83
|
+
dataset_id: dataset id from ``create_dataset``.
|
|
84
|
+
hyperparams: training hyper-params.
|
|
85
|
+
auto_intervention: optional auto-intervention rules.
|
|
86
|
+
"""
|
|
87
|
+
if job_type not in JOB_TYPES:
|
|
88
|
+
raise ValueError(f"job_type must be one of {JOB_TYPES}")
|
|
89
|
+
body: dict[str, Any] = {"job_type": job_type, "base_model": base_model}
|
|
90
|
+
if lora is not None:
|
|
91
|
+
body["lora_config"] = (
|
|
92
|
+
lora.to_dict() if isinstance(lora, LoraConfig) else lora
|
|
93
|
+
)
|
|
94
|
+
if dataset_id is not None:
|
|
95
|
+
body["dataset_id"] = dataset_id
|
|
96
|
+
if hyperparams is not None:
|
|
97
|
+
body["hyperparams"] = (
|
|
98
|
+
hyperparams.to_dict() if isinstance(hyperparams, Hyperparams) else hyperparams
|
|
99
|
+
)
|
|
100
|
+
if auto_intervention is not None:
|
|
101
|
+
body["auto_intervention"] = (
|
|
102
|
+
auto_intervention.to_dict()
|
|
103
|
+
if isinstance(auto_intervention, AutoIntervention)
|
|
104
|
+
else auto_intervention
|
|
105
|
+
)
|
|
106
|
+
data = self._http.request("POST", "/v1/jobs", json_body=body) or {}
|
|
107
|
+
return Job(self, data.get("job_id", ""), data)
|
|
108
|
+
|
|
109
|
+
def list_jobs(self) -> list[dict[str, Any]]:
|
|
110
|
+
return self._http.request("GET", "/v1/jobs") or []
|
|
111
|
+
|
|
112
|
+
def get_job(self, job_id: str) -> Job:
|
|
113
|
+
return Job(self, job_id, self.get_job_raw(job_id))
|
|
114
|
+
|
|
115
|
+
def get_job_raw(self, job_id: str) -> dict[str, Any]:
|
|
116
|
+
return self._http.request("GET", f"/v1/jobs/{job_id}") or {}
|
|
117
|
+
|
|
118
|
+
def cancel_job(self, job_id: str) -> dict[str, Any]:
|
|
119
|
+
return self._http.request("POST", f"/v1/jobs/{job_id}/cancel") or {}
|
|
120
|
+
|
|
121
|
+
# --- job lifecycle (start/pause/resume/stop/delete) -----------------
|
|
122
|
+
def start_job(self, job_id: str) -> dict[str, Any]:
|
|
123
|
+
return self._http.request("POST", f"/v1/jobs/{job_id}/start") or {}
|
|
124
|
+
|
|
125
|
+
def pause_job(self, job_id: str) -> dict[str, Any]:
|
|
126
|
+
return self._http.request("POST", f"/v1/jobs/{job_id}/pause") or {}
|
|
127
|
+
|
|
128
|
+
def resume_job(self, job_id: str) -> dict[str, Any]:
|
|
129
|
+
return self._http.request("POST", f"/v1/jobs/{job_id}/resume") or {}
|
|
130
|
+
|
|
131
|
+
def stop_job(self, job_id: str) -> dict[str, Any]:
|
|
132
|
+
return self._http.request("POST", f"/v1/jobs/{job_id}/stop") or {}
|
|
133
|
+
|
|
134
|
+
def delete_job(self, job_id: str) -> dict[str, Any]:
|
|
135
|
+
return self._http.request("DELETE", f"/v1/jobs/{job_id}") or {}
|
|
136
|
+
|
|
137
|
+
# --- GPU nodes ------------------------------------------------------
|
|
138
|
+
def list_nodes(self) -> list[dict[str, Any]]:
|
|
139
|
+
data = self._http.request("GET", "/v1/gpu-nodes") or {}
|
|
140
|
+
return data.get("items", []) if isinstance(data, dict) else data
|
|
141
|
+
|
|
142
|
+
def get_node(self, node_id: str) -> dict[str, Any]:
|
|
143
|
+
return self._http.request("GET", f"/v1/gpu-nodes/{node_id}") or {}
|
|
144
|
+
|
|
145
|
+
def create_node(self, *, name: str, ip: str, ssh_user: str, ssh_port: int = 22,
|
|
146
|
+
ssh_key: str | None = None, ssh_password: str | None = None,
|
|
147
|
+
spec: str = "", labels: list[str] | None = None) -> dict[str, Any]:
|
|
148
|
+
body: dict[str, Any] = {
|
|
149
|
+
"name": name, "ip": ip, "ssh_user": ssh_user, "ssh_port": ssh_port,
|
|
150
|
+
"spec": spec, "labels": labels or [],
|
|
151
|
+
}
|
|
152
|
+
if ssh_key is not None:
|
|
153
|
+
body["ssh_key"] = ssh_key
|
|
154
|
+
if ssh_password is not None:
|
|
155
|
+
body["ssh_password"] = ssh_password
|
|
156
|
+
return self._http.request("POST", "/v1/gpu-nodes", json_body=body) or {}
|
|
157
|
+
|
|
158
|
+
def update_node(self, node_id: str, **fields) -> dict[str, Any]:
|
|
159
|
+
"""Partial update. Supported kwargs: name, ip, ssh_port, ssh_user,
|
|
160
|
+
ssh_key, ssh_password, spec, labels, status."""
|
|
161
|
+
body = {k: v for k, v in fields.items() if v is not None}
|
|
162
|
+
return self._http.request("PUT", f"/v1/gpu-nodes/{node_id}", json_body=body) or {}
|
|
163
|
+
|
|
164
|
+
def delete_node(self, node_id: str) -> dict[str, Any]:
|
|
165
|
+
return self._http.request("DELETE", f"/v1/gpu-nodes/{node_id}") or {}
|
|
166
|
+
|
|
167
|
+
# --- agents ---------------------------------------------------------
|
|
168
|
+
def list_node_agents(self, node_id: str) -> list[dict[str, Any]]:
|
|
169
|
+
data = self._http.request("GET", f"/v1/gpu-nodes/{node_id}/agents") or {}
|
|
170
|
+
return data.get("items", []) if isinstance(data, dict) else data
|
|
171
|
+
|
|
172
|
+
def deploy_agent(self, node_id: str, name: str | None = None) -> dict[str, Any]:
|
|
173
|
+
body = {"name": name} if name else {}
|
|
174
|
+
return self._http.request("POST", f"/v1/gpu-nodes/{node_id}/agents", json_body=body) or {}
|
|
175
|
+
|
|
176
|
+
def get_agent(self, agent_id: str) -> dict[str, Any]:
|
|
177
|
+
return self._http.request("GET", f"/v1/agents/{agent_id}") or {}
|
|
178
|
+
|
|
179
|
+
def update_agent(self, agent_id: str) -> dict[str, Any]:
|
|
180
|
+
return self._http.request("POST", f"/v1/agents/{agent_id}/update") or {}
|
|
181
|
+
|
|
182
|
+
def restart_agent(self, agent_id: str) -> dict[str, Any]:
|
|
183
|
+
return self._http.request("POST", f"/v1/agents/{agent_id}/restart") or {}
|
|
184
|
+
|
|
185
|
+
def delete_agent(self, agent_id: str) -> dict[str, Any]:
|
|
186
|
+
return self._http.request("DELETE", f"/v1/agents/{agent_id}") or {}
|
|
187
|
+
|
|
188
|
+
# --- interventions -------------------------------------------------
|
|
189
|
+
def intervene(
|
|
190
|
+
self,
|
|
191
|
+
job_id: str,
|
|
192
|
+
action: str,
|
|
193
|
+
reason: str = "",
|
|
194
|
+
hyperparams: dict[str, Any] | None = None,
|
|
195
|
+
) -> dict[str, Any]:
|
|
196
|
+
body: dict[str, Any] = {"action": action}
|
|
197
|
+
if reason:
|
|
198
|
+
body["reason"] = reason
|
|
199
|
+
if hyperparams is not None:
|
|
200
|
+
body["hyperparams"] = hyperparams
|
|
201
|
+
return self._http.request("POST", f"/v1/jobs/{job_id}/intervene", json_body=body) or {}
|
|
202
|
+
|
|
203
|
+
def set_auto_intervention(
|
|
204
|
+
self, job_id: str, ai: AutoIntervention | dict[str, Any]
|
|
205
|
+
) -> dict[str, Any]:
|
|
206
|
+
body = ai.to_dict() if isinstance(ai, AutoIntervention) else ai
|
|
207
|
+
return self._http.request("PUT", f"/v1/jobs/{job_id}/auto-intervention", json_body=body) or {}
|
|
208
|
+
|
|
209
|
+
def job_interventions(self, job_id: str) -> list[dict[str, Any]]:
|
|
210
|
+
data = self._http.request("GET", f"/v1/jobs/{job_id}/interventions") or {}
|
|
211
|
+
if isinstance(data, dict):
|
|
212
|
+
return data.get("items", [])
|
|
213
|
+
return data
|
|
214
|
+
|
|
215
|
+
def job_analysis(self, job_id: str) -> dict[str, Any]:
|
|
216
|
+
return self._http.request("GET", f"/v1/jobs/{job_id}/analysis") or {}
|
|
217
|
+
|
|
218
|
+
# --- models --------------------------------------------------------
|
|
219
|
+
def models(self) -> list[dict[str, Any]]:
|
|
220
|
+
data = self._http.request("GET", "/v1/models") or {}
|
|
221
|
+
return data.get("items", []) if isinstance(data, dict) else data
|
|
222
|
+
|
|
223
|
+
def get_model(self, model_id: str) -> dict[str, Any]:
|
|
224
|
+
return self._http.request("GET", f"/v1/models/{model_id}") or {}
|
|
225
|
+
|
|
226
|
+
def merge_model(self, model_id: str) -> dict[str, Any]:
|
|
227
|
+
return self._http.request("POST", f"/v1/models/{model_id}/merge") or {}
|
|
228
|
+
|
|
229
|
+
def download(self, model_id: str, artifact: str, dest: str) -> str:
|
|
230
|
+
"""``GET /v1/models/{id}/download?artifact=`` — stream a zip to ``dest``."""
|
|
231
|
+
url = f"/v1/models/{model_id}/download"
|
|
232
|
+
from urllib.parse import urlencode
|
|
233
|
+
|
|
234
|
+
full = url + "?" + urlencode({"artifact": artifact})
|
|
235
|
+
raw = self._http.request("GET", full)
|
|
236
|
+
if isinstance(raw, bytes):
|
|
237
|
+
with open(dest, "wb") as f:
|
|
238
|
+
f.write(raw)
|
|
239
|
+
return dest
|
|
240
|
+
# If the backend wrapped the response, the binary body is gone; raise.
|
|
241
|
+
raise AriaPinError("download did not return a binary stream", status=0)
|
|
242
|
+
|
|
243
|
+
def export_model(self, model_id: str, fmt: str) -> dict[str, Any]:
|
|
244
|
+
return self._http.request(
|
|
245
|
+
"POST", f"/v1/models/{model_id}/export", json_body={"format": fmt}
|
|
246
|
+
) or {}
|
|
247
|
+
|
|
248
|
+
def exports(self) -> list[dict[str, Any]]:
|
|
249
|
+
return self._http.request("GET", "/v1/exports") or []
|
|
250
|
+
|
|
251
|
+
def get_export(self, export_id: str) -> dict[str, Any]:
|
|
252
|
+
return self._http.request("GET", f"/v1/exports/{export_id}") or {}
|
|
253
|
+
|
|
254
|
+
# --- evaluation (RL quality regression) ----------------------------
|
|
255
|
+
def evaluate_model(
|
|
256
|
+
self, model_id: str, dataset_id: str, baseline: str | None = None
|
|
257
|
+
) -> dict[str, Any]:
|
|
258
|
+
body: dict[str, Any] = {"dataset_id": dataset_id}
|
|
259
|
+
if baseline is not None:
|
|
260
|
+
body["baseline"] = baseline
|
|
261
|
+
return self._http.request("POST", f"/v1/models/{model_id}/evaluate", json_body=body) or {}
|
|
262
|
+
|
|
263
|
+
def model_evaluations(self, model_id: str) -> list[dict[str, Any]]:
|
|
264
|
+
data = self._http.request("GET", f"/v1/models/{model_id}/evaluations") or {}
|
|
265
|
+
return data.get("items", []) if isinstance(data, dict) else data
|
|
266
|
+
|
|
267
|
+
def get_evaluation(self, eval_id: str) -> dict[str, Any]:
|
|
268
|
+
return self._http.request("GET", f"/v1/evaluations/{eval_id}") or {}
|
|
269
|
+
|
|
270
|
+
# --- billing -------------------------------------------------------
|
|
271
|
+
def billing_summary(self) -> dict[str, Any]:
|
|
272
|
+
return self._http.request("GET", "/v1/billing/summary") or {}
|
|
273
|
+
|
|
274
|
+
def billing_usage(self, period: str = "30d", limit: int = 50) -> list[dict[str, Any]]:
|
|
275
|
+
data = self._http.request(
|
|
276
|
+
"GET", "/v1/billing/usage", query={"period": period, "limit": limit}
|
|
277
|
+
) or {}
|
|
278
|
+
return data.get("items", []) if isinstance(data, dict) else data
|
|
279
|
+
|
|
280
|
+
# --- raw passthrough (for advanced endpoints) ----------------------
|
|
281
|
+
def raw_request(self, method: str, path: str, **kw) -> Any:
|
|
282
|
+
return self._http.request(method, path, **kw)
|
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
"""Exceptions for the ariapin Python client."""
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
class AriaPinError(Exception):
|
|
5
|
+
"""Raised for any non-2xx response from the ariapin API.
|
|
6
|
+
|
|
7
|
+
Carries the HTTP ``status`` code, the backend ``code`` (from the
|
|
8
|
+
unified ``{"code", "data", "message"}`` envelope) and ``message``.
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
def __init__(self, message: str, status: int = 0, code: int = 0, data=None):
|
|
12
|
+
super().__init__(message)
|
|
13
|
+
self.message = message
|
|
14
|
+
self.status = status
|
|
15
|
+
self.code = code
|
|
16
|
+
self.data = data
|
|
17
|
+
|
|
18
|
+
def __str__(self) -> str:
|
|
19
|
+
if self.status or self.code:
|
|
20
|
+
return f"[status={self.status} code={self.code}] {self.message}"
|
|
21
|
+
return self.message
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
class AriaPinConnectionError(AriaPinError):
|
|
25
|
+
"""Raised when the API cannot be reached (DNS, TLS, network)."""
|
|
@@ -0,0 +1,334 @@
|
|
|
1
|
+
"""Data transfer objects mirroring the ariapin REST contract.
|
|
2
|
+
|
|
3
|
+
Field names match the backend JSON tags in ``requirements.md §4.2`` so the
|
|
4
|
+
client can serialize requests without translation. The ``Job`` object is a
|
|
5
|
+
lazy handle returned by ``AriaPinClient.train`` / ``get_job``; its methods
|
|
6
|
+
delegate back to the client.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
from typing import Any
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
class LoraConfig:
|
|
15
|
+
"""LoRA / QLoRA / DoRA configuration (requirements.md §4.2)."""
|
|
16
|
+
|
|
17
|
+
def __init__(
|
|
18
|
+
self,
|
|
19
|
+
method: str = "lora",
|
|
20
|
+
rank: int = 16,
|
|
21
|
+
alpha: int = 32,
|
|
22
|
+
lora_dropout: float = 0.05,
|
|
23
|
+
target_modules: list[str] | None = None,
|
|
24
|
+
quantization: str | None = None,
|
|
25
|
+
):
|
|
26
|
+
self.method = method
|
|
27
|
+
self.rank = rank
|
|
28
|
+
self.alpha = alpha
|
|
29
|
+
self.lora_dropout = lora_dropout
|
|
30
|
+
self.target_modules = target_modules or ["q_proj", "v_proj"]
|
|
31
|
+
self.quantization = quantization
|
|
32
|
+
|
|
33
|
+
def to_dict(self) -> dict[str, Any]:
|
|
34
|
+
d: dict[str, Any] = {
|
|
35
|
+
"method": self.method,
|
|
36
|
+
"rank": self.rank,
|
|
37
|
+
"alpha": self.alpha,
|
|
38
|
+
"lora_dropout": self.lora_dropout,
|
|
39
|
+
"target_modules": self.target_modules,
|
|
40
|
+
}
|
|
41
|
+
if self.quantization is not None:
|
|
42
|
+
d["quantization"] = self.quantization
|
|
43
|
+
return d
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
class Hyperparams:
|
|
47
|
+
"""Training hyper-parameters (whitelist:
|
|
48
|
+
epochs/lr/max_length/group_size/beta/num_generations/kl_coef/reward_type)."""
|
|
49
|
+
|
|
50
|
+
def __init__(
|
|
51
|
+
self,
|
|
52
|
+
epochs: int = 1,
|
|
53
|
+
lr: float = 4e-5,
|
|
54
|
+
max_length: int = 1024,
|
|
55
|
+
group_size: int = 4,
|
|
56
|
+
beta: float = 0.04,
|
|
57
|
+
num_generations: int = 8,
|
|
58
|
+
kl_coef: float = 0.05,
|
|
59
|
+
reward_type: str = "length",
|
|
60
|
+
):
|
|
61
|
+
self.epochs = epochs
|
|
62
|
+
self.lr = lr
|
|
63
|
+
self.max_length = max_length
|
|
64
|
+
self.group_size = group_size
|
|
65
|
+
self.beta = beta
|
|
66
|
+
self.num_generations = num_generations
|
|
67
|
+
self.kl_coef = kl_coef
|
|
68
|
+
self.reward_type = reward_type
|
|
69
|
+
|
|
70
|
+
def to_dict(self) -> dict[str, Any]:
|
|
71
|
+
return {
|
|
72
|
+
"epochs": self.epochs,
|
|
73
|
+
"lr": self.lr,
|
|
74
|
+
"max_length": self.max_length,
|
|
75
|
+
"group_size": self.group_size,
|
|
76
|
+
"beta": self.beta,
|
|
77
|
+
"num_generations": self.num_generations,
|
|
78
|
+
"kl_coef": self.kl_coef,
|
|
79
|
+
"reward_type": self.reward_type,
|
|
80
|
+
}
|
|
81
|
+
|
|
82
|
+
@classmethod
|
|
83
|
+
def from_dict(cls, d: dict[str, Any] | None) -> "Hyperparams":
|
|
84
|
+
if not d:
|
|
85
|
+
return cls()
|
|
86
|
+
return cls(
|
|
87
|
+
epochs=d.get("epochs", 1),
|
|
88
|
+
lr=d.get("lr", 4e-5),
|
|
89
|
+
max_length=d.get("max_length", 1024),
|
|
90
|
+
group_size=d.get("group_size", 4),
|
|
91
|
+
beta=d.get("beta", 0.04),
|
|
92
|
+
num_generations=d.get("num_generations", 8),
|
|
93
|
+
kl_coef=d.get("kl_coef", 0.05),
|
|
94
|
+
reward_type=d.get("reward_type", "length"),
|
|
95
|
+
)
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
class AutoInterventionRule:
|
|
99
|
+
"""A single reward-hacking detection rule (requirements.md §4.2)."""
|
|
100
|
+
|
|
101
|
+
def __init__(
|
|
102
|
+
self,
|
|
103
|
+
id: str,
|
|
104
|
+
metric: str,
|
|
105
|
+
condition: str,
|
|
106
|
+
threshold: float,
|
|
107
|
+
action: str,
|
|
108
|
+
window_steps: int = 50,
|
|
109
|
+
min_consecutive: int = 5,
|
|
110
|
+
hyperparams: dict[str, Any] | None = None,
|
|
111
|
+
):
|
|
112
|
+
self.id = id
|
|
113
|
+
self.metric = metric
|
|
114
|
+
self.condition = condition
|
|
115
|
+
self.threshold = threshold
|
|
116
|
+
self.action = action
|
|
117
|
+
self.window_steps = window_steps
|
|
118
|
+
self.min_consecutive = min_consecutive
|
|
119
|
+
self.hyperparams = hyperparams
|
|
120
|
+
|
|
121
|
+
def to_dict(self) -> dict[str, Any]:
|
|
122
|
+
d: dict[str, Any] = {
|
|
123
|
+
"id": self.id,
|
|
124
|
+
"metric": self.metric,
|
|
125
|
+
"condition": self.condition,
|
|
126
|
+
"threshold": self.threshold,
|
|
127
|
+
"action": self.action,
|
|
128
|
+
"window_steps": self.window_steps,
|
|
129
|
+
"min_consecutive": self.min_consecutive,
|
|
130
|
+
}
|
|
131
|
+
if self.hyperparams is not None:
|
|
132
|
+
d["hyperparams"] = self.hyperparams
|
|
133
|
+
return d
|
|
134
|
+
|
|
135
|
+
|
|
136
|
+
class AutoIntervention:
|
|
137
|
+
"""Auto-intervention policy attached to a job."""
|
|
138
|
+
|
|
139
|
+
def __init__(
|
|
140
|
+
self,
|
|
141
|
+
enabled: bool,
|
|
142
|
+
rules: list[AutoInterventionRule] | None = None,
|
|
143
|
+
cooldown_seconds: int = 600,
|
|
144
|
+
):
|
|
145
|
+
self.enabled = enabled
|
|
146
|
+
self.rules = rules or []
|
|
147
|
+
self.cooldown_seconds = cooldown_seconds
|
|
148
|
+
|
|
149
|
+
def to_dict(self) -> dict[str, Any]:
|
|
150
|
+
return {
|
|
151
|
+
"enabled": self.enabled,
|
|
152
|
+
"rules": [r.to_dict() for r in self.rules],
|
|
153
|
+
"cooldown_seconds": self.cooldown_seconds,
|
|
154
|
+
}
|
|
155
|
+
|
|
156
|
+
|
|
157
|
+
class GPUNode:
|
|
158
|
+
"""A managed GPU node (SSH-deployable host for agents)."""
|
|
159
|
+
|
|
160
|
+
def __init__(
|
|
161
|
+
self,
|
|
162
|
+
id: str,
|
|
163
|
+
name: str,
|
|
164
|
+
ip: str,
|
|
165
|
+
ssh_user: str,
|
|
166
|
+
ssh_port: int = 22,
|
|
167
|
+
spec: str = "",
|
|
168
|
+
labels: list[str] | None = None,
|
|
169
|
+
status: str = "active",
|
|
170
|
+
agent_count: int = 0,
|
|
171
|
+
):
|
|
172
|
+
self.id = id
|
|
173
|
+
self.name = name
|
|
174
|
+
self.ip = ip
|
|
175
|
+
self.ssh_user = ssh_user
|
|
176
|
+
self.ssh_port = ssh_port
|
|
177
|
+
self.spec = spec
|
|
178
|
+
self.labels = labels or []
|
|
179
|
+
self.status = status
|
|
180
|
+
self.agent_count = agent_count
|
|
181
|
+
|
|
182
|
+
@classmethod
|
|
183
|
+
def from_dict(cls, d: dict[str, Any]) -> "GPUNode":
|
|
184
|
+
return cls(
|
|
185
|
+
id=d.get("id", ""),
|
|
186
|
+
name=d.get("name", ""),
|
|
187
|
+
ip=d.get("ip", ""),
|
|
188
|
+
ssh_user=d.get("ssh_user", ""),
|
|
189
|
+
ssh_port=d.get("ssh_port", 22),
|
|
190
|
+
spec=d.get("spec", ""),
|
|
191
|
+
labels=d.get("labels") or [],
|
|
192
|
+
status=d.get("status", "active"),
|
|
193
|
+
agent_count=d.get("agent_count", 0),
|
|
194
|
+
)
|
|
195
|
+
|
|
196
|
+
|
|
197
|
+
class Agent:
|
|
198
|
+
"""A deployed agent on a GPU node."""
|
|
199
|
+
|
|
200
|
+
def __init__(
|
|
201
|
+
self,
|
|
202
|
+
id: str,
|
|
203
|
+
name: str,
|
|
204
|
+
node_id: str = "",
|
|
205
|
+
version: str = "",
|
|
206
|
+
deploy_state: str = "",
|
|
207
|
+
last_error: str = "",
|
|
208
|
+
status: str = "",
|
|
209
|
+
gpu_count: int = 0,
|
|
210
|
+
):
|
|
211
|
+
self.id = id
|
|
212
|
+
self.name = name
|
|
213
|
+
self.node_id = node_id
|
|
214
|
+
self.version = version
|
|
215
|
+
self.deploy_state = deploy_state
|
|
216
|
+
self.last_error = last_error
|
|
217
|
+
self.status = status
|
|
218
|
+
self.gpu_count = gpu_count
|
|
219
|
+
|
|
220
|
+
@classmethod
|
|
221
|
+
def from_dict(cls, d: dict[str, Any]) -> "Agent":
|
|
222
|
+
return cls(
|
|
223
|
+
id=d.get("id", ""),
|
|
224
|
+
name=d.get("name", ""),
|
|
225
|
+
node_id=d.get("node_id", ""),
|
|
226
|
+
version=d.get("version", ""),
|
|
227
|
+
deploy_state=d.get("deploy_state", ""),
|
|
228
|
+
last_error=d.get("last_error", ""),
|
|
229
|
+
status=d.get("status", ""),
|
|
230
|
+
gpu_count=d.get("gpu_count", 0),
|
|
231
|
+
)
|
|
232
|
+
|
|
233
|
+
|
|
234
|
+
_TERMINAL = {"succeeded", "failed", "cancelled", "stopped"}
|
|
235
|
+
|
|
236
|
+
|
|
237
|
+
class Job:
|
|
238
|
+
"""Lazy handle for a training job returned by ``train`` / ``get_job``.
|
|
239
|
+
|
|
240
|
+
Methods delegate back to the owning :class:`AriaPinClient`.
|
|
241
|
+
"""
|
|
242
|
+
|
|
243
|
+
def __init__(self, client: "AriaPinClient", job_id: str, raw: dict[str, Any] | None = None):
|
|
244
|
+
self._client = client
|
|
245
|
+
self.job_id = job_id
|
|
246
|
+
self._raw = raw or {}
|
|
247
|
+
|
|
248
|
+
# --- introspection -------------------------------------------------
|
|
249
|
+
def refresh(self) -> "Job":
|
|
250
|
+
self._raw = self._client.get_job_raw(self.job_id)
|
|
251
|
+
return self
|
|
252
|
+
|
|
253
|
+
@property
|
|
254
|
+
def status(self) -> str:
|
|
255
|
+
if not self._raw:
|
|
256
|
+
self.refresh()
|
|
257
|
+
return str(self._raw.get("status", ""))
|
|
258
|
+
|
|
259
|
+
@property
|
|
260
|
+
def progress(self) -> dict[str, Any]:
|
|
261
|
+
if not self._raw:
|
|
262
|
+
self.refresh()
|
|
263
|
+
return self._raw.get("progress") or {}
|
|
264
|
+
|
|
265
|
+
@property
|
|
266
|
+
def mlflow_url(self) -> str:
|
|
267
|
+
if not self._raw:
|
|
268
|
+
self.refresh()
|
|
269
|
+
return str(self._raw.get("mlflow_url", "") or "")
|
|
270
|
+
|
|
271
|
+
@property
|
|
272
|
+
def raw(self) -> dict[str, Any]:
|
|
273
|
+
if not self._raw:
|
|
274
|
+
self.refresh()
|
|
275
|
+
return self._raw
|
|
276
|
+
|
|
277
|
+
def is_terminal(self) -> bool:
|
|
278
|
+
return self.status in _TERMINAL
|
|
279
|
+
|
|
280
|
+
# --- blocking wait --------------------------------------------------
|
|
281
|
+
def wait(self, interval: int = 10, timeout: int | None = None) -> "Job":
|
|
282
|
+
"""Poll until the job reaches a terminal state.
|
|
283
|
+
|
|
284
|
+
``interval`` is the sleep seconds between polls; ``timeout`` is the
|
|
285
|
+
maximum wall-clock seconds (``None`` = wait indefinitely).
|
|
286
|
+
"""
|
|
287
|
+
import time
|
|
288
|
+
|
|
289
|
+
waited = 0
|
|
290
|
+
self.refresh()
|
|
291
|
+
while not self.is_terminal():
|
|
292
|
+
time.sleep(interval)
|
|
293
|
+
waited += interval
|
|
294
|
+
if timeout is not None and waited >= timeout:
|
|
295
|
+
raise TimeoutError(
|
|
296
|
+
f"job {self.job_id} did not finish within {timeout}s (status={self.status})"
|
|
297
|
+
)
|
|
298
|
+
self.refresh()
|
|
299
|
+
return self
|
|
300
|
+
|
|
301
|
+
# --- actions -------------------------------------------------------
|
|
302
|
+
def cancel(self) -> dict[str, Any]:
|
|
303
|
+
return self._client.cancel_job(self.job_id)
|
|
304
|
+
|
|
305
|
+
def intervene(self, action: str, reason: str = "", hyperparams: dict[str, Any] | None = None) -> dict[str, Any]:
|
|
306
|
+
return self._client.intervene(self.job_id, action, reason, hyperparams)
|
|
307
|
+
|
|
308
|
+
def set_auto_intervention(
|
|
309
|
+
self, enabled: bool, rules: list[AutoInterventionRule] | None = None, cooldown_seconds: int = 600
|
|
310
|
+
) -> dict[str, Any]:
|
|
311
|
+
ai = AutoIntervention(enabled=enabled, rules=rules or [], cooldown_seconds=cooldown_seconds)
|
|
312
|
+
return self._client.set_auto_intervention(self.job_id, ai)
|
|
313
|
+
|
|
314
|
+
def interventions(self) -> list[dict[str, Any]]:
|
|
315
|
+
return self._client.job_interventions(self.job_id)
|
|
316
|
+
|
|
317
|
+
def analysis(self) -> dict[str, Any]:
|
|
318
|
+
return self._client.job_analysis(self.job_id)
|
|
319
|
+
|
|
320
|
+
# --- lifecycle (start/pause/resume/stop/delete) ---------------------
|
|
321
|
+
def start(self) -> dict[str, Any]:
|
|
322
|
+
return self._client.start_job(self.job_id)
|
|
323
|
+
|
|
324
|
+
def pause(self) -> dict[str, Any]:
|
|
325
|
+
return self._client.pause_job(self.job_id)
|
|
326
|
+
|
|
327
|
+
def resume(self) -> dict[str, Any]:
|
|
328
|
+
return self._client.resume_job(self.job_id)
|
|
329
|
+
|
|
330
|
+
def stop(self) -> dict[str, Any]:
|
|
331
|
+
return self._client.stop_job(self.job_id)
|
|
332
|
+
|
|
333
|
+
def delete(self) -> dict[str, Any]:
|
|
334
|
+
return self._client.delete_job(self.job_id)
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: ariapin
|
|
3
|
+
Version: 1.7.0
|
|
4
|
+
Summary: Python client for the ariapin end-side post-training API (LoRA/QLoRA + RL with SFT/DPO/GRPO)
|
|
5
|
+
Author: aria compute
|
|
6
|
+
License: Apache-2.0
|
|
7
|
+
Project-URL: Homepage, https://github.com/ariacompute/harness
|
|
8
|
+
Keywords: llm,post-training,lora,qlora,dora,grpo,dpo,sft,rl
|
|
9
|
+
Classifier: Programming Language :: Python :: 3
|
|
10
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
11
|
+
Classifier: License :: OSI Approved :: Apache Software License
|
|
12
|
+
Requires-Python: >=3.10
|
|
13
|
+
Description-Content-Type: text/markdown
|
|
14
|
+
Provides-Extra: dev
|
|
15
|
+
Requires-Dist: pytest>=7; extra == "dev"
|
|
@@ -0,0 +1,12 @@
|
|
|
1
|
+
pyproject.toml
|
|
2
|
+
ariapin/__init__.py
|
|
3
|
+
ariapin/_http.py
|
|
4
|
+
ariapin/client.py
|
|
5
|
+
ariapin/exceptions.py
|
|
6
|
+
ariapin/models.py
|
|
7
|
+
ariapin.egg-info/PKG-INFO
|
|
8
|
+
ariapin.egg-info/SOURCES.txt
|
|
9
|
+
ariapin.egg-info/dependency_links.txt
|
|
10
|
+
ariapin.egg-info/requires.txt
|
|
11
|
+
ariapin.egg-info/top_level.txt
|
|
12
|
+
tests/test_client.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
ariapin
|
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=64"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "ariapin"
|
|
7
|
+
version = "1.7.0"
|
|
8
|
+
description = "Python client for the ariapin end-side post-training API (LoRA/QLoRA + RL with SFT/DPO/GRPO)"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.10"
|
|
11
|
+
license = { text = "Apache-2.0" }
|
|
12
|
+
authors = [{ name = "aria compute" }]
|
|
13
|
+
keywords = ["llm", "post-training", "lora", "qlora", "dora", "grpo", "dpo", "sft", "rl"]
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Programming Language :: Python :: 3",
|
|
16
|
+
"Programming Language :: Python :: 3.10",
|
|
17
|
+
"License :: OSI Approved :: Apache Software License",
|
|
18
|
+
]
|
|
19
|
+
dependencies = []
|
|
20
|
+
|
|
21
|
+
[project.optional-dependencies]
|
|
22
|
+
dev = ["pytest>=7"]
|
|
23
|
+
|
|
24
|
+
[project.urls]
|
|
25
|
+
Homepage = "https://github.com/ariacompute/harness"
|
|
26
|
+
|
|
27
|
+
[tool.setuptools.packages.find]
|
|
28
|
+
where = ["."]
|
|
29
|
+
include = ["ariapin*"]
|
ariapin-1.7.0/setup.cfg
ADDED
|
@@ -0,0 +1,252 @@
|
|
|
1
|
+
"""Unit tests for the ariapin Python client.
|
|
2
|
+
|
|
3
|
+
Uses only stdlib (unittest) and a tiny mock transport so it runs without the
|
|
4
|
+
network, GPU, or the ariapin backend. Verifies request construction, response
|
|
5
|
+
parsing, error mapping, and terminal-state polling.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
import io
|
|
9
|
+
import json
|
|
10
|
+
import time
|
|
11
|
+
import unittest
|
|
12
|
+
from urllib.error import HTTPError
|
|
13
|
+
from urllib.request import Request
|
|
14
|
+
|
|
15
|
+
from ariapin import (
|
|
16
|
+
AriaPinClient,
|
|
17
|
+
AriaPinError,
|
|
18
|
+
AutoIntervention,
|
|
19
|
+
AutoInterventionRule,
|
|
20
|
+
Hyperparams,
|
|
21
|
+
Job,
|
|
22
|
+
LoraConfig,
|
|
23
|
+
)
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class _FakeHTTP:
|
|
27
|
+
"""Replaces _HTTPClient.request with scripted responses.
|
|
28
|
+
|
|
29
|
+
Mirrors the real transport's envelope unwrapping: a returned dict with a
|
|
30
|
+
``code`` key is unwrapped to its ``data`` field, and a non-zero ``code``
|
|
31
|
+
raises AriaPinError.
|
|
32
|
+
"""
|
|
33
|
+
|
|
34
|
+
def __init__(self):
|
|
35
|
+
self.calls = []
|
|
36
|
+
self.handler = None # callable(method, path, kwargs) -> payload
|
|
37
|
+
|
|
38
|
+
@staticmethod
|
|
39
|
+
def _unwrap(payload):
|
|
40
|
+
if isinstance(payload, dict) and "code" in payload:
|
|
41
|
+
if payload.get("code") not in (0, None):
|
|
42
|
+
raise AriaPinError(
|
|
43
|
+
str(payload.get("message", "")), status=200, code=int(payload.get("code", 0))
|
|
44
|
+
)
|
|
45
|
+
return payload.get("data")
|
|
46
|
+
return payload
|
|
47
|
+
|
|
48
|
+
def request(self, method, path, **kwargs):
|
|
49
|
+
self.calls.append((method, path, kwargs))
|
|
50
|
+
if self.handler is None:
|
|
51
|
+
return None
|
|
52
|
+
try:
|
|
53
|
+
payload = self.handler(method, path, kwargs)
|
|
54
|
+
except HTTPError as e:
|
|
55
|
+
# Mimic _HTTPClient: convert HTTPError -> AriaPinError (4xx immediate,
|
|
56
|
+
# 5xx would retry in real transport; here we surface directly).
|
|
57
|
+
raw = e.read() if e.fp else b""
|
|
58
|
+
msg = ""
|
|
59
|
+
code = 0
|
|
60
|
+
if raw:
|
|
61
|
+
try:
|
|
62
|
+
p = json.loads(raw.decode("utf-8", "replace"))
|
|
63
|
+
msg = str(p.get("message", ""))
|
|
64
|
+
code = int(p.get("code", 0))
|
|
65
|
+
except Exception:
|
|
66
|
+
msg = raw.decode("utf-8", "replace")
|
|
67
|
+
raise AriaPinError(msg or f"HTTP {e.code}", status=e.code, code=code)
|
|
68
|
+
return self._unwrap(payload)
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
def _envelope(data, code=0, message=""):
|
|
72
|
+
return {"code": code, "data": data, "message": message}
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
def _make_client():
|
|
76
|
+
c = AriaPinClient(api_key="sk-test", base_url="http://test")
|
|
77
|
+
c._http = _FakeHTTP()
|
|
78
|
+
return c, c._http
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
class TestTrainRequest(unittest.TestCase):
|
|
82
|
+
def test_train_builds_correct_body_and_returns_job(self):
|
|
83
|
+
client, http = _make_client()
|
|
84
|
+
|
|
85
|
+
def handler(method, path, kwargs):
|
|
86
|
+
self.assertEqual(method, "POST")
|
|
87
|
+
self.assertEqual(path, "/v1/jobs")
|
|
88
|
+
body = kwargs["json_body"]
|
|
89
|
+
self.assertEqual(body["job_type"], "grpo")
|
|
90
|
+
self.assertEqual(body["base_model"], "Qwen/Qwen2.5-0.5B-Instruct")
|
|
91
|
+
self.assertEqual(body["dataset_id"], "ds_x")
|
|
92
|
+
self.assertEqual(body["lora_config"]["method"], "lora")
|
|
93
|
+
self.assertEqual(body["lora_config"]["rank"], 16)
|
|
94
|
+
self.assertEqual(body["hyperparams"]["group_size"], 4)
|
|
95
|
+
self.assertEqual(body["hyperparams"]["beta"], 0.04)
|
|
96
|
+
self.assertTrue(body["auto_intervention"]["enabled"])
|
|
97
|
+
self.assertEqual(len(body["auto_intervention"]["rules"]), 1)
|
|
98
|
+
return _envelope({"job_id": "jb_1", "status": "queued"})
|
|
99
|
+
|
|
100
|
+
http.handler = handler
|
|
101
|
+
job = client.train(
|
|
102
|
+
base_model="Qwen/Qwen2.5-0.5B-Instruct",
|
|
103
|
+
job_type="grpo",
|
|
104
|
+
lora=LoraConfig(method="lora", rank=16),
|
|
105
|
+
dataset_id="ds_x",
|
|
106
|
+
hyperparams=Hyperparams(epochs=1, lr=4e-5, group_size=4, beta=0.04),
|
|
107
|
+
auto_intervention=AutoIntervention(
|
|
108
|
+
enabled=True,
|
|
109
|
+
rules=[
|
|
110
|
+
AutoInterventionRule(
|
|
111
|
+
id="reward_spike", metric="reward", condition="spike",
|
|
112
|
+
threshold=2.0, action="update_hyperparams",
|
|
113
|
+
hyperparams={"lr": 2e-5, "beta": 0.08},
|
|
114
|
+
)
|
|
115
|
+
],
|
|
116
|
+
),
|
|
117
|
+
)
|
|
118
|
+
self.assertIsInstance(job, Job)
|
|
119
|
+
self.assertEqual(job.job_id, "jb_1")
|
|
120
|
+
self.assertEqual(job.status, "queued")
|
|
121
|
+
|
|
122
|
+
|
|
123
|
+
class TestErrorMapping(unittest.TestCase):
|
|
124
|
+
def _raise(self, status, payload):
|
|
125
|
+
raise HTTPError(
|
|
126
|
+
url="http://test/v1/jobs",
|
|
127
|
+
code=status,
|
|
128
|
+
msg="err",
|
|
129
|
+
hdrs=None,
|
|
130
|
+
fp=io.BytesIO(json.dumps(payload).encode()),
|
|
131
|
+
)
|
|
132
|
+
|
|
133
|
+
def test_402_insufficient_balance(self):
|
|
134
|
+
client, http = _make_client()
|
|
135
|
+
http.handler = lambda m, p, k: self._raise(
|
|
136
|
+
402, _envelope(None, code=402, message="insufficient balance")
|
|
137
|
+
)
|
|
138
|
+
with self.assertRaises(AriaPinError) as ctx:
|
|
139
|
+
client.train("m", "sft", dataset_id="d")
|
|
140
|
+
self.assertEqual(ctx.exception.status, 402)
|
|
141
|
+
|
|
142
|
+
def test_422_invalid_lora(self):
|
|
143
|
+
client, http = _make_client()
|
|
144
|
+
http.handler = lambda m, p, k: self._raise(
|
|
145
|
+
422, _envelope(None, code=422, message="rank out of range [1,64]")
|
|
146
|
+
)
|
|
147
|
+
with self.assertRaises(AriaPinError) as ctx:
|
|
148
|
+
client.train("m", "sft", dataset_id="d", lora=LoraConfig(rank=0))
|
|
149
|
+
self.assertEqual(ctx.exception.status, 422)
|
|
150
|
+
|
|
151
|
+
def test_envelope_nonzero_code_raises(self):
|
|
152
|
+
client, http = _make_client()
|
|
153
|
+
http.handler = lambda m, p, k: _envelope(None, code=409, message="conflict")
|
|
154
|
+
with self.assertRaises(AriaPinError) as ctx:
|
|
155
|
+
client.cancel_job("jb_1")
|
|
156
|
+
self.assertEqual(ctx.exception.code, 409)
|
|
157
|
+
|
|
158
|
+
|
|
159
|
+
class TestJobWait(unittest.TestCase):
|
|
160
|
+
def test_wait_polls_until_terminal(self):
|
|
161
|
+
client, http = _make_client()
|
|
162
|
+
states = iter(["running", "running", "running", "running", "succeeded"])
|
|
163
|
+
|
|
164
|
+
def handler(method, path, kwargs):
|
|
165
|
+
# GET /v1/jobs/{id} returns the next simulated state
|
|
166
|
+
return _envelope({"job_id": "jb_1", "status": next(states),
|
|
167
|
+
"progress": {"step": 10, "loss": 0.3}})
|
|
168
|
+
|
|
169
|
+
http.handler = handler
|
|
170
|
+
# speed up: monkeypatch sleep to no-op
|
|
171
|
+
real_sleep = time.sleep
|
|
172
|
+
time.sleep = lambda *_: None
|
|
173
|
+
try:
|
|
174
|
+
job = client.get_job("jb_1")
|
|
175
|
+
final = job.wait()
|
|
176
|
+
finally:
|
|
177
|
+
time.sleep = real_sleep
|
|
178
|
+
self.assertEqual(final.status, "succeeded")
|
|
179
|
+
self.assertEqual(final.progress["step"], 10)
|
|
180
|
+
# GET calls: get_job_raw(1) + wait refresh(2) + 3 loop refreshes(3,4,5)
|
|
181
|
+
self.assertEqual(len(http.calls), 5)
|
|
182
|
+
|
|
183
|
+
|
|
184
|
+
class TestInterventionsAndAnalysis(unittest.TestCase):
|
|
185
|
+
def test_intervene_constructs_body(self):
|
|
186
|
+
client, http = _make_client()
|
|
187
|
+
captured = {}
|
|
188
|
+
|
|
189
|
+
def handler(method, path, kwargs):
|
|
190
|
+
captured["method"] = method
|
|
191
|
+
captured["path"] = path
|
|
192
|
+
captured["body"] = kwargs.get("json_body")
|
|
193
|
+
return _envelope({"job_id": "jb_1", "action": "stop", "status": "stopping"})
|
|
194
|
+
|
|
195
|
+
http.handler = handler
|
|
196
|
+
res = client.intervene("jb_1", "stop", reason="reward spike")
|
|
197
|
+
self.assertEqual(captured["method"], "POST")
|
|
198
|
+
self.assertEqual(captured["path"], "/v1/jobs/jb_1/intervene")
|
|
199
|
+
self.assertEqual(captured["body"]["action"], "stop")
|
|
200
|
+
self.assertEqual(captured["body"]["reason"], "reward spike")
|
|
201
|
+
self.assertEqual(res["status"], "stopping")
|
|
202
|
+
|
|
203
|
+
def test_set_auto_intervention_put(self):
|
|
204
|
+
client, http = _make_client()
|
|
205
|
+
captured = {}
|
|
206
|
+
|
|
207
|
+
def handler(method, path, kwargs):
|
|
208
|
+
captured["method"] = method
|
|
209
|
+
captured["body"] = kwargs.get("json_body")
|
|
210
|
+
return _envelope({"job_id": "jb_1", "auto_intervention": kwargs.get("json_body")})
|
|
211
|
+
|
|
212
|
+
http.handler = handler
|
|
213
|
+
client.set_auto_intervention(
|
|
214
|
+
"jb_1", AutoIntervention(enabled=False, cooldown_seconds=300)
|
|
215
|
+
)
|
|
216
|
+
self.assertEqual(captured["method"], "PUT")
|
|
217
|
+
self.assertEqual(captured["body"]["enabled"], False)
|
|
218
|
+
self.assertEqual(captured["body"]["cooldown_seconds"], 300)
|
|
219
|
+
|
|
220
|
+
def test_analysis_and_interventions(self):
|
|
221
|
+
client, http = _make_client()
|
|
222
|
+
|
|
223
|
+
def handler(method, path, kwargs):
|
|
224
|
+
if path.endswith("/analysis"):
|
|
225
|
+
return _envelope({"risk": "ok", "reward": {"mean": 0.5}})
|
|
226
|
+
if path.endswith("/interventions"):
|
|
227
|
+
return _envelope({"items": [{"action": "stop", "trigger": "auto"}]})
|
|
228
|
+
return None
|
|
229
|
+
|
|
230
|
+
http.handler = handler
|
|
231
|
+
self.assertEqual(client.job_analysis("jb_1")["risk"], "ok")
|
|
232
|
+
self.assertEqual(len(client.job_interventions("jb_1")), 1)
|
|
233
|
+
|
|
234
|
+
|
|
235
|
+
class TestExportsAndBilling(unittest.TestCase):
|
|
236
|
+
def test_export_and_billing(self):
|
|
237
|
+
client, http = _make_client()
|
|
238
|
+
|
|
239
|
+
def handler(method, path, kwargs):
|
|
240
|
+
if path.endswith("/export"):
|
|
241
|
+
return _envelope({"export_id": "ex_1", "status": "queued"})
|
|
242
|
+
if path == "/v1/billing/summary":
|
|
243
|
+
return _envelope({"balance": 9.9, "currency": "USD"})
|
|
244
|
+
return None
|
|
245
|
+
|
|
246
|
+
http.handler = handler
|
|
247
|
+
self.assertEqual(client.export_model("md_1", "gguf")["export_id"], "ex_1")
|
|
248
|
+
self.assertEqual(client.billing_summary()["balance"], 9.9)
|
|
249
|
+
|
|
250
|
+
|
|
251
|
+
if __name__ == "__main__":
|
|
252
|
+
unittest.main()
|