argus-cortex 0.4.2__tar.gz → 0.4.4__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.
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/CHANGELOG.md +43 -1
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/CLAUDE.md +1 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/PKG-INFO +16 -1
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/README.md +12 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/pyproject.toml +4 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/_version.py +2 -2
- argus_cortex-0.4.4/src/argus_cortex/gpu.py +242 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/moderation/singguard.py +3 -1
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/store/lineage.py +15 -5
- argus_cortex-0.4.4/tests/test_gpu.py +227 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_moderation.py +18 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_store.py +5 -3
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_store_pg_integration.py +14 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/.copier-answers.yml +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/.env.example +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/.github/workflows/ci.yml +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/.github/workflows/release.yml +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/.gitignore +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/LICENSE +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/Makefile +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/__init__.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/backends.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/moderation/__init__.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/moderation/_http.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/moderation/nudenet.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/moderation/policies.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/py.typed +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/server.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/store/__init__.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/store/blob.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/store/config.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/store/errors.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/store/models.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/store/schema.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/store/vector.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/taxonomy.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/wire.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/xmp.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_backends.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_blob.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_policies.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_server.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_smoke.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_taxonomy.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_vector.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_wire.py +0 -0
- {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_xmp.py +0 -0
|
@@ -7,6 +7,46 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|
|
7
7
|
|
|
8
8
|
## [Unreleased]
|
|
9
9
|
|
|
10
|
+
## [0.4.4] - 2026-10-31
|
|
11
|
+
|
|
12
|
+
### Fixed
|
|
13
|
+
- **`SingGuardClient` keeps a longer writeup.** `to_result` clipped `explanation` at 4000
|
|
14
|
+
characters, which cut cascade/slow reasoning mid-sentence. The cap is now 16384, enough
|
|
15
|
+
for a 2048-token slow writeup plus a cascade preface.
|
|
16
|
+
- **Re-recording an asset no longer erases what another stage recorded on it.** `record_asset`'s
|
|
17
|
+
upsert on `sha256` replaced `immich_id` and `metadata` outright, so argus-lens recording a caption
|
|
18
|
+
for an image argus-quarry had already recorded (licence, author, source) wiped that provenance,
|
|
19
|
+
and a call without an `immich_id` cleared the stored one. `metadata` now merges key by key
|
|
20
|
+
(`jsonb ||`: the new call's keys win, others are kept) and a `NULL` `immich_id` keeps the stored
|
|
21
|
+
one. `uri` is still last-write-wins: it is the asset's latest location. A caller that needs to
|
|
22
|
+
remove a metadata key can no longer do it through `record_asset`.
|
|
23
|
+
|
|
24
|
+
## [0.4.3] - 2026-10-30
|
|
25
|
+
|
|
26
|
+
### Added
|
|
27
|
+
- **`argus_cortex.gpu`**: GPU tenancy shared across the stages, generalised
|
|
28
|
+
from argus-lens's `gpu` coordinator.
|
|
29
|
+
- `GpuCoordinator`: the protocol. `lease(caller=, min_vram_mb=)` is a
|
|
30
|
+
context manager that holds a GPU slot for its body. `GpuLeaseTimeout` is a
|
|
31
|
+
`TimeoutError`.
|
|
32
|
+
- `NullCoordinator`: no gating, and the default.
|
|
33
|
+
- `LeaseCoordinator`: an exclusive `flock` that every argus process on the
|
|
34
|
+
host takes, then, while still holding it, a wait until `free_vram_mb()`
|
|
35
|
+
reports `min_vram_mb`. Unlike lens's lease, `min_vram_mb` is honoured. NVML
|
|
36
|
+
counts every process on the card, so ComfyUI and Ollama are accounted for
|
|
37
|
+
without taking part. One deadline covers both waits. A probe that cannot
|
|
38
|
+
run (no driver, no `[gpu]` extra, a CPU host) lets the lease through.
|
|
39
|
+
- `free_vram_mb(device)`: NVML free memory in MB, or `None`.
|
|
40
|
+
- `build_coordinator()` / `coordinator_from_env()`: `ARGUS_GPU_COORDINATOR`
|
|
41
|
+
(`none` / `lease`), `ARGUS_GPU_LEASE_PATH`, `ARGUS_GPU_LEASE_TIMEOUT_S`
|
|
42
|
+
(300 s) and `ARGUS_GPU_DEVICE`. These are the names lens already reads. An
|
|
43
|
+
unusable number logs a warning and falls back to the default.
|
|
44
|
+
- The default lock is `<tempdir>/argus-gpu.lock`. Containers have their own
|
|
45
|
+
`/tmp`, so a shared lock needs `ARGUS_GPU_LEASE_PATH` on a volume every
|
|
46
|
+
GPU service mounts. The lock file is `chmod 0o666` after create so a
|
|
47
|
+
restrictive umask does not block other container UIDs from opening it.
|
|
48
|
+
- `[gpu]` extra: `nvidia-ml-py`, imported only when probed.
|
|
49
|
+
|
|
10
50
|
## [0.4.2] - 2026-10-29
|
|
11
51
|
|
|
12
52
|
### Added
|
|
@@ -138,7 +178,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|
|
138
178
|
(`LocalBackend`, `RemoteBackend`, `RemoteProvider`, `with_retries`
|
|
139
179
|
retrying only transient errors, `resolve_device`).
|
|
140
180
|
|
|
141
|
-
[Unreleased]: https://github.com/smk762/argus-cortex/compare/v0.4.
|
|
181
|
+
[Unreleased]: https://github.com/smk762/argus-cortex/compare/v0.4.4...HEAD
|
|
182
|
+
[0.4.4]: https://github.com/smk762/argus-cortex/compare/v0.4.3...v0.4.4
|
|
183
|
+
[0.4.3]: https://github.com/smk762/argus-cortex/compare/v0.4.2...v0.4.3
|
|
142
184
|
[0.4.2]: https://github.com/smk762/argus-cortex/compare/v0.4.1...v0.4.2
|
|
143
185
|
[0.4.1]: https://github.com/smk762/argus-cortex/compare/v0.4.0...v0.4.1
|
|
144
186
|
[0.4.0]: https://github.com/smk762/argus-cortex/compare/v0.3.0...v0.4.0
|
|
@@ -18,6 +18,7 @@ Because everything here is imported by the downstream stages, **a change ripples
|
|
|
18
18
|
- `wire.py` — the versioned wire-schema toolkit consumers build their contracts on: `make_versioned_base(field_name, current, majors)` (a Pydantic base that stamps + version-checks a field like `proof_version`/`manifest_version`), `check_version()` (the major-compatibility gate), `render_schema()`/`wire_schema()` (the canonical committed-schema string behind each consumer's `schema --check` CLI), and `VersionError`.
|
|
19
19
|
- `backends.py` — the AI-backend contract: `Backend`/`LocalBackend`/`RemoteBackend` (point at a hosted service via `from_host(host, port)` or `from_provider(RemoteProvider.…)`), plus `with_retries()` and `resolve_device()`. httpx is lazy, behind `[remote]`.
|
|
20
20
|
- `server.py` — the shared FastAPI/ASGI scaffolding other repos' micro-servers mount: `WriteGuard` (pure-ASGI method gate), the `constant_refuse()` (read-only/replay) and `cross_site_refuse()` (cross-origin write) predicates, `env_flag()`, and the `UNSAFE_METHODS`/`TRUTHY`/`FALSY` constants. Starlette is behind `[server]`.
|
|
21
|
+
- `gpu.py` — GPU tenancy (from argus-lens): the `GpuCoordinator` protocol (`lease(caller=, min_vram_mb=)`), `NullCoordinator`, `LeaseCoordinator` (host-wide `flock` **then** a wait for NVML free VRAM, so tenants outside the lock like ComfyUI and Ollama are accounted for), `free_vram_mb()`, `coordinator_from_env()` (`ARGUS_GPU_*` — host-wide names on purpose: every process must agree). NVML lazily, behind `[gpu]`; a probe that cannot run grants the lease. **The lock file must be on a volume every GPU container mounts** — per-container `/tmp` coordinates nothing.
|
|
21
22
|
- `xmp.py` — the suite's XMP sidecar read/merge/write (`write_xmp`/`read_xmp`/`merge_xmp`/`parse_xmp`, `XmpData`/`XmpRegions`, `XmpError`), used by lens (keywords/caption) and curator (faces, content rating). Stdlib `xml.etree` only. **Writes merge — latest wins per property, foreign properties preserved; keep that.** `None` = leave untouched, empty = remove. A DTD is refused and an unparseable sidecar raises rather than being overwritten. The `argus` namespace URI is a wire contract — don't change it.
|
|
22
23
|
- `moderation/` — content-moderation clients over HTTP: `nudenet.py` (`NudeNetClient`, `NsfwResult`, `DEFAULT_EXPOSED_LABELS` — **the suite's one definition of nudity; don't widen it casually**, every stage's NSFW numbers move with it), `singguard.py` (`SingGuardClient` incl. text-only `check_text()`, `GuardResult`, `rule_matches`), `policies.py` (the named policy texts, the per-policy illegal-rule table and `mode()`: `default`/`sfw-*` are SFW — **a policy's text and its illegal letters change together**), `_http.py` (shared POST, proxy-free opener, `CheckFailed`/`ModerationConfigError`). Stdlib + pydantic; Pillow lazily, behind `[images]`. Never import NudeNet (AGPL) — HTTP only. URLs are passed in; stages own their env config.
|
|
23
24
|
- `store/` — the optional stateful-services layer, opt-in per `CORTEX_*` env var: `lineage.py` (Phase 1 Postgres DAG, `open_lineage_store()`), `vector.py` (Phase 2 Qdrant, `open_vector_store()`), `blob.py` (Phase 3 MinIO/S3, `open_blob_store()`), over `config.py`/`models.py`/`errors.py`. Each `open_*()` returns a `Null*` no-op when its URL is unset. Drivers are behind `[postgres]`/`[qdrant]`/`[s3]`.
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: argus-cortex
|
|
3
|
-
Version: 0.4.
|
|
3
|
+
Version: 0.4.4
|
|
4
4
|
Summary: Shared foundation for the Argus suite: taxonomy, versioned wire-schema tooling, and the local/remote AI backend contract
|
|
5
5
|
Project-URL: Repository, https://github.com/smk762/argus-cortex
|
|
6
6
|
Author: smk762
|
|
@@ -20,6 +20,7 @@ Requires-Dist: certifi>=2026.10.19; extra == 'dev'
|
|
|
20
20
|
Requires-Dist: httpx2>=2.7; extra == 'dev'
|
|
21
21
|
Requires-Dist: httpx>=0.27; extra == 'dev'
|
|
22
22
|
Requires-Dist: minio>=7.2; extra == 'dev'
|
|
23
|
+
Requires-Dist: nvidia-ml-py>=12.535; extra == 'dev'
|
|
23
24
|
Requires-Dist: pillow>=10.0; extra == 'dev'
|
|
24
25
|
Requires-Dist: psycopg[binary,pool]>=3.1; extra == 'dev'
|
|
25
26
|
Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
|
|
@@ -27,6 +28,8 @@ Requires-Dist: pytest>=8.0; extra == 'dev'
|
|
|
27
28
|
Requires-Dist: qdrant-client>=1.10; extra == 'dev'
|
|
28
29
|
Requires-Dist: ruff==0.15.16; extra == 'dev'
|
|
29
30
|
Requires-Dist: starlette>=0.37; extra == 'dev'
|
|
31
|
+
Provides-Extra: gpu
|
|
32
|
+
Requires-Dist: nvidia-ml-py>=12.535; extra == 'gpu'
|
|
30
33
|
Provides-Extra: images
|
|
31
34
|
Requires-Dist: pillow>=10.0; extra == 'images'
|
|
32
35
|
Provides-Extra: postgres
|
|
@@ -58,6 +61,7 @@ Part of the [Argus suite](https://github.com/smk762?tab=repositories&q=argus) (q
|
|
|
58
61
|
- **`argus_cortex.server`** — shared write-guard + env-flag scaffolding for the suite's FastAPI micro-servers: `WriteGuard` (a pure-ASGI method-gate that never buffers streaming responses), the `constant_refuse()` / `cross_site_refuse()` predicates (read-only/replay mode and cross-site-write protection are two configs of the one guard), the `UNSAFE_METHODS` / `TRUTHY` / `FALSY` constants, and `env_flag()` (warns on a mistyped protection flag rather than silently disabling it). Starlette lives behind the `[server]` extra.
|
|
59
62
|
- **`argus_cortex.xmp`** — the suite's one XMP sidecar implementation (hoisted from argus-lens): `write_xmp()` / `read_xmp()` / `merge_xmp()` / `parse_xmp()` over an `XmpData` model covering `dc` (keywords, description), `xmp` (Rating, Label), `lr` (hierarchical keywords), `Iptc4xmpExt` (PersonInImage), `mwg-rs` face/pet regions, and `argus:ContentRating`. Writes **merge**: each property you set replaces that property (latest write wins), everything else in the sidecar — Lightroom develop settings, GPS, another stage's fields — is preserved. `None` leaves a property alone, an empty value removes it. Writes are atomic; an unparseable existing sidecar raises `XmpError` instead of being clobbered. Stdlib only.
|
|
60
63
|
- **`argus_cortex.moderation`** — the suite's content-moderation clients, over HTTP to servers the deployer runs (nothing here imports a detector): `NudeNetClient` (a NudeNet server's `/infer` → `NsfwResult`: detections, per-class maxima, and an **exposed score**) and `SingGuardClient` (a SingGuard server's `/api/guard/check` → `GuardResult`: verdict, the policy rules that fired, explanation; optional custom policy text). `DEFAULT_EXPOSED_LABELS` is the suite's **one definition of nudity**: the five exposed buttocks/breast/genitalia/anus classes — faces, covered parts, and bare feet, belly, armpits or male chest never count. Stdlib + pydantic; Pillow (for SingGuard's downscaling of large images, on by default) is behind `[images]`: without it `SingGuardClient(...)` raises `ModerationConfigError` at construction unless given `downscale=False`. A `user:password@` in a URL is sent as HTTP Basic auth. URLs are passed in, never read from the environment here: each stage owns its config. NudeNet itself is AGPL-3.0, which is why it stays behind an HTTP boundary.
|
|
64
|
+
- **`argus_cortex.gpu`**: GPU tenancy, generalised from argus-lens. `coordinator_from_env()` returns a coordinator whose `lease(caller=, min_vram_mb=)` holds a GPU slot for a block of work. `none` (the default) grants every lease at once. `lease` takes an exclusive `flock` that every argus process on the host shares (one lock for the host, not per GPU), then waits until `free_vram_mb()` (NVML, behind `[gpu]`) reports `min_vram_mb` free across *every* process on the card, ComfyUI and Ollama included, which never take the lock. It raises `GpuLeaseTimeout` after `ARGUS_GPU_LEASE_TIMEOUT_S` (300 s). The VRAM check is a snapshot, not a reservation, and a probe that cannot run lets the lease through. Containers each have their own `/tmp`, so set `ARGUS_GPU_LEASE_PATH` to a file on a volume every GPU service mounts, or the lock coordinates nothing.
|
|
61
65
|
- **`argus_cortex.store`** — the optional **stateful services layer** (the suite's persistent "memory"). External and opt-in — every store is configured by `CORTEX_*` env vars and no-ops when its URL is unset.
|
|
62
66
|
- **Phase 1 — Postgres lineage store.** Persists the `source_asset → caption → human_edit → dataset_membership → training_run` DAG for LoRA reproducibility, edit capture, and the feedback loop. Backed by a psycopg connection **pool** (concurrent `asyncio.to_thread` calls are safe; a dropped connection is replaced), with a `transaction()` boundary for atomic multi-step writes. `open_lineage_store()` returns a no-op store when `CORTEX_PG_URL` is unset; psycopg lives behind the `[postgres]` extra.
|
|
63
67
|
- **Phase 2 — Qdrant vector store.** Stores image + tag-set embeddings for retrieval-augmented few-shot and near-duplicate curation; put a `caption_id`/`asset_id` in a point's payload and a search hit joins straight back to the lineage. `open_vector_store()` no-ops when `CORTEX_QDRANT_URL` is unset; qdrant-client lives behind the `[qdrant]` extra. cortex stores vectors it's *handed* — computing embeddings is the caller's job.
|
|
@@ -119,6 +123,17 @@ policies.mode("adult-granular"), policies.mode("default") # "nsfw", "sfw"
|
|
|
119
123
|
|
|
120
124
|
Stages ask for a policy by name; `policies` holds each one's text, the rules that mean illegal content under it, and its SFW/NSFW mode (`default` and `sfw-*` are SFW).
|
|
121
125
|
|
|
126
|
+
```python
|
|
127
|
+
from argus_cortex.gpu import GpuLeaseTimeout, coordinator_from_env
|
|
128
|
+
|
|
129
|
+
gpu = coordinator_from_env() # ARGUS_GPU_COORDINATOR=lease, ARGUS_GPU_LEASE_PATH=/shared/argus-gpu.lock
|
|
130
|
+
try:
|
|
131
|
+
with gpu.lease(caller="faces", min_vram_mb=1500): # one argus job at a time, and 1.5 GB free
|
|
132
|
+
run_the_model()
|
|
133
|
+
except GpuLeaseTimeout:
|
|
134
|
+
... # no slot in time: back off, or answer 503
|
|
135
|
+
```
|
|
136
|
+
|
|
122
137
|
Each `check()` raises `CheckFailed` (`NudeNetCheckFailed` / `GuardCheckFailed`) for a network error or an unusable answer, so a caller checking many images can record "unchecked" against the one image — and should treat unchecked as *not* safe. Bad configuration raises `ModerationConfigError`, whose message never echoes the URL.
|
|
123
138
|
|
|
124
139
|
```python
|
|
@@ -12,6 +12,7 @@ Part of the [Argus suite](https://github.com/smk762?tab=repositories&q=argus) (q
|
|
|
12
12
|
- **`argus_cortex.server`** — shared write-guard + env-flag scaffolding for the suite's FastAPI micro-servers: `WriteGuard` (a pure-ASGI method-gate that never buffers streaming responses), the `constant_refuse()` / `cross_site_refuse()` predicates (read-only/replay mode and cross-site-write protection are two configs of the one guard), the `UNSAFE_METHODS` / `TRUTHY` / `FALSY` constants, and `env_flag()` (warns on a mistyped protection flag rather than silently disabling it). Starlette lives behind the `[server]` extra.
|
|
13
13
|
- **`argus_cortex.xmp`** — the suite's one XMP sidecar implementation (hoisted from argus-lens): `write_xmp()` / `read_xmp()` / `merge_xmp()` / `parse_xmp()` over an `XmpData` model covering `dc` (keywords, description), `xmp` (Rating, Label), `lr` (hierarchical keywords), `Iptc4xmpExt` (PersonInImage), `mwg-rs` face/pet regions, and `argus:ContentRating`. Writes **merge**: each property you set replaces that property (latest write wins), everything else in the sidecar — Lightroom develop settings, GPS, another stage's fields — is preserved. `None` leaves a property alone, an empty value removes it. Writes are atomic; an unparseable existing sidecar raises `XmpError` instead of being clobbered. Stdlib only.
|
|
14
14
|
- **`argus_cortex.moderation`** — the suite's content-moderation clients, over HTTP to servers the deployer runs (nothing here imports a detector): `NudeNetClient` (a NudeNet server's `/infer` → `NsfwResult`: detections, per-class maxima, and an **exposed score**) and `SingGuardClient` (a SingGuard server's `/api/guard/check` → `GuardResult`: verdict, the policy rules that fired, explanation; optional custom policy text). `DEFAULT_EXPOSED_LABELS` is the suite's **one definition of nudity**: the five exposed buttocks/breast/genitalia/anus classes — faces, covered parts, and bare feet, belly, armpits or male chest never count. Stdlib + pydantic; Pillow (for SingGuard's downscaling of large images, on by default) is behind `[images]`: without it `SingGuardClient(...)` raises `ModerationConfigError` at construction unless given `downscale=False`. A `user:password@` in a URL is sent as HTTP Basic auth. URLs are passed in, never read from the environment here: each stage owns its config. NudeNet itself is AGPL-3.0, which is why it stays behind an HTTP boundary.
|
|
15
|
+
- **`argus_cortex.gpu`**: GPU tenancy, generalised from argus-lens. `coordinator_from_env()` returns a coordinator whose `lease(caller=, min_vram_mb=)` holds a GPU slot for a block of work. `none` (the default) grants every lease at once. `lease` takes an exclusive `flock` that every argus process on the host shares (one lock for the host, not per GPU), then waits until `free_vram_mb()` (NVML, behind `[gpu]`) reports `min_vram_mb` free across *every* process on the card, ComfyUI and Ollama included, which never take the lock. It raises `GpuLeaseTimeout` after `ARGUS_GPU_LEASE_TIMEOUT_S` (300 s). The VRAM check is a snapshot, not a reservation, and a probe that cannot run lets the lease through. Containers each have their own `/tmp`, so set `ARGUS_GPU_LEASE_PATH` to a file on a volume every GPU service mounts, or the lock coordinates nothing.
|
|
15
16
|
- **`argus_cortex.store`** — the optional **stateful services layer** (the suite's persistent "memory"). External and opt-in — every store is configured by `CORTEX_*` env vars and no-ops when its URL is unset.
|
|
16
17
|
- **Phase 1 — Postgres lineage store.** Persists the `source_asset → caption → human_edit → dataset_membership → training_run` DAG for LoRA reproducibility, edit capture, and the feedback loop. Backed by a psycopg connection **pool** (concurrent `asyncio.to_thread` calls are safe; a dropped connection is replaced), with a `transaction()` boundary for atomic multi-step writes. `open_lineage_store()` returns a no-op store when `CORTEX_PG_URL` is unset; psycopg lives behind the `[postgres]` extra.
|
|
17
18
|
- **Phase 2 — Qdrant vector store.** Stores image + tag-set embeddings for retrieval-augmented few-shot and near-duplicate curation; put a `caption_id`/`asset_id` in a point's payload and a search hit joins straight back to the lineage. `open_vector_store()` no-ops when `CORTEX_QDRANT_URL` is unset; qdrant-client lives behind the `[qdrant]` extra. cortex stores vectors it's *handed* — computing embeddings is the caller's job.
|
|
@@ -73,6 +74,17 @@ policies.mode("adult-granular"), policies.mode("default") # "nsfw", "sfw"
|
|
|
73
74
|
|
|
74
75
|
Stages ask for a policy by name; `policies` holds each one's text, the rules that mean illegal content under it, and its SFW/NSFW mode (`default` and `sfw-*` are SFW).
|
|
75
76
|
|
|
77
|
+
```python
|
|
78
|
+
from argus_cortex.gpu import GpuLeaseTimeout, coordinator_from_env
|
|
79
|
+
|
|
80
|
+
gpu = coordinator_from_env() # ARGUS_GPU_COORDINATOR=lease, ARGUS_GPU_LEASE_PATH=/shared/argus-gpu.lock
|
|
81
|
+
try:
|
|
82
|
+
with gpu.lease(caller="faces", min_vram_mb=1500): # one argus job at a time, and 1.5 GB free
|
|
83
|
+
run_the_model()
|
|
84
|
+
except GpuLeaseTimeout:
|
|
85
|
+
... # no slot in time: back off, or answer 503
|
|
86
|
+
```
|
|
87
|
+
|
|
76
88
|
Each `check()` raises `CheckFailed` (`NudeNetCheckFailed` / `GuardCheckFailed`) for a network error or an unusable answer, so a caller checking many images can record "unchecked" against the one image — and should treat unchecked as *not* safe. Bad configuration raises `ModerationConfigError`, whose message never echoes the URL.
|
|
77
89
|
|
|
78
90
|
```python
|
|
@@ -35,6 +35,9 @@ server = ["starlette>=0.37"]
|
|
|
35
35
|
# Pillow, for SingGuardClient's downscaling of large images
|
|
36
36
|
# (argus_cortex.moderation.singguard), imported only when it downscales.
|
|
37
37
|
images = ["pillow>=10.0"]
|
|
38
|
+
# NVML (nvidia-ml-py, imported as pynvml), for argus_cortex.gpu's free-VRAM probe;
|
|
39
|
+
# imported only when probed. Without it the lease skips the VRAM wait.
|
|
40
|
+
gpu = ["nvidia-ml-py>=12.535"]
|
|
38
41
|
# The store drivers are in dev so each store's real client path (pool/enum/client
|
|
39
42
|
# construction, error wrapping) is exercised in CI, not just via fakes.
|
|
40
43
|
# starlette + httpx power the pure-ASGI guard tests (Starlette TestClient).
|
|
@@ -49,6 +52,7 @@ dev = [
|
|
|
49
52
|
"httpx>=0.27",
|
|
50
53
|
"httpx2>=2.7",
|
|
51
54
|
"pillow>=10.0",
|
|
55
|
+
"nvidia-ml-py>=12.535",
|
|
52
56
|
"certifi>=2026.10.19",
|
|
53
57
|
]
|
|
54
58
|
|
|
@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
|
|
|
18
18
|
commit_id: str | None
|
|
19
19
|
__commit_id__: str | None
|
|
20
20
|
|
|
21
|
-
__version__ = version = '0.4.
|
|
22
|
-
__version_tuple__ = version_tuple = (0, 4,
|
|
21
|
+
__version__ = version = '0.4.4'
|
|
22
|
+
__version_tuple__ = version_tuple = (0, 4, 4)
|
|
23
23
|
|
|
24
24
|
__commit_id__ = commit_id = None
|
|
@@ -0,0 +1,242 @@
|
|
|
1
|
+
"""GPU tenancy shared across the suite: one heavy argus job on a GPU at a time,
|
|
2
|
+
and only once the card has room for it.
|
|
3
|
+
|
|
4
|
+
Generalised from argus-lens's ``gpu`` coordinator so every stage that runs a
|
|
5
|
+
model on a GPU it shares (lens captioning, curator's faces, alongside ComfyUI
|
|
6
|
+
and an Ollama-backed SingGuard) takes the same lease.
|
|
7
|
+
|
|
8
|
+
* :class:`GpuCoordinator`: the protocol: ``lease(caller=, min_vram_mb=)`` is a
|
|
9
|
+
context manager that holds a GPU slot for its body.
|
|
10
|
+
* :class:`NullCoordinator`: no gating, the default.
|
|
11
|
+
* :class:`LeaseCoordinator`: an advisory file lock (POSIX ``flock``) taken by
|
|
12
|
+
every argus process on the host, then a wait until :func:`free_vram_mb` reports
|
|
13
|
+
``min_vram_mb`` free. The lock orders argus jobs among themselves; the VRAM wait
|
|
14
|
+
is what accounts for tenants that never take the lock (ComfyUI, Ollama), since
|
|
15
|
+
NVML reports free memory across every process on the card.
|
|
16
|
+
* :func:`coordinator_from_env`: ``ARGUS_GPU_COORDINATOR`` (``none``/``lease``),
|
|
17
|
+
``ARGUS_GPU_LEASE_PATH``, ``ARGUS_GPU_LEASE_TIMEOUT_S``, ``ARGUS_GPU_DEVICE``.
|
|
18
|
+
The names are host-wide, not per stage: every process must agree on them.
|
|
19
|
+
One lock serialises every argus GPU job on the host; ``ARGUS_GPU_DEVICE``
|
|
20
|
+
only selects which card NVML probes.
|
|
21
|
+
|
|
22
|
+
The VRAM check is a snapshot, not a reservation: a tenant outside the lock can
|
|
23
|
+
still allocate between the check and the job. Each stage owns its footprint
|
|
24
|
+
estimates (``min_vram_mb``); a probe that cannot run (no driver, no NVML, a CPU
|
|
25
|
+
host) lets the lease through rather than blocking forever.
|
|
26
|
+
|
|
27
|
+
Containers have their own ``/tmp``: the lock only coordinates processes that see
|
|
28
|
+
the same file, so point ``ARGUS_GPU_LEASE_PATH`` at a volume every GPU service
|
|
29
|
+
mounts. ``flock`` works across containers on one host kernel. NVML is behind the
|
|
30
|
+
``[gpu]`` extra (``nvidia-ml-py``) and imported only when probed.
|
|
31
|
+
"""
|
|
32
|
+
|
|
33
|
+
from __future__ import annotations
|
|
34
|
+
|
|
35
|
+
import contextlib
|
|
36
|
+
import logging
|
|
37
|
+
import math
|
|
38
|
+
import os
|
|
39
|
+
import tempfile
|
|
40
|
+
import threading
|
|
41
|
+
import time
|
|
42
|
+
from collections.abc import Callable, Iterator
|
|
43
|
+
from typing import Protocol, runtime_checkable
|
|
44
|
+
|
|
45
|
+
_log = logging.getLogger(__name__)
|
|
46
|
+
|
|
47
|
+
DEFAULT_TIMEOUT_S = 300.0
|
|
48
|
+
DEFAULT_POLL_S = 0.5
|
|
49
|
+
COORDINATORS = ("none", "lease")
|
|
50
|
+
|
|
51
|
+
_nvml_lock = threading.Lock()
|
|
52
|
+
_nvml_ready: bool | None = None # None: not tried yet; False: NVML cannot run here
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
class GpuLeaseTimeout(TimeoutError):
|
|
56
|
+
"""No GPU slot within the timeout: the lock stayed held, or the card never had room.
|
|
57
|
+
|
|
58
|
+
A ``TimeoutError``, so a caller can catch either, and tell "no GPU slot" from
|
|
59
|
+
a failed job and back off.
|
|
60
|
+
"""
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
@runtime_checkable
|
|
64
|
+
class GpuCoordinator(Protocol):
|
|
65
|
+
"""Gates a block of GPU work behind a lease."""
|
|
66
|
+
|
|
67
|
+
name: str
|
|
68
|
+
|
|
69
|
+
def lease(self, *, caller: str, min_vram_mb: int) -> contextlib.AbstractContextManager[None]:
|
|
70
|
+
"""A context manager that holds a GPU slot for its body."""
|
|
71
|
+
...
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
class NullCoordinator:
|
|
75
|
+
"""No gating: every lease is granted at once. The default."""
|
|
76
|
+
|
|
77
|
+
name = "none"
|
|
78
|
+
|
|
79
|
+
@contextlib.contextmanager
|
|
80
|
+
def lease(self, *, caller: str, min_vram_mb: int) -> Iterator[None]:
|
|
81
|
+
yield
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def free_vram_mb(device: int = 0) -> int | None:
|
|
85
|
+
"""Free VRAM on NVML device *device*, in MB, across every process on it.
|
|
86
|
+
|
|
87
|
+
``None`` when it cannot be read: no ``[gpu]`` extra, no driver, no such
|
|
88
|
+
device. NVML numbers devices as ``nvidia-smi`` does, which inside a container
|
|
89
|
+
means the GPUs the container was given.
|
|
90
|
+
"""
|
|
91
|
+
global _nvml_ready
|
|
92
|
+
try:
|
|
93
|
+
import pynvml # noqa: PLC0415 - optional, behind the [gpu] extra
|
|
94
|
+
except ImportError:
|
|
95
|
+
return None
|
|
96
|
+
with _nvml_lock:
|
|
97
|
+
if _nvml_ready is None:
|
|
98
|
+
try:
|
|
99
|
+
pynvml.nvmlInit()
|
|
100
|
+
_nvml_ready = True
|
|
101
|
+
except Exception: # noqa: BLE001 - any NVML failure means "cannot probe"
|
|
102
|
+
_nvml_ready = False
|
|
103
|
+
if not _nvml_ready:
|
|
104
|
+
return None
|
|
105
|
+
try:
|
|
106
|
+
info = pynvml.nvmlDeviceGetMemoryInfo(pynvml.nvmlDeviceGetHandleByIndex(device))
|
|
107
|
+
except Exception: # noqa: BLE001
|
|
108
|
+
return None
|
|
109
|
+
return max(0, int(info.free) // (1024 * 1024))
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
def default_lock_path() -> str:
|
|
113
|
+
"""``<tempdir>/argus-gpu.lock``: shared only by processes that share that directory."""
|
|
114
|
+
return os.path.join(tempfile.gettempdir(), "argus-gpu.lock")
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
class LeaseCoordinator:
|
|
118
|
+
"""One argus GPU job at a time on this host, started only when the card has room.
|
|
119
|
+
|
|
120
|
+
Takes an exclusive ``flock`` on *lock_path*, then (when *min_vram_mb* > 0)
|
|
121
|
+
waits until *probe* reports that much free, polling every *poll_s*. The lock
|
|
122
|
+
is held through the wait, so the next argus job cannot take the room first.
|
|
123
|
+
Both waits share one *timeout_s* deadline; past it :class:`GpuLeaseTimeout`
|
|
124
|
+
is raised and the lock released. A probe answering ``None`` grants the lease.
|
|
125
|
+
"""
|
|
126
|
+
|
|
127
|
+
name = "lease"
|
|
128
|
+
|
|
129
|
+
def __init__(
|
|
130
|
+
self,
|
|
131
|
+
lock_path: str | None = None,
|
|
132
|
+
timeout_s: float = DEFAULT_TIMEOUT_S,
|
|
133
|
+
poll_s: float = DEFAULT_POLL_S,
|
|
134
|
+
device: int = 0,
|
|
135
|
+
probe: Callable[[int], int | None] = free_vram_mb,
|
|
136
|
+
) -> None:
|
|
137
|
+
if not math.isfinite(timeout_s):
|
|
138
|
+
raise ValueError("timeout_s must be a finite number")
|
|
139
|
+
self.lock_path = lock_path or default_lock_path()
|
|
140
|
+
self.timeout_s = max(0.0, timeout_s)
|
|
141
|
+
self.poll_s = poll_s
|
|
142
|
+
self.device = device
|
|
143
|
+
self.probe = probe
|
|
144
|
+
|
|
145
|
+
@contextlib.contextmanager
|
|
146
|
+
def lease(self, *, caller: str, min_vram_mb: int) -> Iterator[None]:
|
|
147
|
+
try:
|
|
148
|
+
import fcntl # noqa: PLC0415 - POSIX only
|
|
149
|
+
except ImportError as exc: # pragma: no cover - non-POSIX host
|
|
150
|
+
raise RuntimeError("the 'lease' GPU coordinator needs POSIX fcntl; use 'none'") from exc
|
|
151
|
+
|
|
152
|
+
deadline = time.monotonic() + self.timeout_s
|
|
153
|
+
directory = os.path.dirname(self.lock_path)
|
|
154
|
+
if directory:
|
|
155
|
+
os.makedirs(directory, exist_ok=True)
|
|
156
|
+
# 0o666: services in different containers may run as different users.
|
|
157
|
+
# os.open's mode is masked by umask, so chmod the fd or the file is
|
|
158
|
+
# 0644/0600 and the next UID cannot open it.
|
|
159
|
+
fd = os.open(self.lock_path, os.O_CREAT | os.O_RDWR, 0o666)
|
|
160
|
+
try:
|
|
161
|
+
with contextlib.suppress(OSError):
|
|
162
|
+
os.fchmod(fd, 0o666)
|
|
163
|
+
while True:
|
|
164
|
+
try:
|
|
165
|
+
fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
166
|
+
break
|
|
167
|
+
except BlockingIOError:
|
|
168
|
+
if time.monotonic() >= deadline:
|
|
169
|
+
raise GpuLeaseTimeout(
|
|
170
|
+
f"GPU lease not acquired within {self.timeout_s:g}s ({self.lock_path})"
|
|
171
|
+
) from None
|
|
172
|
+
time.sleep(self.poll_s)
|
|
173
|
+
try:
|
|
174
|
+
self._wait_for_room(caller, min_vram_mb, deadline)
|
|
175
|
+
_log.debug("gpu lease acquired by %s (min %d MB, %s)", caller, min_vram_mb, self.lock_path)
|
|
176
|
+
yield
|
|
177
|
+
finally:
|
|
178
|
+
with contextlib.suppress(OSError):
|
|
179
|
+
fcntl.flock(fd, fcntl.LOCK_UN)
|
|
180
|
+
finally:
|
|
181
|
+
os.close(fd)
|
|
182
|
+
|
|
183
|
+
def _wait_for_room(self, caller: str, min_vram_mb: int, deadline: float) -> None:
|
|
184
|
+
if min_vram_mb <= 0:
|
|
185
|
+
return
|
|
186
|
+
while True:
|
|
187
|
+
free = self.probe(self.device)
|
|
188
|
+
if free is None or free >= min_vram_mb:
|
|
189
|
+
return
|
|
190
|
+
if time.monotonic() >= deadline:
|
|
191
|
+
raise GpuLeaseTimeout(
|
|
192
|
+
f"GPU {self.device} had {free} MB free, {caller} needs {min_vram_mb} MB, after {self.timeout_s:g}s"
|
|
193
|
+
)
|
|
194
|
+
time.sleep(self.poll_s)
|
|
195
|
+
|
|
196
|
+
|
|
197
|
+
def build_coordinator(
|
|
198
|
+
name: str,
|
|
199
|
+
*,
|
|
200
|
+
lock_path: str | None = None,
|
|
201
|
+
timeout_s: float | None = None,
|
|
202
|
+
device: int = 0,
|
|
203
|
+
) -> GpuCoordinator:
|
|
204
|
+
"""A coordinator by name, one of :data:`COORDINATORS` (empty is ``none``)."""
|
|
205
|
+
key = (name or "none").strip().lower()
|
|
206
|
+
if key == "none":
|
|
207
|
+
return NullCoordinator()
|
|
208
|
+
if key == "lease":
|
|
209
|
+
return LeaseCoordinator(
|
|
210
|
+
lock_path=lock_path,
|
|
211
|
+
timeout_s=DEFAULT_TIMEOUT_S if timeout_s is None else timeout_s,
|
|
212
|
+
device=device,
|
|
213
|
+
)
|
|
214
|
+
raise ValueError(f"unknown GPU coordinator {name!r}; choose from: {', '.join(COORDINATORS)}")
|
|
215
|
+
|
|
216
|
+
|
|
217
|
+
def _env_number(name: str, kind: type[float] | type[int]) -> float | int | None:
|
|
218
|
+
raw = os.environ.get(name, "").strip()
|
|
219
|
+
if not raw:
|
|
220
|
+
return None
|
|
221
|
+
try:
|
|
222
|
+
value = kind(raw)
|
|
223
|
+
except (ValueError, OverflowError):
|
|
224
|
+
value = -1
|
|
225
|
+
# inf/nan parse as float and never compare < 0, so the deadline never fires.
|
|
226
|
+
if isinstance(value, float) and not math.isfinite(value):
|
|
227
|
+
value = -1
|
|
228
|
+
if value < 0:
|
|
229
|
+
_log.warning("%s set to an unusable value %r; using the default", name, raw)
|
|
230
|
+
return None
|
|
231
|
+
return value
|
|
232
|
+
|
|
233
|
+
|
|
234
|
+
def coordinator_from_env() -> GpuCoordinator:
|
|
235
|
+
"""The coordinator the environment names; ``none`` when ``ARGUS_GPU_COORDINATOR`` is unset."""
|
|
236
|
+
device = _env_number("ARGUS_GPU_DEVICE", int)
|
|
237
|
+
return build_coordinator(
|
|
238
|
+
os.environ.get("ARGUS_GPU_COORDINATOR", "none"),
|
|
239
|
+
lock_path=os.environ.get("ARGUS_GPU_LEASE_PATH") or None,
|
|
240
|
+
timeout_s=_env_number("ARGUS_GPU_LEASE_TIMEOUT_S", float),
|
|
241
|
+
device=0 if device is None else int(device),
|
|
242
|
+
)
|
|
@@ -47,7 +47,9 @@ MAX_EDGE = 1536
|
|
|
47
47
|
MAX_BYTES = 4 * 1024 * 1024
|
|
48
48
|
_MIN_EDGE = 256 # below this, give up rather than send SingGuard a postage stamp
|
|
49
49
|
_NEEDS_PILLOW = "downscaling images needs: pip install argus-cortex[images] (or pass downscale=False)"
|
|
50
|
-
|
|
50
|
+
# Fits a 2048-token slow writeup plus a short cascade preface. 4000
|
|
51
|
+
# cut recorded reasoning mid-sentence on cascade scans.
|
|
52
|
+
_EXPLANATION_MAX = 16384
|
|
51
53
|
_FIELD_MAX = 200
|
|
52
54
|
|
|
53
55
|
|
|
@@ -70,7 +70,13 @@ class LineageStore(Protocol):
|
|
|
70
70
|
...
|
|
71
71
|
|
|
72
72
|
def record_asset(self, asset: SourceAsset) -> str | None:
|
|
73
|
-
"""Insert (or reuse, by ``sha256``) a source asset; return its id.
|
|
73
|
+
"""Insert (or reuse, by ``sha256``) a source asset; return its id.
|
|
74
|
+
|
|
75
|
+
On reuse, ``uri`` becomes this call's; ``metadata`` merges into the
|
|
76
|
+
stored bag (this call's keys win); a ``None`` ``immich_id`` keeps the
|
|
77
|
+
stored one. So stages recording the same bytes add to the row rather
|
|
78
|
+
than erase each other's fields.
|
|
79
|
+
"""
|
|
74
80
|
...
|
|
75
81
|
|
|
76
82
|
def record_caption(self, asset_id: str, caption: Caption) -> str | None:
|
|
@@ -302,8 +308,12 @@ class PostgresLineageStore:
|
|
|
302
308
|
|
|
303
309
|
def record_asset(self, asset: SourceAsset) -> str | None:
|
|
304
310
|
# De-dupe on content hash: re-ingesting identical bytes returns the
|
|
305
|
-
# existing row (DO UPDATE ensures RETURNING fires even on conflict)
|
|
306
|
-
#
|
|
311
|
+
# existing row (DO UPDATE ensures RETURNING fires even on conflict).
|
|
312
|
+
# Several stages record the same asset (quarry with provenance, lens
|
|
313
|
+
# with a caption), so a later call must not erase what an earlier one
|
|
314
|
+
# knew: uri is the latest location (last-write-wins), immich_id keeps
|
|
315
|
+
# the stored id when this call has none, and metadata merges key by
|
|
316
|
+
# key (jsonb ||: this call's keys win, the others are kept).
|
|
307
317
|
# NOTE: dedup only works when sha256 is set — NULLs are distinct under the
|
|
308
318
|
# UNIQUE index, so a hash-less asset inserts a fresh row every call.
|
|
309
319
|
return self._insert(
|
|
@@ -312,8 +322,8 @@ class PostgresLineageStore:
|
|
|
312
322
|
VALUES (%s, %s, %s, %s::jsonb)
|
|
313
323
|
ON CONFLICT (sha256) DO UPDATE
|
|
314
324
|
SET uri = EXCLUDED.uri,
|
|
315
|
-
immich_id = EXCLUDED.immich_id,
|
|
316
|
-
metadata = EXCLUDED.metadata
|
|
325
|
+
immich_id = COALESCE(EXCLUDED.immich_id, source_asset.immich_id),
|
|
326
|
+
metadata = source_asset.metadata || EXCLUDED.metadata
|
|
317
327
|
RETURNING id
|
|
318
328
|
""",
|
|
319
329
|
(asset.uri, asset.sha256, asset.immich_id, self._json(asset.metadata)),
|
|
@@ -0,0 +1,227 @@
|
|
|
1
|
+
"""GPU coordination: the file lease, the VRAM wait, and configuration from the environment."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import builtins
|
|
6
|
+
import os
|
|
7
|
+
import threading
|
|
8
|
+
from pathlib import Path
|
|
9
|
+
|
|
10
|
+
import pytest
|
|
11
|
+
|
|
12
|
+
from argus_cortex import gpu
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
@pytest.fixture(autouse=True)
|
|
16
|
+
def _reset_nvml_ready(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
17
|
+
monkeypatch.setattr(gpu, "_nvml_ready", None)
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def _lease(tmp_path: Path, probe=lambda device: None, timeout_s: float = 5.0) -> gpu.LeaseCoordinator: # noqa: ANN001
|
|
21
|
+
return gpu.LeaseCoordinator(
|
|
22
|
+
lock_path=str(tmp_path / "locks" / "gpu.lock"), timeout_s=timeout_s, poll_s=0.01, probe=probe
|
|
23
|
+
)
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def test_null_coordinator_grants_at_once() -> None:
|
|
27
|
+
coordinator = gpu.NullCoordinator()
|
|
28
|
+
assert isinstance(coordinator, gpu.GpuCoordinator)
|
|
29
|
+
with coordinator.lease(caller="t", min_vram_mb=10**9):
|
|
30
|
+
pass
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def test_lock_file_is_world_writable_despite_umask(tmp_path: Path) -> None:
|
|
34
|
+
"""os.open's mode is umask-masked; without fchmod the next container UID cannot open the lock."""
|
|
35
|
+
path = tmp_path / "locks" / "gpu.lock"
|
|
36
|
+
previous = os.umask(0o077)
|
|
37
|
+
try:
|
|
38
|
+
with gpu.LeaseCoordinator(lock_path=str(path), timeout_s=1, probe=lambda device: None).lease(
|
|
39
|
+
caller="a", min_vram_mb=0
|
|
40
|
+
):
|
|
41
|
+
mode = path.stat().st_mode & 0o777
|
|
42
|
+
finally:
|
|
43
|
+
os.umask(previous)
|
|
44
|
+
assert mode == 0o666
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
def test_lease_released_when_the_body_raises(tmp_path: Path) -> None:
|
|
48
|
+
with (
|
|
49
|
+
pytest.raises(RuntimeError, match="boom"),
|
|
50
|
+
_lease(tmp_path).lease(caller="a", min_vram_mb=0),
|
|
51
|
+
):
|
|
52
|
+
raise RuntimeError("boom")
|
|
53
|
+
with _lease(tmp_path, timeout_s=0.05).lease(caller="b", min_vram_mb=0):
|
|
54
|
+
pass
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def test_lease_is_exclusive_and_released(tmp_path: Path) -> None:
|
|
58
|
+
"""A second holder waits for the first; a holder that never lets go times the other out."""
|
|
59
|
+
holder, waiter = _lease(tmp_path), _lease(tmp_path, timeout_s=0.1)
|
|
60
|
+
assert isinstance(holder, gpu.GpuCoordinator)
|
|
61
|
+
with (
|
|
62
|
+
holder.lease(caller="a", min_vram_mb=0),
|
|
63
|
+
pytest.raises(gpu.GpuLeaseTimeout, match="not acquired"),
|
|
64
|
+
waiter.lease(caller="b", min_vram_mb=0),
|
|
65
|
+
):
|
|
66
|
+
pass
|
|
67
|
+
with waiter.lease(caller="b", min_vram_mb=0): # free again once the first exits
|
|
68
|
+
pass
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
def test_lease_waits_for_the_holder_to_finish(tmp_path: Path) -> None:
|
|
72
|
+
order: list[str] = []
|
|
73
|
+
entered, release = threading.Event(), threading.Event()
|
|
74
|
+
|
|
75
|
+
def first() -> None:
|
|
76
|
+
with _lease(tmp_path).lease(caller="a", min_vram_mb=0):
|
|
77
|
+
entered.set()
|
|
78
|
+
release.wait(5)
|
|
79
|
+
order.append("a done")
|
|
80
|
+
|
|
81
|
+
thread = threading.Thread(target=first)
|
|
82
|
+
thread.start()
|
|
83
|
+
assert entered.wait(5)
|
|
84
|
+
threading.Timer(0.05, release.set).start()
|
|
85
|
+
with _lease(tmp_path).lease(caller="b", min_vram_mb=0):
|
|
86
|
+
order.append("b in")
|
|
87
|
+
thread.join(5)
|
|
88
|
+
assert order == ["a done", "b in"]
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
def test_lease_waits_for_free_vram_while_holding_the_lock(tmp_path: Path) -> None:
|
|
92
|
+
readings = iter([1000, 2000, 9000])
|
|
93
|
+
seen: list[int] = []
|
|
94
|
+
|
|
95
|
+
def probe(device: int) -> int | None:
|
|
96
|
+
seen.append(device)
|
|
97
|
+
return next(readings)
|
|
98
|
+
|
|
99
|
+
coordinator = _lease(tmp_path, probe=probe)
|
|
100
|
+
# The lock is held during and after the wait: another argus job cannot take the room.
|
|
101
|
+
with (
|
|
102
|
+
coordinator.lease(caller="faces", min_vram_mb=4000),
|
|
103
|
+
pytest.raises(gpu.GpuLeaseTimeout),
|
|
104
|
+
_lease(tmp_path, timeout_s=0.05).lease(caller="lens", min_vram_mb=0),
|
|
105
|
+
):
|
|
106
|
+
pass
|
|
107
|
+
assert seen == [0, 0, 0]
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
def test_lease_times_out_when_the_card_never_has_room_and_lets_go(tmp_path: Path) -> None:
|
|
111
|
+
coordinator = _lease(tmp_path, probe=lambda device: 500, timeout_s=0.05)
|
|
112
|
+
with (
|
|
113
|
+
pytest.raises(gpu.GpuLeaseTimeout, match="500 MB free, faces needs 4000 MB"),
|
|
114
|
+
coordinator.lease(caller="faces", min_vram_mb=4000),
|
|
115
|
+
):
|
|
116
|
+
pass
|
|
117
|
+
with _lease(tmp_path, timeout_s=0.05).lease(caller="next", min_vram_mb=0): # the lock was released
|
|
118
|
+
pass
|
|
119
|
+
|
|
120
|
+
|
|
121
|
+
def test_lease_proceeds_when_vram_cannot_be_read_or_is_not_needed(tmp_path: Path) -> None:
|
|
122
|
+
calls: list[int] = []
|
|
123
|
+
with _lease(tmp_path, probe=lambda device: calls.append(device)).lease(caller="cpu host", min_vram_mb=4000):
|
|
124
|
+
pass
|
|
125
|
+
assert calls == [0] # asked once, got None, went ahead
|
|
126
|
+
with _lease(tmp_path, probe=lambda device: pytest.fail("probed")).lease(caller="tiny", min_vram_mb=0):
|
|
127
|
+
pass
|
|
128
|
+
|
|
129
|
+
|
|
130
|
+
def test_free_vram_reads_bytes_for_the_asked_device(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
131
|
+
import sys
|
|
132
|
+
import types
|
|
133
|
+
|
|
134
|
+
fake = types.ModuleType("pynvml")
|
|
135
|
+
seen: list[int] = []
|
|
136
|
+
|
|
137
|
+
class _Info:
|
|
138
|
+
free = 3 * 1024 * 1024 * 1024
|
|
139
|
+
|
|
140
|
+
def handle(index: int) -> str:
|
|
141
|
+
seen.append(index)
|
|
142
|
+
return f"dev{index}"
|
|
143
|
+
|
|
144
|
+
def memory_info(dev: str) -> _Info:
|
|
145
|
+
assert dev == "dev2"
|
|
146
|
+
return _Info()
|
|
147
|
+
|
|
148
|
+
fake.nvmlInit = lambda: None # type: ignore[attr-defined]
|
|
149
|
+
fake.nvmlDeviceGetHandleByIndex = handle # type: ignore[attr-defined]
|
|
150
|
+
fake.nvmlDeviceGetMemoryInfo = memory_info # type: ignore[attr-defined]
|
|
151
|
+
monkeypatch.setitem(sys.modules, "pynvml", fake)
|
|
152
|
+
assert gpu.free_vram_mb(2) == 3072
|
|
153
|
+
assert seen == [2]
|
|
154
|
+
|
|
155
|
+
|
|
156
|
+
def test_free_vram_is_none_without_nvml(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
157
|
+
real_import = builtins.__import__
|
|
158
|
+
|
|
159
|
+
def no_pynvml(name: str, *args, **kwargs): # noqa: ANN002, ANN003, ANN202
|
|
160
|
+
if name == "pynvml":
|
|
161
|
+
raise ImportError(name)
|
|
162
|
+
return real_import(name, *args, **kwargs)
|
|
163
|
+
|
|
164
|
+
monkeypatch.setattr(builtins, "__import__", no_pynvml)
|
|
165
|
+
assert gpu.free_vram_mb() is None
|
|
166
|
+
|
|
167
|
+
|
|
168
|
+
def test_free_vram_is_none_when_nvml_cannot_start(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
169
|
+
"""No driver: NVML init fails once and is not retried on every probe."""
|
|
170
|
+
import sys
|
|
171
|
+
import types
|
|
172
|
+
|
|
173
|
+
attempts: list[int] = []
|
|
174
|
+
fake = types.ModuleType("pynvml")
|
|
175
|
+
|
|
176
|
+
def init() -> None:
|
|
177
|
+
attempts.append(1)
|
|
178
|
+
raise RuntimeError("NVML Shared Library Not Found")
|
|
179
|
+
|
|
180
|
+
fake.nvmlInit = init # type: ignore[attr-defined]
|
|
181
|
+
monkeypatch.setitem(sys.modules, "pynvml", fake)
|
|
182
|
+
monkeypatch.setattr(gpu, "_nvml_ready", None)
|
|
183
|
+
assert gpu.free_vram_mb() is None and gpu.free_vram_mb() is None
|
|
184
|
+
assert attempts == [1]
|
|
185
|
+
|
|
186
|
+
|
|
187
|
+
def test_build_coordinator() -> None:
|
|
188
|
+
assert isinstance(gpu.build_coordinator(""), gpu.NullCoordinator)
|
|
189
|
+
assert isinstance(gpu.build_coordinator(" NONE "), gpu.NullCoordinator)
|
|
190
|
+
lease = gpu.build_coordinator("lease", lock_path="/x/gpu.lock", timeout_s=7, device=1)
|
|
191
|
+
assert isinstance(lease, gpu.LeaseCoordinator)
|
|
192
|
+
assert (lease.lock_path, lease.timeout_s, lease.device) == ("/x/gpu.lock", 7, 1)
|
|
193
|
+
assert gpu.build_coordinator("lease").lock_path == gpu.default_lock_path()
|
|
194
|
+
with pytest.raises(ValueError, match="none, lease"):
|
|
195
|
+
gpu.build_coordinator("gothmog")
|
|
196
|
+
with pytest.raises(ValueError, match="finite"):
|
|
197
|
+
gpu.LeaseCoordinator(timeout_s=float("inf"))
|
|
198
|
+
with pytest.raises(ValueError, match="finite"):
|
|
199
|
+
gpu.LeaseCoordinator(timeout_s=float("nan"))
|
|
200
|
+
|
|
201
|
+
|
|
202
|
+
def test_coordinator_from_env(monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture) -> None:
|
|
203
|
+
for name in ("ARGUS_GPU_COORDINATOR", "ARGUS_GPU_LEASE_PATH", "ARGUS_GPU_LEASE_TIMEOUT_S", "ARGUS_GPU_DEVICE"):
|
|
204
|
+
monkeypatch.delenv(name, raising=False)
|
|
205
|
+
assert isinstance(gpu.coordinator_from_env(), gpu.NullCoordinator)
|
|
206
|
+
monkeypatch.setenv("ARGUS_GPU_COORDINATOR", "lease")
|
|
207
|
+
monkeypatch.setenv("ARGUS_GPU_LEASE_PATH", "/shared/argus-gpu.lock")
|
|
208
|
+
monkeypatch.setenv("ARGUS_GPU_LEASE_TIMEOUT_S", "30")
|
|
209
|
+
monkeypatch.setenv("ARGUS_GPU_DEVICE", "1")
|
|
210
|
+
lease = gpu.coordinator_from_env()
|
|
211
|
+
assert (lease.lock_path, lease.timeout_s, lease.device) == ("/shared/argus-gpu.lock", 30.0, 1)
|
|
212
|
+
monkeypatch.setenv("ARGUS_GPU_LEASE_TIMEOUT_S", "soon")
|
|
213
|
+
monkeypatch.setenv("ARGUS_GPU_DEVICE", "-1")
|
|
214
|
+
with caplog.at_level("WARNING"):
|
|
215
|
+
lease = gpu.coordinator_from_env()
|
|
216
|
+
assert (lease.timeout_s, lease.device) == (gpu.DEFAULT_TIMEOUT_S, 0)
|
|
217
|
+
assert "ARGUS_GPU_LEASE_TIMEOUT_S" in caplog.text and "ARGUS_GPU_DEVICE" in caplog.text
|
|
218
|
+
caplog.clear()
|
|
219
|
+
monkeypatch.setenv("ARGUS_GPU_LEASE_TIMEOUT_S", "inf")
|
|
220
|
+
with caplog.at_level("WARNING"):
|
|
221
|
+
assert gpu.coordinator_from_env().timeout_s == gpu.DEFAULT_TIMEOUT_S
|
|
222
|
+
assert "ARGUS_GPU_LEASE_TIMEOUT_S" in caplog.text
|
|
223
|
+
caplog.clear()
|
|
224
|
+
monkeypatch.setenv("ARGUS_GPU_LEASE_TIMEOUT_S", "nan")
|
|
225
|
+
with caplog.at_level("WARNING"):
|
|
226
|
+
assert gpu.coordinator_from_env().timeout_s == gpu.DEFAULT_TIMEOUT_S
|
|
227
|
+
assert "ARGUS_GPU_LEASE_TIMEOUT_S" in caplog.text
|
|
@@ -253,6 +253,24 @@ def test_singguard_config() -> None:
|
|
|
253
253
|
SingGuardClient("gopher://g")
|
|
254
254
|
|
|
255
255
|
|
|
256
|
+
def test_to_result_keeps_a_long_explanation() -> None:
|
|
257
|
+
writeup = "Cascade: fast flagged this; the slow pass is the verdict.\n\n" + ("A posed adult. " * 400)
|
|
258
|
+
writeup = writeup.strip()
|
|
259
|
+
assert 4000 < len(writeup) < singguard._EXPLANATION_MAX
|
|
260
|
+
result = singguard.to_result(
|
|
261
|
+
{"safe": True, "verdict": "safe", "rule": "Safe", "explanation": writeup, "categories": []}
|
|
262
|
+
)
|
|
263
|
+
assert result.explanation == writeup
|
|
264
|
+
|
|
265
|
+
|
|
266
|
+
def test_to_result_clips_explanation_at_max() -> None:
|
|
267
|
+
exact = "y" * singguard._EXPLANATION_MAX
|
|
268
|
+
over = exact + "z"
|
|
269
|
+
payload = {"safe": True, "verdict": "safe", "rule": "Safe", "categories": []}
|
|
270
|
+
assert singguard.to_result({**payload, "explanation": exact}).explanation == exact
|
|
271
|
+
assert singguard.to_result({**payload, "explanation": over}).explanation == exact
|
|
272
|
+
|
|
273
|
+
|
|
256
274
|
def test_singguard_check(singguard_server) -> None:
|
|
257
275
|
fake, url = singguard_server
|
|
258
276
|
client = SingGuardClient(url, timeout=5, thinking="fast-slow", policy="### A. Minors")
|
|
@@ -206,9 +206,11 @@ def test_record_asset_returns_id_and_upserts_on_sha() -> None:
|
|
|
206
206
|
sql, params = conn.executed[-1]
|
|
207
207
|
assert "INSERT INTO source_asset" in sql
|
|
208
208
|
assert "ON CONFLICT (sha256)" in sql
|
|
209
|
-
# on a dedup hit
|
|
210
|
-
|
|
211
|
-
assert "
|
|
209
|
+
# on a dedup hit uri is refreshed, but a call without an immich_id or a
|
|
210
|
+
# metadata key keeps the stored one (another stage may have recorded it)
|
|
211
|
+
assert "uri = EXCLUDED.uri" in sql
|
|
212
|
+
assert "immich_id = COALESCE(EXCLUDED.immich_id, source_asset.immich_id)" in sql
|
|
213
|
+
assert "metadata = source_asset.metadata || EXCLUDED.metadata" in sql
|
|
212
214
|
# metadata is serialised to a JSON string for the ::jsonb cast
|
|
213
215
|
assert params == ("immich://1", "abc", None, '{"k": "v"}')
|
|
214
216
|
|
|
@@ -81,6 +81,20 @@ def test_transaction_rolls_back_on_error(store: PostgresLineageStore) -> None:
|
|
|
81
81
|
assert reused is not None
|
|
82
82
|
|
|
83
83
|
|
|
84
|
+
def test_reupsert_merges_metadata_and_keeps_immich_id(store: PostgresLineageStore) -> None:
|
|
85
|
+
# quarry records provenance first; lens later records the same bytes with
|
|
86
|
+
# none: the provenance and immich id survive, the uri moves, shared keys update.
|
|
87
|
+
sha = uuid.uuid4().hex
|
|
88
|
+
first = store.record_asset(
|
|
89
|
+
SourceAsset(uri="quarry://a", sha256=sha, immich_id="im-1", metadata={"license": "cc0", "seen": 1})
|
|
90
|
+
)
|
|
91
|
+
again = store.record_asset(SourceAsset(uri="/data/images/a.png", sha256=sha, metadata={"seen": 2}))
|
|
92
|
+
assert again == first
|
|
93
|
+
with store._acquire() as conn:
|
|
94
|
+
row = conn.execute("SELECT uri, immich_id, metadata FROM source_asset WHERE id = %s", (first,)).fetchone()
|
|
95
|
+
assert row == ("/data/images/a.png", "im-1", {"license": "cc0", "seen": 2})
|
|
96
|
+
|
|
97
|
+
|
|
84
98
|
def test_concurrent_writes_are_isolated(store: PostgresLineageStore) -> None:
|
|
85
99
|
# Many concurrent record_asset calls each borrow their own pooled connection;
|
|
86
100
|
# all succeed with distinct ids (no "another operation in progress").
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|