Spaces:
Running
Running
File size: 5,056 Bytes
650ee43 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | from typing import Iterator
from trackio.registry import Registry
from trackio.sqlite_storage import SQLiteStorage
class Run:
def __init__(self, project: str, name: str, run_id: str | None = None):
self.project = project
self.name = name
self._id = run_id or name
self._config = None
@property
def id(self) -> str:
return self._id
@property
def config(self) -> dict | None:
if self._config is None:
self._config = SQLiteStorage.get_run_config(
self.project, self.name, run_id=self.id
)
return self._config
def alerts(self, level: str | None = None, since: str | None = None) -> list[dict]:
return SQLiteStorage.get_alerts(
self.project, run_name=self.name, run_id=self.id, level=level, since=since
)
def delete(self) -> bool:
return SQLiteStorage.delete_run(self.project, self.name, run_id=self.id)
def move(self, new_project: str) -> bool:
success = SQLiteStorage.move_run(
self.project, self.name, new_project, run_id=self.id
)
if success:
self.project = new_project
return success
def rename(self, new_name: str) -> "Run":
SQLiteStorage.rename_run(self.project, self.name, new_name, run_id=self.id)
self.name = new_name
return self
def __repr__(self) -> str:
return f"<Run {self.name} in project {self.project}>"
class Runs:
def __init__(self, project: str):
self.project = project
self._runs = None
def _load_runs(self):
if self._runs is None:
records = SQLiteStorage.get_run_records(self.project)
self._runs = [
Run(
self.project,
str(record["name"]),
run_id=str(record["id"]) if record["id"] is not None else None,
)
for record in records
]
def __iter__(self) -> Iterator[Run]:
self._load_runs()
return iter(self._runs)
def __getitem__(self, index: int) -> Run:
self._load_runs()
return self._runs[index]
def __len__(self) -> int:
self._load_runs()
return len(self._runs)
def __repr__(self) -> str:
self._load_runs()
return f"<Runs project={self.project} count={len(self._runs)}>"
class Api:
def runs(self, project: str) -> Runs:
if not SQLiteStorage.get_project_db_path(project).exists():
raise ValueError(f"Project '{project}' does not exist")
return Runs(project)
def alerts(
self,
project: str,
run: str | None = None,
level: str | None = None,
since: str | None = None,
) -> list[dict]:
if not SQLiteStorage.get_project_db_path(project).exists():
raise ValueError(f"Project '{project}' does not exist")
return SQLiteStorage.get_alerts(project, run_name=run, level=level, since=since)
def create_registry(
self,
name: str,
description: str | None = None,
bucket_id: str | None = None,
) -> Registry:
"""Create a new registry and return a handle on it.
Raises `ValueError` if a registry with this name already exists.
Registries are never created implicitly: linking into a registry
that does not exist raises an error.
Args:
name (`str`):
Registry name, e.g. `"models"`. Must match
`^[A-Za-z0-9_-]+$`.
description (`str`, *optional*):
Human-readable description of the registry.
bucket_id (`str`, *optional*):
Hugging Face bucket to hold the registry, e.g.
`"my-org/models-registry"`. The bucket is created (private) if
it does not exist. A bucket-backed registry is reachable from
any machine with access to it, including runs that log to a
Space. Omit for a local registry; defaults to
`TRACKIO_REGISTRY_BUCKET_ID` when that is set.
Returns:
A [`Registry`] handle on the new registry.
"""
registry = Registry(name, bucket_id=bucket_id)
registry._storage.create_registry(name, description=description)
return registry
def registry(self, name: str, bucket_id: str | None = None) -> Registry:
"""Fetch a handle on an existing registry.
Pass `bucket_id` (or set `TRACKIO_REGISTRY_BUCKET_ID`) for a
bucket-backed registry. Raises `ValueError` if no registry with this
name exists."""
registry = Registry(name, bucket_id=bucket_id)
if not registry._storage.registry_exists(name):
where = (
""
if registry.bucket_id is None
else f" in bucket '{registry.bucket_id}'"
)
raise ValueError(f"Registry '{name}' does not exist{where}")
return registry
|