import warnings
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any, Literal, TypeAlias
from urllib.parse import urljoin
import requests
from pydantic import BaseModel, ConfigDict
from pydantic.alias_generators import to_camel
from typing_extensions import Self
@dataclass
class ApiConfig:
"""Configuration for interacting with the Embedl Hub REST API."""
base_url: str
api_key: str
@property
def headers(self) -> dict[str, str]:
return {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
"Accept": "application/json",
}
class Model(BaseModel):
"""Base model with camel case aliasing that ignores fields it does not know.
See ``docs/adr/0002-the-sdk-reads-additive-hub-changes.md``.
"""
model_config = ConfigDict(
alias_generator=to_camel, validate_by_name=True, extra="ignore"
)
[docs]
class Project(Model):
id: str
name: str
[docs]
class RunType(Enum):
COMPILE = "COMPILE"
CUSTOM = "CUSTOM"
EVAL = "EVAL"
GRAPH = "GRAPH"
INFERENCE = "INFERENCE"
PROFILE = "PROFILE"
def coerce_run_type(s: str) -> "RunType | str":
"""Map a string to a :class:`RunType` enum member if one matches.
The match is case-insensitive and whitespace is stripped. If no
enum member matches, the original string is returned unchanged so
the caller can treat it as a custom type label.
:param s: A run-type string supplied by the user.
:return: The matching :class:`RunType` member, or *s* unchanged.
"""
try:
return RunType(s.strip().upper())
except ValueError:
return s
[docs]
class RunStatus(Enum):
PENDING = "PENDING"
RUNNING = "RUNNING"
FINISHED = "FINISHED"
FAILED = "FAILED"
KILLED = "KILLED"
[docs]
class Run(Model):
id: str
name: str
description: str | None = None
type: RunType
custom_type: str | None = None
status: RunStatus
parent_run_id: str | None = None
created_at: datetime
started_at: datetime
ended_at: datetime | None
class RunSummary(Model):
id: str
name: str
description: str | None = None
type: RunType
custom_type: str | None = None
status: RunStatus
parent_run_id: str | None = None
created_at: datetime
started_at: datetime
ended_at: datetime | None
class RunWithProject(Run):
project: Project
class RunListPagination(Model):
limit: int
offset: int
total: int
has_more: bool
next_offset: int | None
[docs]
class RunListPage(Model):
items: list[RunSummary]
pagination: RunListPagination
[docs]
class Parameter(Model):
name: str
value: str
measured_at: datetime | None = None
[docs]
class ExternalLink(Model):
id: int | None = None
label: str
url: str
created_at: datetime | None = None
class Metric(Model):
name: str
value: float
step: int | None = None
measured_at: datetime | None = None
[docs]
class Tag(Model):
name: str
value: str
[docs]
class RunDashboard(Model):
id: str
name: str
description: str | None = None
config: dict[str, Any]
is_default: bool = False
opens_run: bool = False
created_at: datetime
updated_at: datetime
class ArtifactStatus(Enum):
PENDING = "PENDING"
UPLOADED = "UPLOADED"
FAILED = "FAILED"
[docs]
class Artifact(Model):
id: str
run_id: str
file_name: str
file_size: str
status: ArtifactStatus
created_at: datetime
updated_at: datetime
class ArtifactUploadResponse(Model):
url: str
expires_at: datetime | None
max_size: str
upload_mode: Literal["API", "SIGNED_URL"]
CompletedRunStatus: TypeAlias = Literal[
RunStatus.FINISHED, RunStatus.KILLED, RunStatus.FAILED
]
class _RemoveParent:
"""Sentinel value to explicitly remove a parent run."""
pass
REMOVE_PARENT = _RemoveParent()
JSONObj: TypeAlias = dict[str, Any]
JSONItems: TypeAlias = list[JSONObj]
JSONData: TypeAlias = JSONObj | JSONItems
def create_project(config: ApiConfig, name: str) -> Project:
"""Create a new project."""
data = _request(config, "POST", "/api/projects", json={"name": name})
data = _expect_dict(data)
return Project(**data)
def get_project_by_name(config: ApiConfig, name: str) -> Project | None:
"""Get project by name, or None if not found."""
try:
data = _request(
config,
"GET",
"/api/projects",
params={"name": name},
)
except ApiError as err:
if err.status_code == 404:
return None
raise
data = _expect_dict(data)
return Project(**data)
def create_run(
config: ApiConfig,
project_id: str,
type: RunType,
started_at: datetime | None = None,
name: str | None = None,
parent_run_id: str | None = None,
custom_type: str | None = None,
description: str | None = None,
pending: bool = False,
) -> Run:
"""Create a new run, pending rather than running when asked to."""
payload = {"type": type.value}
if pending:
payload["status"] = RunStatus.PENDING.value
if name:
payload["name"] = name
if parent_run_id:
payload["parentRunId"] = parent_run_id
if custom_type:
payload["customType"] = custom_type
if description is not None:
payload["description"] = description
if started_at:
payload["startedAt"] = started_at.isoformat()
data = _request(
config,
"POST",
f"/api/projects/{project_id}/runs",
json=payload,
)
data = _expect_dict(data)
return Run(**data)
def get_run(
config: ApiConfig,
project_id: str,
run_id: str,
) -> Run:
"""Get a run by ID."""
data = _request(
config,
"GET",
f"/api/projects/{project_id}/runs/{run_id}",
)
data = _expect_dict(data)
return Run(**data)
def start_pending_run(config: ApiConfig, project_id: str, run_id: str) -> None:
"""Start a pending run; the Hub sets its start time."""
_request(
config,
"PATCH",
f"/api/projects/{project_id}/runs/{run_id}",
json={"status": RunStatus.RUNNING.value},
)
def resolve_run(config: ApiConfig, run_id: str) -> RunWithProject:
"""Get a run by ID alone, with the project it belongs to."""
data = _request(config, "GET", f"/api/runs/{run_id}")
data = _expect_dict(data)
return RunWithProject(**data)
def list_runs(
config: ApiConfig,
project_id: str,
limit: int | None = None,
offset: int | None = None,
statuses: list[RunStatus] | None = None,
types: list[RunType] | None = None,
custom_types: list[str] | None = None,
params: Mapping[str, str] | None = None,
) -> RunListPage:
"""List the runs of a project, newest first.
:param config: API configuration.
:param project_id: ID of the project to list runs for.
:param limit: Maximum number of runs to return (server default and
cap apply).
:param offset: Number of runs to skip, for pagination.
:param statuses: When given, only runs with these statuses.
:param types: When given, only runs of these types.
:param custom_types: When given, only custom runs with these labels.
:param params: When given, only runs that logged every one of these
parameters with exactly that value, sent as given.
:return: One page of runs with pagination info.
"""
query: dict[str, str | int] = {}
if limit is not None:
query["limit"] = limit
if offset is not None:
query["offset"] = offset
if statuses:
query["status"] = ",".join(status.value for status in statuses)
if types:
query["type"] = ",".join(run_type.value for run_type in types)
if custom_types:
query["customType"] = ",".join(custom_types)
for name, value in (params or {}).items():
query[f"param:{name}"] = value
data = _request(
config,
"GET",
f"/api/projects/{project_id}/runs",
params=query or None,
)
data = _expect_dict(data)
return RunListPage(**data)
def get_run_metrics(
config: ApiConfig,
project_id: str,
run_id: str,
) -> list[Metric]:
"""Get the metrics logged on a run."""
data = _request(
config,
"GET",
f"/api/projects/{project_id}/runs/{run_id}/metrics",
)
data = _expect_list(data)
return [Metric(**item) for item in data]
def get_run_params(
config: ApiConfig,
project_id: str,
run_id: str,
) -> list[Parameter]:
"""Get the parameters logged on a run."""
data = _request(
config,
"GET",
f"/api/projects/{project_id}/runs/{run_id}/params",
)
data = _expect_list(data)
return [Parameter(**item) for item in data]
def get_run_tags(
config: ApiConfig,
project_id: str,
run_id: str,
) -> list[Tag]:
"""Get the tags set on a run."""
data = _request(
config,
"GET",
f"/api/projects/{project_id}/runs/{run_id}/tags",
)
data = _expect_list(data)
return [Tag(**item) for item in data]
def get_run_external_links(
config: ApiConfig,
project_id: str,
run_id: str,
) -> list[ExternalLink]:
"""Get the external links logged on a run."""
data = _request(
config,
"GET",
f"/api/projects/{project_id}/runs/{run_id}/external-links",
)
data = _expect_list(data)
return [ExternalLink(**item) for item in data]
def get_run_artifacts(
config: ApiConfig,
project_id: str,
run_id: str,
) -> list[Artifact]:
"""Get the artifacts recorded on a run."""
data = _request(
config,
"GET",
f"/api/projects/{project_id}/runs/{run_id}/artifacts",
)
data = _expect_list(data)
return [Artifact(**item) for item in data]
def update_run(
config: ApiConfig,
project_id: str,
run_id: str,
status: CompletedRunStatus | None,
ended_at: datetime | None,
metrics: list[Metric] | None = None,
params: list[Parameter] | None = None,
external_links: list[ExternalLink] | None = None,
parent_run_id: str | _RemoveParent | None = None,
description: str | None = None,
) -> None:
"""Update run metadata, status, attributes, and end time."""
payload = {}
if status:
payload["status"] = status.value
if ended_at:
payload["endedAt"] = ended_at.isoformat()
if metrics:
payload["metrics"] = [
metric.model_dump(
mode="json", by_alias=True, exclude_defaults=True
)
for metric in metrics
]
if params:
payload["params"] = [
param.model_dump(mode="json", by_alias=True, exclude_defaults=True)
for param in params
]
if external_links:
payload["externalLinks"] = [
link.model_dump(
mode="json",
by_alias=True,
exclude={"id"},
exclude_defaults=True,
)
for link in external_links
]
if parent_run_id is not None:
if isinstance(parent_run_id, _RemoveParent):
payload["parentRunId"] = None
else:
payload["parentRunId"] = parent_run_id
if description is not None:
payload["description"] = description
_request(
config,
"PATCH",
f"/api/projects/{project_id}/runs/{run_id}",
json=payload,
)
PARAM_VALUE_LIMIT = 255
"""The hub's ParamValue schema caps values at 255 UTF-16 code units."""
JS_WHITESPACE = (
" \t\n\r\v\f\u00a0\u1680\u2000\u2001\u2002\u2003\u2004\u2005"
"\u2006\u2007\u2008\u2009\u200a\u2028\u2029\u202f\u205f\u3000"
"\ufeff"
)
"""The characters JavaScript's trim() removes — notably U+FEFF, which
Python's strip() keeps."""
def js_trim(value: str) -> str:
"""Trim with JavaScript semantics, as the hub's validators do."""
return value.strip(JS_WHITESPACE)
def utf16_length(value: str) -> int:
"""String length in UTF-16 code units — the hub's JavaScript
validators count these, so non-BMP characters count double."""
return sum(2 if ord(c) > 0xFFFF else 1 for c in value)
def _take_utf16_units(value: str, units: int) -> str:
"""Take a prefix of at most *units* UTF-16 code units."""
taken = []
used = 0
for c in value:
used += 2 if ord(c) > 0xFFFF else 1
if used > units:
break
taken.append(c)
return "".join(taken)
def truncate_param_value(name: str, value: str) -> str:
"""Clamp an over-long parameter value, warning about the cut.
Shared by the immediate and the batched logging paths so a queued
``log param`` behaves exactly like an immediate one. Lengths are
measured in UTF-16 code units to match the hub's validators. (The
previous client-side limit of 511 had drifted from the server's
255 — a "truncated" value was still rejected with a 400.)
"""
value = js_trim(value)
if utf16_length(value) <= PARAM_VALUE_LIMIT:
return value
warnings.warn(
f'Parameter "{name}" value too long: {utf16_length(value)}. '
f"Limit: {PARAM_VALUE_LIMIT}.",
RuntimeWarning,
)
head = _take_utf16_units(value, 125)
tail = _take_utf16_units(value[::-1], 125)[::-1]
return head + " ... " + tail
def log_param(
config: ApiConfig,
name: str,
value: str,
project_id: str,
run_id: str,
) -> Parameter:
"""Log a parameter for a run."""
value = truncate_param_value(name, value)
payload = {"name": name, "value": value}
data = _request(
config,
"POST",
f"/api/projects/{project_id}/runs/{run_id}/params",
json=payload,
)
data = _expect_dict(data)
return Parameter(**data)
def log_external_link(
config: ApiConfig,
label: str,
url: str,
project_id: str,
run_id: str,
) -> ExternalLink:
"""Log an external link for a run."""
payload = {"label": label, "url": url}
data = _request(
config,
"POST",
f"/api/projects/{project_id}/runs/{run_id}/external-links",
json=payload,
)
data = _expect_dict(data)
return ExternalLink(**data)
def log_metric(
config: ApiConfig,
name: str,
value: float,
project_id: str,
run_id: str,
step: int | None = None,
) -> Metric:
"""Log a metric for a run."""
payload = {"name": name, "value": value}
if step is not None:
payload["step"] = step
data = _request(
config,
"POST",
f"/api/projects/{project_id}/runs/{run_id}/metrics",
json=payload,
)
data = _expect_dict(data)
return Metric(**data)
def log_tag(
config: ApiConfig,
name: str,
value: str,
project_id: str,
run_id: str,
) -> Tag:
"""Log a tag for a run."""
payload = {"name": name, "value": value}
data = _request(
config,
"POST",
f"/api/projects/{project_id}/runs/{run_id}/tags",
json=payload,
)
data = _expect_dict(data)
return Tag(**data)
def create_run_dashboard(
config: ApiConfig,
project_id: str,
run_id: str,
dashboard: JSONObj,
) -> RunDashboard:
"""Store a dashboard for a run from its upload body."""
data = _request(
config,
"POST",
f"/api/projects/{project_id}/runs/{run_id}/dashboards",
json=dashboard,
)
data = _expect_dict(data)
return RunDashboard(**data)
def create_artifact(
config: ApiConfig,
run_id: str,
file_name: str,
file_size: int,
) -> Artifact:
"""Create an artifact for a given run.
Uploading the file for the artifact must be done separately. See
`create_artifact_upload_url()`.
"""
payload = {
"runId": run_id,
"fileName": file_name,
"fileSize": file_size,
}
data = _request(
config,
"POST",
"/api/artifacts",
json=payload,
)
data = _expect_dict(data)
return Artifact(**data)
def create_artifact_upload_url(
config: ApiConfig,
artifact_id: str,
) -> ArtifactUploadResponse:
"""Create a temporary URL for uploading a file for an artifact."""
data = _request(
config,
"POST",
f"/api/artifacts/{artifact_id}/upload-url",
)
data = _expect_dict(data)
return ArtifactUploadResponse(**data)
def upload_file(
file_path: Path,
url: str,
headers: dict[str, str] | None = None,
) -> None:
"""Upload a file to the given URL."""
try:
with file_path.open("rb") as f:
upload_response = requests.put(
url,
data=f,
timeout=60,
headers=headers,
)
upload_response.raise_for_status()
except requests.exceptions.RequestException as exc:
raise NetworkRequestError(
f"File upload to {url} failed: {exc}"
) from exc
def upload_file_to_gcs(
file_path: Path,
upload_url: str,
file_size: int | None = None,
) -> None:
"""Upload a file to the given Google Cloud Storage signed URL."""
file_size = file_size or file_path.stat().st_size
headers = {
"Content-Type": "application/octet-stream",
"x-goog-content-length-range": f"0,{file_size}",
}
upload_file(file_path, upload_url, headers=headers)
def update_artifact(
config: ApiConfig, artifact_id: str, status: ArtifactStatus
) -> None:
"""Update an artifact."""
_request(
config,
"PATCH",
f"/api/artifacts/{artifact_id}",
json={"status": status.value},
)
def delete_artifact(
config: ApiConfig,
artifact_id: str,
) -> None:
"""Delete an artifact, cleaning up storage."""
_request(
config,
"DELETE",
f"/api/artifacts/{artifact_id}",
)
def _expect_dict(data: JSONData | None) -> JSONObj:
"""Ensure data is a dict, else raise error."""
if not isinstance(data, dict):
raise RuntimeError("Unexpected response shape: expected object")
return data
def _expect_list(data: JSONData | None) -> list[JSONObj]:
"""Ensure data is a list, else raise error."""
if not isinstance(data, list):
raise RuntimeError("Unexpected response shape: expected array")
return data
def _request(
config: ApiConfig,
method: str,
url: str,
json: JSONObj | None = None,
params: dict[str, str | int] | None = None,
) -> JSONData | None:
"""Send HTTP request and handle API response."""
# Endpoint paths are written absolute ("/api/..."), which
# urljoin would resolve against the host, discarding the path
# prefix of a Hub served under one. Joining them relative to
# the base keeps that prefix.
full_url = urljoin(config.base_url, url.lstrip("/"))
try:
resp = requests.request(
method=method,
url=full_url,
headers=config.headers,
json=json,
params=params,
timeout=10,
)
except requests.exceptions.RequestException as exc:
raise NetworkRequestError(
f"Request to {full_url} failed: {exc}"
) from exc
try:
payload: JSONObj = resp.json() if resp.content else {}
except ValueError:
payload = {}
if resp.ok:
if resp.status_code == 204:
return None
if "data" in payload:
return payload["data"]
raise ApiError(resp.status_code, "Missing `data` field", [], resp)
errors = [
ApiErrorDetail.from_dict(err) for err in payload.get("errors", [])
]
if resp.status_code == 401:
message = (
"Invalid or expired API key. Generate a new one at "
f"{urljoin(config.base_url, 'profile')} and store it with "
"'embedl-hub auth --api-key -', or set the EMBEDL_HUB_API_KEY "
"environment variable."
)
else:
message = payload.get("message") or resp.reason
raise ApiError(resp.status_code, message, errors, resp)
class ErrorCode(Enum):
"""Specific error codes returned by the API."""
STORAGE_QUOTA_EXCEEDED = "storage_quota_exceeded"
FILE_TOO_LARGE = "file_too_large"
JOB_QUOTA_EXCEEDED = "job_quota_exceeded"
INVALID_RUN_LINEAGE = "invalid_run_lineage"
@dataclass
class ApiErrorDetail:
"""An individual error contained in an API error response."""
title: str | None = None
status: int | None = None
code: str | None = None
detail: str | None = None
source: dict[str, str] | None = None
meta: dict[str, Any] | None = None
@classmethod
def from_dict(cls, data: dict[str, Any]) -> Self:
return cls(
title=data.get("title"),
status=data.get("status"),
code=data.get("code"),
detail=data.get("detail"),
source=data.get("source"),
meta=data.get("meta"),
)
def __str__(self) -> str:
parts: list[str] = []
if self.title:
parts.append(self.title)
if self.code:
parts.append(f"(code: {self.code})")
if self.status:
parts.append(f"[{self.status}]")
if self.detail:
parts.append(self.detail)
if self.source:
parts.append(f"-> {self.source}")
return " ".join(parts) or "<empty>"
[docs]
class ApiError(Exception):
"""Raised for API errors with JSON body."""
def __init__(
self,
status_code: int,
message: str,
errors: list[ApiErrorDetail] | None = None,
response: requests.Response | None = None,
) -> None:
super().__init__(f"{status_code} {message}")
self.status_code = status_code
self.message = message
self.errors = errors or []
self.response = response
def __str__(self) -> str:
if not self.errors:
return super().__str__()
joined = "\n".join(map(str, self.errors))
return f"{super().__str__()}\n{joined}"
[docs]
class NetworkRequestError(Exception):
"""Raised when HTTP request fails."""
