petabyte-client 0.1.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.
- petabyte_client-0.1.0/LICENSE +26 -0
- petabyte_client-0.1.0/PKG-INFO +78 -0
- petabyte_client-0.1.0/README.md +57 -0
- petabyte_client-0.1.0/modelhub/__init__.py +25 -0
- petabyte_client-0.1.0/modelhub/__main__.py +5 -0
- petabyte_client-0.1.0/modelhub/cache.py +301 -0
- petabyte_client-0.1.0/modelhub/cli.py +463 -0
- petabyte_client-0.1.0/modelhub/download.py +226 -0
- petabyte_client-0.1.0/modelhub/hardware.py +151 -0
- petabyte_client-0.1.0/modelhub/ids.py +142 -0
- petabyte_client-0.1.0/modelhub/manager.py +244 -0
- petabyte_client-0.1.0/modelhub/manifest.py +138 -0
- petabyte_client-0.1.0/modelhub/providers/__init__.py +7 -0
- petabyte_client-0.1.0/modelhub/providers/base.py +101 -0
- petabyte_client-0.1.0/modelhub/providers/http.py +49 -0
- petabyte_client-0.1.0/modelhub/providers/huggingface.py +250 -0
- petabyte_client-0.1.0/modelhub/providers/registry.py +118 -0
- petabyte_client-0.1.0/petabyte.py +430 -0
- petabyte_client-0.1.0/petabyte_client.egg-info/PKG-INFO +78 -0
- petabyte_client-0.1.0/petabyte_client.egg-info/SOURCES.txt +24 -0
- petabyte_client-0.1.0/petabyte_client.egg-info/dependency_links.txt +1 -0
- petabyte_client-0.1.0/petabyte_client.egg-info/entry_points.txt +2 -0
- petabyte_client-0.1.0/petabyte_client.egg-info/requires.txt +1 -0
- petabyte_client-0.1.0/petabyte_client.egg-info/top_level.txt +2 -0
- petabyte_client-0.1.0/pyproject.toml +39 -0
- petabyte_client-0.1.0/setup.cfg +4 -0
|
@@ -0,0 +1,26 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Petabyte, Inc.
|
|
4
|
+
|
|
5
|
+
This license covers ONLY the `petabyte-client` distribution — the `petabyte` CLI
|
|
6
|
+
module and the bundled `modelhub` package that ship in that wheel. The rest of the
|
|
7
|
+
Petabyte repository (the server and everything under the repository-root LICENSE)
|
|
8
|
+
remains proprietary and is NOT covered by this license.
|
|
9
|
+
|
|
10
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
11
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
12
|
+
in the Software without restriction, including without limitation the rights
|
|
13
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
14
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
15
|
+
furnished to do so, subject to the following conditions:
|
|
16
|
+
|
|
17
|
+
The above copyright notice and this permission notice shall be included in all
|
|
18
|
+
copies or substantial portions of the Software.
|
|
19
|
+
|
|
20
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
21
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
22
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
23
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
24
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
25
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
26
|
+
SOFTWARE.
|
|
@@ -0,0 +1,78 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: petabyte-client
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Petabyte CLI — rent verified GPU compute, launch apps, and pull AI models from your terminal.
|
|
5
|
+
Author: Petabyte, Inc.
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://petabyte.market
|
|
8
|
+
Project-URL: Documentation, https://petabyte.market/wiki
|
|
9
|
+
Project-URL: Source, https://github.com/BDR-Pro/petabyte-client
|
|
10
|
+
Keywords: gpu,compute,cloud,marketplace,cli,ai,models
|
|
11
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
12
|
+
Classifier: Environment :: Console
|
|
13
|
+
Classifier: Intended Audience :: Developers
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
16
|
+
Requires-Python: >=3.9
|
|
17
|
+
Description-Content-Type: text/markdown
|
|
18
|
+
License-File: LICENSE
|
|
19
|
+
Requires-Dist: httpx>=0.24
|
|
20
|
+
Dynamic: license-file
|
|
21
|
+
|
|
22
|
+
<!-- Source of truth: the petabyte CLI is developed in the (private) petabyte monorepo and
|
|
23
|
+
mirrored here by scripts/build_cli_package.py. Open issues/PRs against this repo for the
|
|
24
|
+
client; the server is closed-source. -->
|
|
25
|
+
|
|
26
|
+
# Petabyte CLI & Dashboard
|
|
27
|
+
|
|
28
|
+
## CLI
|
|
29
|
+
Installed from PyPI, the `petabyte` command is a thin client (only needs `httpx` — it just talks
|
|
30
|
+
to the API over HTTPS, so it never pulls in the server):
|
|
31
|
+
```bash
|
|
32
|
+
pip install petabyte-client # the command it installs is `petabyte`
|
|
33
|
+
export PETABYTE_API_URL=https://petabyte.market # default; or pass --api / omit for localhost
|
|
34
|
+
petabyte register -u alice -p secret
|
|
35
|
+
petabyte login -u alice -p secret
|
|
36
|
+
petabyte deposit 100
|
|
37
|
+
petabyte specs # a readable, cheapest-first GPU table
|
|
38
|
+
petabyte launch ollama --hours 2 # one-click app: cheapest verified GPU, started
|
|
39
|
+
petabyte run hello.ipynb --gpu H100 --hours 1 # run a notebook/.py on a rented GPU
|
|
40
|
+
petabyte ask "explain attention" --model llama3.2 # pay-per-token Inference API (OpenAI-compatible)
|
|
41
|
+
petabyte wallet
|
|
42
|
+
```
|
|
43
|
+
`petabyte ask` uses an `inference`-scoped API key — pass `--key`, set `PETABYTE_API_KEY`, or
|
|
44
|
+
save it as `api_key` in your CLI config; the answer prints to stdout (the token/cost line to stderr).
|
|
45
|
+
Model management is included — `petabyte pull <publisher/model>`, `petabyte model list/inspect/remove`,
|
|
46
|
+
and `petabyte run <model-id>` work straight from the pip install (the model hub is pure standard
|
|
47
|
+
library, so it adds no dependency beyond httpx):
|
|
48
|
+
```bash
|
|
49
|
+
petabyte pull Qwen/Qwen3-8B # verified, resumable download into ~/.petabyte
|
|
50
|
+
petabyte model list # what's in your local cache
|
|
51
|
+
petabyte run Qwen/Qwen3-8B # start a model runtime
|
|
52
|
+
```
|
|
53
|
+
The package is built from the repo-root `pyproject.toml` (`name = "petabyte-client"`; the command
|
|
54
|
+
stays `petabyte`), which bundles this
|
|
55
|
+
CLI module plus the `modelhub` package. From a source checkout you can also run it directly with
|
|
56
|
+
`python cli/petabyte.py <cmd>`, or `pip install .` from the repo root.
|
|
57
|
+
`run` books the cheapest matching GPU, escrows funds, dispatches the notebook,
|
|
58
|
+
polls, and prints the result. `.ipynb` (code cells) and `.py` files are supported.
|
|
59
|
+
|
|
60
|
+
### Output & config
|
|
61
|
+
- **Human output:** semantic colour (green=ok, yellow=pending, red=error, cyan=info) with
|
|
62
|
+
aligned tables. Colour turns **off** automatically when stdout is not a TTY or `NO_COLOR`
|
|
63
|
+
is set — safe for scripts and CI. The buyer `petabyte` client uses its own small inline
|
|
64
|
+
colour helpers (no dependency beyond `httpx`); it does **not** import `cli_ui.py`.
|
|
65
|
+
- `PETABYTE_API_URL` (or `--api`) selects the API; `PETABYTE_CONFIG=/path/cli.json`
|
|
66
|
+
isolates the saved token/API (handy in CI or tests).
|
|
67
|
+
|
|
68
|
+
> Note: a stable `--json`/`PETABYTE_JSON` machine-readable mode and a `doctor`
|
|
69
|
+
> health-gate command exist in the **seller agent CLI** (`lumaris_agent/agent_cli.py`),
|
|
70
|
+
> not in this buyer client. The shared `cli_ui.py` presentation layer is likewise used
|
|
71
|
+
> by the agent and desktop app, not by `petabyte.py`.
|
|
72
|
+
|
|
73
|
+
## Dashboard
|
|
74
|
+
Served by the API at `/` (same-origin, no CORS setup). Start the API and open
|
|
75
|
+
`http://localhost:8000/` — live nodes/jobs/GMV stats, wallet + deposit, the GPU
|
|
76
|
+
inventory with a live $/hr-vs-AWS savings column, and one-click job runs.
|
|
77
|
+
|
|
78
|
+
Both need an attested, online seller node (run the agent) to actually execute jobs.
|
|
@@ -0,0 +1,57 @@
|
|
|
1
|
+
<!-- Source of truth: the petabyte CLI is developed in the (private) petabyte monorepo and
|
|
2
|
+
mirrored here by scripts/build_cli_package.py. Open issues/PRs against this repo for the
|
|
3
|
+
client; the server is closed-source. -->
|
|
4
|
+
|
|
5
|
+
# Petabyte CLI & Dashboard
|
|
6
|
+
|
|
7
|
+
## CLI
|
|
8
|
+
Installed from PyPI, the `petabyte` command is a thin client (only needs `httpx` — it just talks
|
|
9
|
+
to the API over HTTPS, so it never pulls in the server):
|
|
10
|
+
```bash
|
|
11
|
+
pip install petabyte-client # the command it installs is `petabyte`
|
|
12
|
+
export PETABYTE_API_URL=https://petabyte.market # default; or pass --api / omit for localhost
|
|
13
|
+
petabyte register -u alice -p secret
|
|
14
|
+
petabyte login -u alice -p secret
|
|
15
|
+
petabyte deposit 100
|
|
16
|
+
petabyte specs # a readable, cheapest-first GPU table
|
|
17
|
+
petabyte launch ollama --hours 2 # one-click app: cheapest verified GPU, started
|
|
18
|
+
petabyte run hello.ipynb --gpu H100 --hours 1 # run a notebook/.py on a rented GPU
|
|
19
|
+
petabyte ask "explain attention" --model llama3.2 # pay-per-token Inference API (OpenAI-compatible)
|
|
20
|
+
petabyte wallet
|
|
21
|
+
```
|
|
22
|
+
`petabyte ask` uses an `inference`-scoped API key — pass `--key`, set `PETABYTE_API_KEY`, or
|
|
23
|
+
save it as `api_key` in your CLI config; the answer prints to stdout (the token/cost line to stderr).
|
|
24
|
+
Model management is included — `petabyte pull <publisher/model>`, `petabyte model list/inspect/remove`,
|
|
25
|
+
and `petabyte run <model-id>` work straight from the pip install (the model hub is pure standard
|
|
26
|
+
library, so it adds no dependency beyond httpx):
|
|
27
|
+
```bash
|
|
28
|
+
petabyte pull Qwen/Qwen3-8B # verified, resumable download into ~/.petabyte
|
|
29
|
+
petabyte model list # what's in your local cache
|
|
30
|
+
petabyte run Qwen/Qwen3-8B # start a model runtime
|
|
31
|
+
```
|
|
32
|
+
The package is built from the repo-root `pyproject.toml` (`name = "petabyte-client"`; the command
|
|
33
|
+
stays `petabyte`), which bundles this
|
|
34
|
+
CLI module plus the `modelhub` package. From a source checkout you can also run it directly with
|
|
35
|
+
`python cli/petabyte.py <cmd>`, or `pip install .` from the repo root.
|
|
36
|
+
`run` books the cheapest matching GPU, escrows funds, dispatches the notebook,
|
|
37
|
+
polls, and prints the result. `.ipynb` (code cells) and `.py` files are supported.
|
|
38
|
+
|
|
39
|
+
### Output & config
|
|
40
|
+
- **Human output:** semantic colour (green=ok, yellow=pending, red=error, cyan=info) with
|
|
41
|
+
aligned tables. Colour turns **off** automatically when stdout is not a TTY or `NO_COLOR`
|
|
42
|
+
is set — safe for scripts and CI. The buyer `petabyte` client uses its own small inline
|
|
43
|
+
colour helpers (no dependency beyond `httpx`); it does **not** import `cli_ui.py`.
|
|
44
|
+
- `PETABYTE_API_URL` (or `--api`) selects the API; `PETABYTE_CONFIG=/path/cli.json`
|
|
45
|
+
isolates the saved token/API (handy in CI or tests).
|
|
46
|
+
|
|
47
|
+
> Note: a stable `--json`/`PETABYTE_JSON` machine-readable mode and a `doctor`
|
|
48
|
+
> health-gate command exist in the **seller agent CLI** (`lumaris_agent/agent_cli.py`),
|
|
49
|
+
> not in this buyer client. The shared `cli_ui.py` presentation layer is likewise used
|
|
50
|
+
> by the agent and desktop app, not by `petabyte.py`.
|
|
51
|
+
|
|
52
|
+
## Dashboard
|
|
53
|
+
Served by the API at `/` (same-origin, no CORS setup). Start the API and open
|
|
54
|
+
`http://localhost:8000/` — live nodes/jobs/GMV stats, wallet + deposit, the GPU
|
|
55
|
+
inventory with a live $/hr-vs-AWS savings column, and one-click job runs.
|
|
56
|
+
|
|
57
|
+
Both need an attested, online seller node (run the agent) to actually execute jobs.
|
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
"""modelhub — Petabyte's provider-independent model discovery / download / management layer.
|
|
2
|
+
|
|
3
|
+
Hugging Face-grade convenience (`petabyte model pull publisher/model`) with a clean local layer:
|
|
4
|
+
content-addressed cache, resumable verified downloads, hardware-aware compatibility, and a provider
|
|
5
|
+
abstraction so models can come from Hugging Face, a mirror, or a future Petabyte registry without the
|
|
6
|
+
rest of Petabyte caring which.
|
|
7
|
+
|
|
8
|
+
Pure standard library at the core (urllib) so the CLI installs light and the whole thing is testable
|
|
9
|
+
offline against a local HTTP server.
|
|
10
|
+
"""
|
|
11
|
+
from .ids import ModelRef, parse, ModelIdError
|
|
12
|
+
from .manifest import Manifest, ModelFile
|
|
13
|
+
from .cache import Cache, default_home, CachePathError
|
|
14
|
+
from .manager import ModelManager, CompatibilityError, ConcurrentPullError
|
|
15
|
+
from .hardware import detect, compatibility
|
|
16
|
+
from . import download
|
|
17
|
+
from .providers import get_provider, search_all, ProviderError, GatedError
|
|
18
|
+
|
|
19
|
+
__all__ = [
|
|
20
|
+
"ModelRef", "parse", "ModelIdError", "Manifest", "ModelFile",
|
|
21
|
+
"Cache", "default_home", "CachePathError",
|
|
22
|
+
"ModelManager", "CompatibilityError", "ConcurrentPullError",
|
|
23
|
+
"detect", "compatibility", "download",
|
|
24
|
+
"get_provider", "search_all", "ProviderError", "GatedError",
|
|
25
|
+
]
|
|
@@ -0,0 +1,301 @@
|
|
|
1
|
+
"""cache.py — a content-addressed local model cache (conceptually Git/Docker/HF).
|
|
2
|
+
|
|
3
|
+
Layout under PETABYTE_HOME (default ~/.petabyte):
|
|
4
|
+
|
|
5
|
+
blobs/sha256/<hash> deduplicated content store — one copy per distinct blob
|
|
6
|
+
manifests/<pub>/<name>/<rev>.json the normalized manifest we resolved
|
|
7
|
+
refs/<pub>/<name>/<tag> a small file naming the revision a tag points at
|
|
8
|
+
models/<pub>/<name>/<rev>/... the materialized tree a runtime consumes (links into blobs)
|
|
9
|
+
tmp/ .partial downloads + per-model locks
|
|
10
|
+
|
|
11
|
+
If two models/revisions share a blob (same sha256) it is stored once and every model dir links to
|
|
12
|
+
it. Removing a model never deletes a blob another model still references; `prune` reclaims only
|
|
13
|
+
blobs nothing points at.
|
|
14
|
+
|
|
15
|
+
Path safety is enforced everywhere a remote-supplied path becomes a filesystem path: a file path
|
|
16
|
+
from a manifest can never escape its model directory (no "..", no absolute paths, no drive letters,
|
|
17
|
+
no backslashes, no symlink games).
|
|
18
|
+
"""
|
|
19
|
+
import hashlib
|
|
20
|
+
import json
|
|
21
|
+
import os
|
|
22
|
+
import shutil
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class CachePathError(ValueError):
|
|
26
|
+
"""A remote-supplied path tried to escape the cache."""
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def default_home() -> str:
|
|
30
|
+
return os.environ.get("PETABYTE_HOME") or os.path.join(os.path.expanduser("~"), ".petabyte")
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def _safe_rel(rel: str) -> str:
|
|
34
|
+
"""Validate a manifest file path and return a normalized, forward-slash relative path that is
|
|
35
|
+
guaranteed to stay inside a model directory. Raises CachePathError on anything unsafe."""
|
|
36
|
+
if rel is None:
|
|
37
|
+
raise CachePathError("empty path")
|
|
38
|
+
p = str(rel).replace("\\", "/").strip()
|
|
39
|
+
if not p or "\x00" in p:
|
|
40
|
+
raise CachePathError(f"empty/invalid path: {rel!r}")
|
|
41
|
+
if p.startswith("/") or (len(p) >= 2 and p[1] == ":"):
|
|
42
|
+
raise CachePathError(f"absolute path not allowed: {rel!r}")
|
|
43
|
+
parts = []
|
|
44
|
+
for seg in p.split("/"):
|
|
45
|
+
if seg in ("", "."):
|
|
46
|
+
continue
|
|
47
|
+
if seg == "..":
|
|
48
|
+
raise CachePathError(f"path traversal not allowed: {rel!r}")
|
|
49
|
+
parts.append(seg)
|
|
50
|
+
if not parts:
|
|
51
|
+
raise CachePathError(f"empty path: {rel!r}")
|
|
52
|
+
return "/".join(parts)
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
class Cache:
|
|
56
|
+
def __init__(self, home=None):
|
|
57
|
+
self.home = os.path.abspath(home or default_home())
|
|
58
|
+
self.blobs = os.path.join(self.home, "blobs", "sha256")
|
|
59
|
+
self.manifests = os.path.join(self.home, "manifests")
|
|
60
|
+
self.refs = os.path.join(self.home, "refs")
|
|
61
|
+
self.models = os.path.join(self.home, "models")
|
|
62
|
+
self.tmp = os.path.join(self.home, "tmp")
|
|
63
|
+
|
|
64
|
+
def ensure(self):
|
|
65
|
+
for d in (self.blobs, self.manifests, self.refs, self.models, self.tmp):
|
|
66
|
+
os.makedirs(d, exist_ok=True)
|
|
67
|
+
|
|
68
|
+
# ---- blobs ----
|
|
69
|
+
def blob_path(self, sha256: str) -> str:
|
|
70
|
+
sha = (sha256 or "").lower()
|
|
71
|
+
if not (len(sha) == 64 and all(c in "0123456789abcdef" for c in sha)):
|
|
72
|
+
raise CachePathError(f"bad sha256: {sha256!r}")
|
|
73
|
+
return os.path.join(self.blobs, sha)
|
|
74
|
+
|
|
75
|
+
def has_blob(self, sha256: str) -> bool:
|
|
76
|
+
try:
|
|
77
|
+
return os.path.exists(self.blob_path(sha256))
|
|
78
|
+
except CachePathError:
|
|
79
|
+
return False
|
|
80
|
+
|
|
81
|
+
# ---- model dir / materialization ----
|
|
82
|
+
def model_dir(self, ref, revision) -> str:
|
|
83
|
+
parts = ref.slug_parts()
|
|
84
|
+
rev = _safe_rel(revision or "main").replace("/", "_")
|
|
85
|
+
d = os.path.join(self.models, *[_safe_rel(p) for p in parts], rev)
|
|
86
|
+
# defence in depth: the resolved path must live under self.models
|
|
87
|
+
if os.path.commonpath([os.path.abspath(d), self.models]) != self.models:
|
|
88
|
+
raise CachePathError("model dir escaped cache root")
|
|
89
|
+
return d
|
|
90
|
+
|
|
91
|
+
def materialize(self, model_dir, rel_path, *, blob_sha=None, src_file=None):
|
|
92
|
+
"""Place one file into the model dir at rel_path. If blob_sha is given, link to the shared
|
|
93
|
+
blob (symlink → hardlink → copy fallback). Otherwise move src_file into place. Returns the
|
|
94
|
+
final absolute path."""
|
|
95
|
+
rel = _safe_rel(rel_path)
|
|
96
|
+
dest = os.path.join(model_dir, rel)
|
|
97
|
+
# the joined path must stay within model_dir
|
|
98
|
+
if os.path.commonpath([os.path.abspath(dest), os.path.abspath(model_dir)]) != os.path.abspath(model_dir):
|
|
99
|
+
raise CachePathError(f"path escaped model dir: {rel_path!r}")
|
|
100
|
+
os.makedirs(os.path.dirname(dest), exist_ok=True)
|
|
101
|
+
if os.path.islink(dest) or os.path.exists(dest):
|
|
102
|
+
os.remove(dest)
|
|
103
|
+
if blob_sha is not None:
|
|
104
|
+
blob = self.blob_path(blob_sha)
|
|
105
|
+
rel_to_blob = os.path.relpath(blob, os.path.dirname(dest))
|
|
106
|
+
try:
|
|
107
|
+
os.symlink(rel_to_blob, dest)
|
|
108
|
+
except (OSError, NotImplementedError):
|
|
109
|
+
try:
|
|
110
|
+
os.link(blob, dest)
|
|
111
|
+
except OSError:
|
|
112
|
+
shutil.copy2(blob, dest)
|
|
113
|
+
elif src_file is not None:
|
|
114
|
+
shutil.move(src_file, dest)
|
|
115
|
+
return dest
|
|
116
|
+
|
|
117
|
+
# ---- manifests + refs ----
|
|
118
|
+
def _manifest_path(self, ref, revision) -> str:
|
|
119
|
+
parts = [_safe_rel(p) for p in ref.slug_parts()]
|
|
120
|
+
rev = _safe_rel(revision).replace("/", "_")
|
|
121
|
+
return os.path.join(self.manifests, *parts, rev + ".json")
|
|
122
|
+
|
|
123
|
+
def write_manifest(self, manifest):
|
|
124
|
+
path = self._manifest_path_from_manifest(manifest)
|
|
125
|
+
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
126
|
+
tmp = path + ".tmp"
|
|
127
|
+
with open(tmp, "w") as f:
|
|
128
|
+
f.write(manifest.to_json())
|
|
129
|
+
os.replace(tmp, path)
|
|
130
|
+
|
|
131
|
+
def _manifest_path_from_manifest(self, manifest):
|
|
132
|
+
pub = manifest.publisher
|
|
133
|
+
name = manifest.name or manifest.id.split("/")[-1]
|
|
134
|
+
parts = [_safe_rel(x) for x in ([pub] if pub else []) + [name]]
|
|
135
|
+
rev = _safe_rel(manifest.revision).replace("/", "_")
|
|
136
|
+
return os.path.join(self.manifests, *parts, rev + ".json")
|
|
137
|
+
|
|
138
|
+
def read_manifest(self, ref, revision):
|
|
139
|
+
from .manifest import Manifest
|
|
140
|
+
path = self._manifest_path(ref, revision)
|
|
141
|
+
if not os.path.exists(path):
|
|
142
|
+
return None
|
|
143
|
+
return Manifest.from_dict(json.load(open(path)))
|
|
144
|
+
|
|
145
|
+
def write_ref(self, ref, revision):
|
|
146
|
+
parts = [_safe_rel(p) for p in ref.slug_parts()]
|
|
147
|
+
tag = _safe_rel(ref.tag or "latest")
|
|
148
|
+
d = os.path.join(self.refs, *parts)
|
|
149
|
+
os.makedirs(d, exist_ok=True)
|
|
150
|
+
with open(os.path.join(d, tag), "w") as f:
|
|
151
|
+
f.write(revision)
|
|
152
|
+
|
|
153
|
+
def read_ref(self, ref):
|
|
154
|
+
parts = [_safe_rel(p) for p in ref.slug_parts()]
|
|
155
|
+
tag = _safe_rel(ref.tag or "latest")
|
|
156
|
+
p = os.path.join(self.refs, *parts, tag)
|
|
157
|
+
return open(p).read().strip() if os.path.exists(p) else None
|
|
158
|
+
|
|
159
|
+
# ---- listing / inspection ----
|
|
160
|
+
def installed(self) -> list:
|
|
161
|
+
"""Every installed (materialized + manifested) model revision."""
|
|
162
|
+
from .manifest import Manifest
|
|
163
|
+
out = []
|
|
164
|
+
for root, _dirs, files in os.walk(self.manifests):
|
|
165
|
+
for fn in files:
|
|
166
|
+
if not fn.endswith(".json"):
|
|
167
|
+
continue
|
|
168
|
+
try:
|
|
169
|
+
m = Manifest.from_dict(json.load(open(os.path.join(root, fn))))
|
|
170
|
+
except Exception: # noqa: BLE001
|
|
171
|
+
continue
|
|
172
|
+
from .ids import ModelRef
|
|
173
|
+
ref = ModelRef(m.publisher, m.name or m.id.split("/")[-1], source=m.source)
|
|
174
|
+
mdir = self.model_dir(ref, m.revision)
|
|
175
|
+
out.append({"id": m.id, "revision": m.revision, "source": m.source,
|
|
176
|
+
"format": m.format, "parameters": m.parameters,
|
|
177
|
+
"total_size": m.total_size, "license": m.license,
|
|
178
|
+
"path": mdir, "installed": os.path.isdir(mdir),
|
|
179
|
+
"requirements": m.requirements})
|
|
180
|
+
out.sort(key=lambda x: x["id"])
|
|
181
|
+
return out
|
|
182
|
+
|
|
183
|
+
def _referenced_shas(self) -> set:
|
|
184
|
+
"""sha256 of every blob referenced by a manifest whose model dir still exists (live)."""
|
|
185
|
+
from .manifest import Manifest
|
|
186
|
+
from .ids import ModelRef
|
|
187
|
+
live = set()
|
|
188
|
+
for root, _dirs, files in os.walk(self.manifests):
|
|
189
|
+
for fn in files:
|
|
190
|
+
if not fn.endswith(".json"):
|
|
191
|
+
continue
|
|
192
|
+
try:
|
|
193
|
+
m = Manifest.from_dict(json.load(open(os.path.join(root, fn))))
|
|
194
|
+
except Exception: # noqa: BLE001
|
|
195
|
+
continue
|
|
196
|
+
ref = ModelRef(m.publisher, m.name or m.id.split("/")[-1], source=m.source)
|
|
197
|
+
if os.path.isdir(self.model_dir(ref, m.revision)):
|
|
198
|
+
live.update(f.sha256 for f in m.files if f.sha256)
|
|
199
|
+
return live
|
|
200
|
+
|
|
201
|
+
def status(self) -> dict:
|
|
202
|
+
total, blob_count, blob_bytes, reclaimable = 0, 0, 0, 0
|
|
203
|
+
referenced = self._referenced_shas()
|
|
204
|
+
ref_count = {}
|
|
205
|
+
# count how many live manifests reference each sha (for "shared")
|
|
206
|
+
from .manifest import Manifest
|
|
207
|
+
from .ids import ModelRef
|
|
208
|
+
for root, _dirs, files in os.walk(self.manifests):
|
|
209
|
+
for fn in files:
|
|
210
|
+
if not fn.endswith(".json"):
|
|
211
|
+
continue
|
|
212
|
+
try:
|
|
213
|
+
m = Manifest.from_dict(json.load(open(os.path.join(root, fn))))
|
|
214
|
+
except Exception: # noqa: BLE001
|
|
215
|
+
continue
|
|
216
|
+
r = ModelRef(m.publisher, m.name or m.id.split("/")[-1], source=m.source)
|
|
217
|
+
if not os.path.isdir(self.model_dir(r, m.revision)):
|
|
218
|
+
continue
|
|
219
|
+
for f in m.files:
|
|
220
|
+
if f.sha256:
|
|
221
|
+
ref_count[f.sha256] = ref_count.get(f.sha256, 0) + 1
|
|
222
|
+
shared = 0
|
|
223
|
+
if os.path.isdir(self.blobs):
|
|
224
|
+
for fn in os.listdir(self.blobs):
|
|
225
|
+
fp = os.path.join(self.blobs, fn)
|
|
226
|
+
if not os.path.isfile(fp):
|
|
227
|
+
continue
|
|
228
|
+
sz = os.path.getsize(fp)
|
|
229
|
+
blob_count += 1
|
|
230
|
+
blob_bytes += sz
|
|
231
|
+
total += sz
|
|
232
|
+
if fn not in referenced:
|
|
233
|
+
reclaimable += sz
|
|
234
|
+
elif ref_count.get(fn, 0) > 1:
|
|
235
|
+
shared += 1
|
|
236
|
+
# non-blob model files (configs/tokenizers stored directly)
|
|
237
|
+
model_extra = 0
|
|
238
|
+
for root, _dirs, files in os.walk(self.models):
|
|
239
|
+
for fn in files:
|
|
240
|
+
fp = os.path.join(root, fn)
|
|
241
|
+
if os.path.islink(fp):
|
|
242
|
+
continue
|
|
243
|
+
try:
|
|
244
|
+
model_extra += os.path.getsize(fp)
|
|
245
|
+
except OSError:
|
|
246
|
+
pass
|
|
247
|
+
total += model_extra
|
|
248
|
+
installed = self.installed()
|
|
249
|
+
return {"home": self.home, "total_bytes": total, "blob_bytes": blob_bytes,
|
|
250
|
+
"blob_count": blob_count, "shared_blobs": shared,
|
|
251
|
+
"reclaimable_bytes": reclaimable, "model_extra_bytes": model_extra,
|
|
252
|
+
"models": len(installed), "installed": installed}
|
|
253
|
+
|
|
254
|
+
def prune(self, dry_run=False) -> dict:
|
|
255
|
+
"""Delete blobs nothing references. Never touches a referenced blob."""
|
|
256
|
+
referenced = self._referenced_shas()
|
|
257
|
+
removed, freed = [], 0
|
|
258
|
+
if os.path.isdir(self.blobs):
|
|
259
|
+
for fn in list(os.listdir(self.blobs)):
|
|
260
|
+
fp = os.path.join(self.blobs, fn)
|
|
261
|
+
if not os.path.isfile(fp) or fn in referenced:
|
|
262
|
+
continue
|
|
263
|
+
freed += os.path.getsize(fp)
|
|
264
|
+
removed.append(fn)
|
|
265
|
+
if not dry_run:
|
|
266
|
+
os.remove(fp)
|
|
267
|
+
return {"removed": removed, "removed_count": len(removed), "freed_bytes": freed,
|
|
268
|
+
"dry_run": dry_run}
|
|
269
|
+
|
|
270
|
+
def remove(self, ref, revision) -> dict:
|
|
271
|
+
"""Remove one model revision (model dir + manifest + tag ref). Blobs are left for prune."""
|
|
272
|
+
removed = []
|
|
273
|
+
mdir = self.model_dir(ref, revision)
|
|
274
|
+
if os.path.isdir(mdir):
|
|
275
|
+
shutil.rmtree(mdir)
|
|
276
|
+
removed.append("model")
|
|
277
|
+
mpath = self._manifest_path(ref, revision)
|
|
278
|
+
if os.path.exists(mpath):
|
|
279
|
+
os.remove(mpath)
|
|
280
|
+
removed.append("manifest")
|
|
281
|
+
# drop any tag ref pointing at this revision
|
|
282
|
+
parts = [_safe_rel(p) for p in ref.slug_parts()]
|
|
283
|
+
rdir = os.path.join(self.refs, *parts)
|
|
284
|
+
if os.path.isdir(rdir):
|
|
285
|
+
for tag in os.listdir(rdir):
|
|
286
|
+
tp = os.path.join(rdir, tag)
|
|
287
|
+
try:
|
|
288
|
+
if open(tp).read().strip() == revision:
|
|
289
|
+
os.remove(tp)
|
|
290
|
+
removed.append(f"ref:{tag}")
|
|
291
|
+
except OSError:
|
|
292
|
+
pass
|
|
293
|
+
return {"removed": removed, "id": ref.id, "revision": revision}
|
|
294
|
+
|
|
295
|
+
@staticmethod
|
|
296
|
+
def sha256_file(path, chunk=1024 * 1024) -> str:
|
|
297
|
+
h = hashlib.sha256()
|
|
298
|
+
with open(path, "rb") as f:
|
|
299
|
+
for block in iter(lambda: f.read(chunk), b""):
|
|
300
|
+
h.update(block)
|
|
301
|
+
return h.hexdigest()
|