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.
Files changed (47) hide show
  1. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/CHANGELOG.md +43 -1
  2. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/CLAUDE.md +1 -0
  3. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/PKG-INFO +16 -1
  4. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/README.md +12 -0
  5. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/pyproject.toml +4 -0
  6. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/_version.py +2 -2
  7. argus_cortex-0.4.4/src/argus_cortex/gpu.py +242 -0
  8. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/moderation/singguard.py +3 -1
  9. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/store/lineage.py +15 -5
  10. argus_cortex-0.4.4/tests/test_gpu.py +227 -0
  11. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_moderation.py +18 -0
  12. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_store.py +5 -3
  13. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_store_pg_integration.py +14 -0
  14. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/.copier-answers.yml +0 -0
  15. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/.env.example +0 -0
  16. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/.github/workflows/ci.yml +0 -0
  17. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/.github/workflows/release.yml +0 -0
  18. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/.gitignore +0 -0
  19. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/LICENSE +0 -0
  20. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/Makefile +0 -0
  21. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/__init__.py +0 -0
  22. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/backends.py +0 -0
  23. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/moderation/__init__.py +0 -0
  24. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/moderation/_http.py +0 -0
  25. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/moderation/nudenet.py +0 -0
  26. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/moderation/policies.py +0 -0
  27. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/py.typed +0 -0
  28. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/server.py +0 -0
  29. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/store/__init__.py +0 -0
  30. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/store/blob.py +0 -0
  31. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/store/config.py +0 -0
  32. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/store/errors.py +0 -0
  33. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/store/models.py +0 -0
  34. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/store/schema.py +0 -0
  35. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/store/vector.py +0 -0
  36. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/taxonomy.py +0 -0
  37. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/wire.py +0 -0
  38. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/src/argus_cortex/xmp.py +0 -0
  39. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_backends.py +0 -0
  40. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_blob.py +0 -0
  41. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_policies.py +0 -0
  42. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_server.py +0 -0
  43. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_smoke.py +0 -0
  44. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_taxonomy.py +0 -0
  45. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_vector.py +0 -0
  46. {argus_cortex-0.4.2 → argus_cortex-0.4.4}/tests/test_wire.py +0 -0
  47. {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.2...HEAD
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.2
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.2'
22
- __version_tuple__ = version_tuple = (0, 4, 2)
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
- _EXPLANATION_MAX = 4000
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) and
306
- # refreshes uri/immich_id/metadata to the latest call (last-write-wins).
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, uri/immich_id/metadata are all refreshed (not just uri)
210
- assert "immich_id = EXCLUDED.immich_id" in sql
211
- assert "metadata = EXCLUDED.metadata" in sql
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