omnirun 0.2.3__tar.gz → 0.3.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (125) hide show
  1. {omnirun-0.2.3 → omnirun-0.3.0}/.envrc +1 -0
  2. {omnirun-0.2.3 → omnirun-0.3.0}/DESIGN.md +348 -58
  3. omnirun-0.2.3/README.md → omnirun-0.3.0/PKG-INFO +91 -30
  4. omnirun-0.2.3/PKG-INFO → omnirun-0.3.0/README.md +67 -50
  5. {omnirun-0.2.3 → omnirun-0.3.0}/TESTING.md +84 -8
  6. omnirun-0.3.0/docs/superpowers/plans/2026-07-11-phase1-discovery.md +1467 -0
  7. omnirun-0.3.0/docs/superpowers/plans/2026-07-11-phase2-state-layer.md +218 -0
  8. omnirun-0.3.0/docs/superpowers/plans/2026-07-11-phase3-scheduler.md +233 -0
  9. omnirun-0.3.0/docs/superpowers/plans/2026-07-12-phase4-lifecycle.md +1652 -0
  10. omnirun-0.3.0/docs/superpowers/specs/2026-07-11-omnirun-scheduler-redesign-design.md +317 -0
  11. omnirun-0.3.0/docs/superpowers/specs/2026-07-13-ssh-everywhere-notebook-transport-design.md +108 -0
  12. omnirun-0.3.0/docs/superpowers/specs/diagrams/components.svg +1 -0
  13. omnirun-0.3.0/docs/superpowers/specs/diagrams/deployment.svg +1 -0
  14. omnirun-0.3.0/docs/superpowers/specs/diagrams/domain.svg +1 -0
  15. omnirun-0.3.0/docs/superpowers/specs/diagrams/redesign-components.puml +46 -0
  16. omnirun-0.3.0/docs/superpowers/specs/diagrams/redesign-deployment.puml +44 -0
  17. omnirun-0.3.0/docs/superpowers/specs/diagrams/redesign-domain.puml +88 -0
  18. omnirun-0.3.0/docs/superpowers/specs/diagrams/redesign-seams.puml +47 -0
  19. omnirun-0.3.0/docs/superpowers/specs/diagrams/redesign-seq-happy.puml +39 -0
  20. omnirun-0.3.0/docs/superpowers/specs/diagrams/redesign-seq-recovery.puml +35 -0
  21. omnirun-0.3.0/docs/superpowers/specs/diagrams/redesign-states.puml +31 -0
  22. omnirun-0.3.0/docs/superpowers/specs/diagrams/seams.svg +1 -0
  23. omnirun-0.3.0/docs/superpowers/specs/diagrams/seq_happy.svg +1 -0
  24. omnirun-0.3.0/docs/superpowers/specs/diagrams/seq_recovery.svg +1 -0
  25. omnirun-0.3.0/docs/superpowers/specs/diagrams/states.svg +1 -0
  26. {omnirun-0.2.3 → omnirun-0.3.0}/pyproject.toml +14 -3
  27. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/__init__.py +1 -1
  28. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/backends/base.py +73 -2
  29. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/backends/colab.py +212 -41
  30. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/backends/jobdir.py +39 -5
  31. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/backends/kaggle.py +320 -63
  32. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/backends/local.py +23 -13
  33. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/backends/marketplace.py +17 -11
  34. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/backends/slurm.py +259 -7
  35. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/backends/ssh.py +45 -11
  36. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/backends/vast.py +15 -2
  37. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/bootstrap.py +189 -35
  38. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/chooser.py +0 -9
  39. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/cli.py +329 -102
  40. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/config.py +102 -7
  41. omnirun-0.3.0/src/omnirun/control.py +447 -0
  42. omnirun-0.3.0/src/omnirun/daemon.py +557 -0
  43. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/execlayer/ssh.py +31 -10
  44. omnirun-0.3.0/src/omnirun/models.py +421 -0
  45. omnirun-0.3.0/src/omnirun/providers/__init__.py +13 -0
  46. omnirun-0.3.0/src/omnirun/providers/adapter.py +250 -0
  47. omnirun-0.3.0/src/omnirun/providers/base.py +71 -0
  48. omnirun-0.3.0/src/omnirun/queue.py +66 -0
  49. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/repo.py +14 -80
  50. omnirun-0.3.0/src/omnirun/scheduler.py +198 -0
  51. omnirun-0.3.0/src/omnirun/sshconn.py +115 -0
  52. omnirun-0.3.0/src/omnirun/state/__init__.py +26 -0
  53. omnirun-0.3.0/src/omnirun/state/schema.py +79 -0
  54. omnirun-0.3.0/src/omnirun/state/store.py +612 -0
  55. omnirun-0.3.0/src/omnirun/transport.py +280 -0
  56. omnirun-0.3.0/tests/fakes.py +204 -0
  57. omnirun-0.3.0/tests/live/__init__.py +0 -0
  58. omnirun-0.3.0/tests/live/conftest.py +203 -0
  59. omnirun-0.3.0/tests/live/test_live_colab.py +107 -0
  60. omnirun-0.3.0/tests/live/test_live_kaggle.py +64 -0
  61. omnirun-0.3.0/tests/live/test_live_slurm.py +111 -0
  62. omnirun-0.3.0/tests/test_backend_discover.py +51 -0
  63. omnirun-0.3.0/tests/test_bootstrap.py +407 -0
  64. {omnirun-0.2.3 → omnirun-0.3.0}/tests/test_chooser.py +0 -24
  65. {omnirun-0.2.3 → omnirun-0.3.0}/tests/test_cli.py +318 -93
  66. {omnirun-0.2.3 → omnirun-0.3.0}/tests/test_colab.py +144 -10
  67. omnirun-0.3.0/tests/test_colab_status_retry.py +49 -0
  68. omnirun-0.3.0/tests/test_config.py +134 -0
  69. omnirun-0.3.0/tests/test_control_e2e.py +409 -0
  70. omnirun-0.3.0/tests/test_execlayer_ssh.py +81 -0
  71. {omnirun-0.2.3 → omnirun-0.3.0}/tests/test_kaggle.py +357 -33
  72. omnirun-0.3.0/tests/test_kaggle_discover.py +45 -0
  73. {omnirun-0.2.3 → omnirun-0.3.0}/tests/test_local_backend.py +34 -1
  74. {omnirun-0.2.3 → omnirun-0.3.0}/tests/test_marketplaces.py +19 -1
  75. omnirun-0.3.0/tests/test_models_facts.py +51 -0
  76. omnirun-0.3.0/tests/test_models_scheduler.py +328 -0
  77. omnirun-0.3.0/tests/test_provider_adapter.py +564 -0
  78. {omnirun-0.2.3 → omnirun-0.3.0}/tests/test_queue.py +65 -99
  79. {omnirun-0.2.3 → omnirun-0.3.0}/tests/test_repo.py +1 -68
  80. omnirun-0.3.0/tests/test_scheduler.py +481 -0
  81. omnirun-0.3.0/tests/test_scheduler_invariants.py +459 -0
  82. {omnirun-0.2.3 → omnirun-0.3.0}/tests/test_slurm.py +38 -7
  83. omnirun-0.3.0/tests/test_slurm_discover.py +393 -0
  84. {omnirun-0.2.3 → omnirun-0.3.0}/tests/test_ssh_backend.py +49 -6
  85. omnirun-0.3.0/tests/test_ssh_everywhere_t3.py +660 -0
  86. {omnirun-0.2.3 → omnirun-0.3.0}/tests/test_ssh_exec.py +6 -2
  87. omnirun-0.3.0/tests/test_sshconn.py +115 -0
  88. omnirun-0.3.0/tests/test_state_store.py +598 -0
  89. omnirun-0.3.0/tests/test_transport.py +84 -0
  90. omnirun-0.3.0/tests/test_vast_cuda_filter.py +37 -0
  91. {omnirun-0.2.3 → omnirun-0.3.0}/uv.lock +233 -3
  92. omnirun-0.2.3/src/omnirun/daemon.py +0 -493
  93. omnirun-0.2.3/src/omnirun/models.py +0 -196
  94. omnirun-0.2.3/src/omnirun/queue.py +0 -106
  95. omnirun-0.2.3/src/omnirun/store.py +0 -107
  96. omnirun-0.2.3/tests/test_bootstrap.py +0 -189
  97. omnirun-0.2.3/tests/test_config.py +0 -37
  98. omnirun-0.2.3/tests/test_store.py +0 -85
  99. {omnirun-0.2.3 → omnirun-0.3.0}/.github/workflows/checks.yml +0 -0
  100. {omnirun-0.2.3 → omnirun-0.3.0}/.github/workflows/ci.yml +0 -0
  101. {omnirun-0.2.3 → omnirun-0.3.0}/.github/workflows/publish.yml +0 -0
  102. {omnirun-0.2.3 → omnirun-0.3.0}/.gitignore +0 -0
  103. {omnirun-0.2.3 → omnirun-0.3.0}/.python-version +0 -0
  104. {omnirun-0.2.3 → omnirun-0.3.0}/AGENTS.md +0 -0
  105. {omnirun-0.2.3 → omnirun-0.3.0}/CLAUDE.md +0 -0
  106. {omnirun-0.2.3 → omnirun-0.3.0}/flake.lock +0 -0
  107. {omnirun-0.2.3 → omnirun-0.3.0}/flake.nix +0 -0
  108. {omnirun-0.2.3 → omnirun-0.3.0}/main.py +0 -0
  109. {omnirun-0.2.3 → omnirun-0.3.0}/research/colab-kaggle.md +0 -0
  110. {omnirun-0.2.3 → omnirun-0.3.0}/research/gpu-marketplaces.md +0 -0
  111. {omnirun-0.2.3 → omnirun-0.3.0}/research/launcher-landscape.md +0 -0
  112. {omnirun-0.2.3 → omnirun-0.3.0}/research/ml-job-queues.md +0 -0
  113. {omnirun-0.2.3 → omnirun-0.3.0}/research/slurm-ssh.md +0 -0
  114. {omnirun-0.2.3 → omnirun-0.3.0}/scripts/check-version.sh +0 -0
  115. {omnirun-0.2.3 → omnirun-0.3.0}/scripts/pre-push +0 -0
  116. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/backends/__init__.py +0 -0
  117. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/backends/runpod.py +0 -0
  118. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/backends/tarsafe.py +0 -0
  119. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/backends/thunder.py +0 -0
  120. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/execlayer/__init__.py +0 -0
  121. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/execlayer/base.py +0 -0
  122. {omnirun-0.2.3 → omnirun-0.3.0}/src/omnirun/execlayer/local.py +0 -0
  123. {omnirun-0.2.3 → omnirun-0.3.0}/tests/__init__.py +0 -0
  124. {omnirun-0.2.3 → omnirun-0.3.0}/tests/conftest.py +0 -0
  125. {omnirun-0.2.3 → omnirun-0.3.0}/tests/test_tarsafe.py +0 -0
@@ -1,3 +1,4 @@
1
1
  if has nix; then
2
2
  use flake .
3
3
  fi
4
+ dotenv_if_exists
@@ -164,15 +164,17 @@ break sibling jobs at the same sha).
164
164
 
165
165
  ## 4. Chooser
166
166
 
167
+ `omnirun offers` probes and ranks — display only, no submission. Actual
168
+ placement is done by the scheduler tick inside `Control` (§10).
169
+
167
170
  1. Probe all enabled backends in parallel (timeout 10s each).
168
171
  2. Partition offers: fit / unfit.
169
172
  3. Score fitting offers. Default policy (configurable weights):
170
173
  - `total_cost = hourly × est_time` (0 for free backends)
171
174
  - `time_to_result = wait_estimate + est_time`
172
- - **Auto-pick** iff a free offer has `wait < auto_wait_threshold` (default 15m),
173
- or exactly one offer fits, or `--yes` with a `--max-cost`.
174
- - Otherwise render the table (rich): backend, GPU, $/hr, est. total $, est. wait,
175
- time-to-result, notes — user picks by number. `--yes` picks the top-ranked.
175
+ - The old interactive offer-pick UI (`--yes` required) is gone; `submit` now
176
+ always routes through the scheduler and auto-places. `omnirun offers` still
177
+ prints the table so you can inspect options before committing.
176
178
 
177
179
  Wait estimation is honest about uncertainty. Slurm: `sinfo` idle-node check
178
180
  ("likely immediate") → `sbatch --test-only` pre-submit estimate + script validation
@@ -189,10 +191,11 @@ explicitly requested premium tier still surfaces, marked non-free.
189
191
 
190
192
  ## 5. State & config
191
193
 
192
- - Client state: `$OMNIRUN_STATE_DIR/jobs/<job_id>/meta.json` (default
193
- `~/.local/share/omnirun/`; spec, handle, last-known status, offer chosen). Plain JSON,
194
- greppable, no daemon. The optional scheduler (§6) persists its queue alongside it under
195
- `$OMNIRUN_STATE_DIR/queue/` (one atomic JSON file per entry) plus `daemon.json`.
194
+ - Client state: `$OMNIRUN_STATE_DIR/omnirun.db` (default `~/.local/share/omnirun/`)
195
+ a SQLite database managed by the `Store` repository (SQLAlchemy Core 2.0). See §9
196
+ for the full SQL state layer description, including the Postgres option and the atomic
197
+ `reserve_entry` concurrency guard. The optional `[state]` config section selects the
198
+ backend/path/url; `omnirun state migrate` imports from legacy JSON files.
196
199
  - Config: `~/.config/omnirun/config.toml` (override `$OMNIRUN_CONFIG`; backends,
197
200
  credentials refs, policy, daemon) + optional per-repo `omnirun.toml` (resource defaults,
198
201
  outputs, env.setup overrides). Backend sections are permissive: common fields are typed,
@@ -229,7 +232,7 @@ gpus = [{ type = "4090", count = 1 }] # static capability declaration
229
232
 
230
233
  [backends.kaggle]
231
234
  type = "kaggle" # creds from ~/.config/kaggle/kaggle.json
232
- weekly_gpu_hours = 30 # local budget (no quota API)
235
+ # weekly GPU quota is read live from the quota API
233
236
 
234
237
  [backends.colab]
235
238
  type = "colab"
@@ -248,8 +251,9 @@ type = "thunder" # TNR_API_TOKEN env
248
251
 
249
252
  ## 6. Repos & credentials
250
253
 
251
- Submit-time invariant: working tree clean (or `--dirty` to auto-stash-commit onto a
252
- `omnirun/<job_id>` ref v1), HEAD pushed to remote (offer to push).
254
+ Submit-time invariant: working tree clean (always enforced a dirty tree is
255
+ refused, with no escape hatch, so a job only ever runs a real, reproducible
256
+ revision), HEAD pushed to remote (offer to push).
253
257
 
254
258
  Worker access to private repos — **no git credentials ever leave the laptop**:
255
259
  - **ssh/slurm/marketplace**: at submit time the client `git push`es the exact sha
@@ -266,7 +270,7 @@ Worker access to private repos — **no git credentials ever leave the laptop**:
266
270
  itself over its own internet connection from an anonymous `https://` URL — no bundle
267
271
  is shipped, no credentials are needed (`CodeSource(kind="remote")`, §3). A plan URL is
268
272
  returned only when *all* hold, else `None` (→ bundle): (a) there is a real origin
269
- remote and the sha is a normal pushed **branch** commit (not `--dirty`, not detached
273
+ remote and the sha is a normal pushed **branch** commit (not a detached
270
274
  HEAD); (b) the origin is anonymously **public** — `remote_is_public()` checks via
271
275
  `gh repo view --json visibility` for GitHub when `gh` is present, else an
272
276
  unauthenticated `curl` smart-http probe of `<url>/info/refs?service=git-upload-pack`
@@ -322,8 +326,9 @@ is already in the revision and is left alone. (`repo.env_file` / `jobdir.stage_e
322
326
  `logs/ outputs/ result.json phase` into `/kaggle/working/omnirun-job.tar.gz` so results
323
327
  persist with the kernel version. Env handling is forced to `system` (§ below) to keep
324
328
  Kaggle's preinstalled CUDA-matched torch. Probe constraints: ~12h session cap → unfit if
325
- `resources.time` exceeds; ~30 GPU-h/week quota is not queryable tracked locally in
326
- state as a budget (`weekly_gpu_hours`). Poll ≥30s. `gc()` is a no-op (nothing
329
+ `resources.time` exceeds; the weekly GPU quota is read live from `KaggleApi.quota_view()`
330
+ (same source as `discover()`) GPU offers are unfit only when the real remaining allowance
331
+ is exhausted (0h). Poll ≥30s. `gc()` is a no-op (nothing
327
332
  worker-side: the bundle rode inside the kernel).
328
333
  - **Colab**: fully automated via the official `google-colab-cli` (v0.6+, June 2026;
329
334
  one-time OAuth). Submit = `colab new --gpu <T4|L4|G4|A100|H100>` → `colab upload`
@@ -353,50 +358,327 @@ ssh → run bootstrap detached → poll. Auto-terminate on completion (configura
353
358
  `keep_alive`), plus `omnirun gc` to reap leaked instances (safety net against
354
359
  billing surprises). Idle-timeout watchdog baked into the payload wrapper.
355
360
 
356
- ## 9. Queue & scheduler daemon (optional)
357
-
358
- Direct `submit` stays daemonless. For fan-out — many jobs, or spreading a batch across
359
- backends with per-backend concurrency caps an *optional* long-lived scheduler is
360
- available (`daemon.py`, `queue.py`). It changes nothing about how a job runs; it just
361
- decides *when* and *where*, then calls the same `Backend.submit`.
362
-
363
- - **`omnirun serve`** runs the daemon in the foreground: a localhost TCP socket (default
364
- `127.0.0.1:8787`), newline-delimited JSON request/response, plus a scheduler thread
365
- ticking every `poll_interval_s` (default 10s). It owns a durable queue persisted at
366
- `$OMNIRUN_STATE_DIR/queue/` (one atomic JSON file per entry; a restart re-reads and
367
- resumes) and records a `daemon.json` (host/port/pid) clients use to find it.
368
- - **Scheduler tick** (under a lock): refresh RUNNING entries against their backend, then
369
- place PENDING ones. Placement reserves a backend slot *synchronously*
370
- (state PLACING, so concurrent placements can't double-book) before dispatching the
371
- blocking `submit` on a thread pool; it backfills as jobs complete, and retries a failed
372
- submit up to 3 attempts before marking it FAILED. Each backend's `max_parallel` caps its
373
- concurrent non-terminal jobs per-partition Slurm limits = one backend section per
374
- partition, each capped.
375
- - **Entry lifecycle**: PENDING → PLACING → RUNNING → SUCCEEDED / FAILED / CANCELLED.
376
- - **Commands**: `enqueue [--count N] [--backend NAME] -- CMD...` (same resource flags as
377
- `submit`), `queue` (show the table), `queue --wait` (poll until all terminal),
378
- `queue --cancel <qid|all>`; the socket also speaks `ping` / `list` / `shutdown`.
379
- - **Placement is greedy** it favors the fastest-freeing backend via the same
380
- `chooser.rank`, with a short offer cache so a batch of identical jobs doesn't re-probe
381
- every tick. Assignment/least-loaded fairness and warm-worker reuse (every placement is
382
- today a cold one-shot `submit`) are deferred refinements, not yet built.
383
-
384
- ## 10. CLI
361
+ ## 9. SQL state layer
362
+
363
+ Client state lives under `$OMNIRUN_STATE_DIR` (default `~/.local/share/omnirun/`),
364
+ managed by a single `Store` repository (class in `state/store.py`) over
365
+ **SQLAlchemy Core 2.0**. On the laptop the engine points at
366
+ `$OMNIRUN_STATE_DIR/omnirun.db` (SQLite, zero-setup, the default). On a VPS or shared
367
+ server you point it at a Postgres database via `[state] url = "postgresql+psycopg://…"` —
368
+ the same schema and interface work on both dialects.
369
+
370
+ ### Schema hybrid document
371
+
372
+ Each table carries a primary key plus the few columns we filter or sort on, and a
373
+ `data` JSON column holding the full `model_dump(mode="json")` of the Pydantic domain
374
+ object (`JobRecord`, `ProviderFacts`, `QueueEntry`). Pydantic stays the serialization
375
+ source of truth; later field additions need no schema migration, only a `STATE_SCHEMA_VERSION`
376
+ bump in `meta`. On Postgres the `data` column is `JSONB` (for indexing performance);
377
+ on SQLite it is plain `JSON` this difference is handled once in `schema.py`
378
+ (`JSONText = JSON().with_variant(JSONB(), "postgresql")`), invisible at call sites.
379
+
380
+ ```
381
+ meta: key TEXT PK, value TEXT
382
+ jobs: job_id TEXT PK, name TEXT, backend TEXT, state TEXT,
383
+ submitted_at TEXT, schema_version INT, data JSON/JSONB
384
+ wait_samples: id INTEGER PK autoincr, backend TEXT, key TEXT,
385
+ wait_s REAL, recorded_at TEXT [index on (backend, key)]
386
+ facts: backend TEXT PK, discovered_at TEXT, ttl_s REAL, health TEXT, data JSON/JSONB
387
+ queue: qid TEXT PK, state TEXT, created_at TEXT, only_backend TEXT,
388
+ backend TEXT, job_id TEXT, data JSON/JSONB
389
+ ```
390
+
391
+ Timestamps are stored as ISO-8601 TEXT (portable and sortable).
392
+
393
+ ### Atomic `reserve_entry` — the #12 double-book guard
394
+
395
+ `Store.reserve_entry(qid, backend, cap)` is the only concurrency-critical operation.
396
+ It must guarantee that no two callers can both pass the `count_active(backend) < cap`
397
+ check and both flip to PLACING — which would over-book the backend.
398
+
399
+ - **SQLite**: `Store.__init__` installs a pair of SQLAlchemy engine events that disable
400
+ pysqlite's implicit `BEGIN` and emit `BEGIN IMMEDIATE` at transaction start instead
401
+ (`_install_sqlite_write_lock`). `BEGIN IMMEDIATE` acquires the reserved write lock
402
+ up front, serializing all `reserve_entry` calls sequentially. `with_for_update()` on
403
+ the re-read query is a no-op clause on SQLite — the serialization comes from the write
404
+ lock. A generous `busy_timeout` (30 s) makes contending callers wait rather than raise.
405
+
406
+ - **Postgres**: `engine.begin()` opens a transaction at the server default (READ COMMITTED).
407
+ A `FOR UPDATE` on the target row (qid) locks that single row, but leaves the count of
408
+ OTHER rows readable by concurrent transactions — two threads reserving different qids
409
+ can both read `count_active("x") == 1` before either commits, and both flip, over-booking.
410
+ The fix is a **transaction-scoped advisory lock per backend**:
411
+ ```sql
412
+ SELECT pg_advisory_xact_lock(hashtext(:b))
413
+ ```
414
+ issued at the top of the `reserve_entry` transaction. This serializes all reservers
415
+ for the same backend string; the second thread blocks until the first commits, then
416
+ re-reads the count and finds the cap full, returning False. The advisory lock is
417
+ auto-released at transaction commit/rollback. On PG 18.1 this eliminated 25/25
418
+ over-books observed without the guard (reproduction in `pg_overbook_raw.py`).
419
+
420
+ The `_upsert` helper is the one place where dialect branching lives:
421
+ `sqlite_insert(table).on_conflict_do_update(...)` vs `pg_insert(table).on_conflict_do_update(...)`.
422
+ Call sites are dialect-unaware.
423
+
424
+ ### Opening a Store
425
+
426
+ `open_store(url)` creates the engine, `create_all()` (idempotent), stamps
427
+ `meta["schema_version"] = 2`, and returns the `Store`. The `[state]` TOML section
428
+ controls which engine is used:
429
+
430
+ ```toml
431
+ [state]
432
+ backend = "sqlite" # or "postgres"
433
+ # path = "/custom/path/omnirun.db" # explicit SQLite path (default: state_dir/omnirun.db)
434
+ # url = "postgresql+psycopg://user:pw@host/db" # overrides backend/path
435
+ ```
436
+
437
+ A one-time JSON→SQL importer (`omnirun state migrate [--from DIR] [--dry-run]`) reads
438
+ the legacy `$OMNIRUN_STATE_DIR/jobs/*/meta.json`, `facts/*.json`, `queue/*.json`, and
439
+ `wait_history.json` files, tolerating `schema_version` 0 and 1, and upserts them into
440
+ the SQL store. `omnirun state path` prints the active database URL.
441
+
442
+ ## 10. Scheduler — Provider seam, pure tick, budget/deadline
443
+
444
+ > **Status (core branch).** What ships today is the **pure cheapest-fitting
445
+ > free→paid tick** plus the atomic reserve and reconcile/orphan-recovery described
446
+ > below. The **budget ledger, per-job deadline (`start_by`/`finish_by`), priority,
447
+ > `omnirun reprioritize`, and `omnirun budget` (daily/weekly caps)** are **deferred
448
+ > to a follow-up branch, pending a real need for spend control** — the `tick`
449
+ > signature below therefore takes no `ledger`, ranks by `submitted_at` only, and
450
+ > invariants **budget_safety** and **deadline_defense** are not active in core.
451
+
452
+ ### Provider seam
453
+
454
+ Between the pure `tick` and the eight concrete `Backend` implementations sits a
455
+ thin protocol layer (`providers/`):
456
+
457
+ ```python
458
+ class Provider(Protocol):
459
+ name: str
460
+ def discover(self) -> ProviderFacts: ...
461
+ def offer(self, req: ResourceSpec) -> list[Slot]: ...
462
+ def place(self, rec: JobRecord, slot: Slot) -> Placement: ...
463
+ def poll(self, p: Placement) -> Status: ...
464
+ def cancel(self, p: Placement, mode: CancelMode) -> None: ...
465
+ def stream_logs(self, p: Placement) -> Iterator[str]: ...
466
+ def collect_outputs(self, p: Placement, dest: Path) -> None: ...
467
+ def gc(self) -> None: ...
468
+ ```
469
+
470
+ `CancelMode` has two values: `GRACEFUL` (ask the job to stop cleanly, then
471
+ hard-kill after a `cancel_grace_s` window) and `FORCE` (tear it down
472
+ immediately). The `BackendProvider` adapter drives the uniform sequence:
473
+ GRACEFUL → poll the backend until terminal or `cancel_grace_s` elapses →
474
+ FORCE (SIGKILL) → **reap** the billable/worker resource (terminate the
475
+ marketplace instance / `scancel` / stop the kernel-session) via `Backend.gc`.
476
+ Cancel is idempotent and complete: after it returns there is no live placement
477
+ or billing instance (invariant 5), even for a job that already looked terminal.
478
+ `omnirun cancel --force` skips the grace window. The per-backend grace duration
479
+ defaults to 30 s and is overridable via `cancel_grace_s` in the backend's
480
+ config block.
481
+
482
+ SSH-family backends (`signal_job` in `backends/jobdir.py`) signal the whole
483
+ process group (the bootstrap records its `pgid` in `$JOB_DIR/pgid`): TERM on
484
+ graceful, KILL on force; Slurm uses `scancel` / `scancel -s KILL`. The shared
485
+ `.trees/<sha>` worktree and `.venv` are **never touched** — a job never owns
486
+ them. Marketplace/notebook backends reap the billing instance / kernel-session
487
+ as an idempotent side-effect of cancel, even when the job already looks
488
+ terminal. Kaggle has no cancel API; the adapter caches `CANCELLED` and logs a
489
+ note rather than raising.
490
+
491
+ **Streaming logs.** `stream_logs` tails the worker's canonical
492
+ `logs/bootstrap.log` (the one ordered merged stream that the bootstrap tees all
493
+ stdout/stderr through) on every backend, so `omnirun logs -f` is uniform.
494
+ Kaggle's batch API exposes a run log only once the kernel completes, so its
495
+ follow mode emits a one-line honesty note (`LIVE_TAIL_NOTE`) and then yields
496
+ the final dump — there is no live mid-run tail. A daemon-side ring buffer that
497
+ fans one stream to many remote followers is Phase 5.
498
+
499
+ `BackendProvider` (`providers/adapter.py`) is the **one bridge** from this seam to
500
+ today's eight `Backend` implementations — it wraps a single `Backend` + a shared
501
+ `Store` and adapts `Backend.probe` into `Slot`s and `Backend.submit` into
502
+ `place`. No backend rewrite was needed; `BackendProvider` is the tractability
503
+ hinge.
504
+
505
+ **At-least-once seam.** The `place`/persist boundary is at-least-once. Phase 4
506
+ closes the marketplace orphan window: `BackendProvider.place` threads
507
+ `on_provisioning` so a billable handle is persisted onto the PLACING placement
508
+ before `place` returns, and `Control._reconcile` ADOPTS (re-polls) a
509
+ partial-handle PLACING instead of reverting and relaunching. The remaining
510
+ concurrent-tick lease (two overlapping ticks reverting each other's fresh
511
+ reservation) is Phase 5.
512
+
513
+ ### Pure tick: `(jobs, slots, ledger, now) -> decisions`
514
+
515
+ `scheduler.tick` is a **pure function** — no I/O, `now` is a parameter (making
516
+ it deterministic and testable), no wall-clock, no backend names, imports only
517
+ `models` and `budget`. Fit is decided solely by
518
+ `slot.capabilities.satisfies(req)`.
519
+
520
+ ```python
521
+ def tick(
522
+ jobs: list[JobRecord],
523
+ slots: list[Slot],
524
+ ledger: BudgetLedger,
525
+ now: datetime,
526
+ *,
527
+ policy: SchedPolicy | None = None,
528
+ ) -> list[Decision]:
529
+ ```
530
+
531
+ **Per-tick match order (for each pending job, priority+urgency ranked):**
532
+
533
+ 1. *Admit/HELD* — if slots exist and none of their capabilities can ever
534
+ satisfy the resource request, the job is HELD (not queued) so the tick
535
+ reports the reason. With no slots we can't prove impossibility, so the
536
+ job stays QUEUED.
537
+ 2. *4a: free-first* — best free slot (smallest `wait_s`) that meets the
538
+ `finish_by` deadline.
539
+ 3. *4b: last-responsible-moment paid escalation* — if no free slot met the
540
+ deadline, pick the cheapest affordable paid slot that does, checked
541
+ against a LOCAL working ledger (this tick's prior paid commitments
542
+ count, so the total paid in one tick never exceeds the cap).
543
+ 4. *4c: run-late liveness* — if neither 4a nor 4b worked (no slot can meet
544
+ the deadline), place on the best FREE slot IGNORING the deadline.
545
+ **A job is never refused for cost** — it waits for free capacity.
546
+
547
+ `SchedPolicy.allow_paid = False` skips step 4b; the job runs free or waits.
548
+
549
+ ### Budget ledger and per-job policy
550
+
551
+ `JobPolicy` (carried in `JobSpec.policy`) holds:
552
+ - `deadline: Deadline | None` — optional `start_by` and/or `finish_by` (UTC).
553
+ - `max_cost: float | None` — **per-job USD ceiling** for a single placement.
554
+ **Breaking change from the old chooser flag:** `--max-cost` is now a
555
+ per-job total cost ceiling (works best with `--time` to bound the estimate),
556
+ NOT a filter that drops offers above a cost threshold. Existing uses of
557
+ `--max-cost` as a coarse filter should be migrated to per-backend
558
+ `max_hourly` in the config.
559
+ - `priority: int` — higher = scheduled sooner; reprioritizable live via
560
+ `omnirun reprioritize`.
561
+
562
+ **Budget ledger** (`budget.py`, `BudgetLedger`): pure, immutable, two kinds of
563
+ entries — `committed` (at placement) and `spent` (at completion). Two calendar-
564
+ aligned windows:
565
+ - `"day"` — entries where `entry.at.date() == now.date()` (UTC).
566
+ - `"week"` — entries in the same ISO year+week.
567
+
568
+ **BOTH windows are enforced simultaneously** (one wallet, two gates):
569
+
570
+ - The **daily cap** is the tick's primary window: `ledger.can_afford(cost, now)`
571
+ gates paid escalation in step 4b.
572
+ - The **weekly cap** is enforced alongside by `Control._enact_place`: before
573
+ reserving, if the weekly ledger cannot afford the estimate, the place is
574
+ skipped (job stays QUEUED, retries a later tick / next week).
575
+
576
+ A paid job is blocked if it cannot afford EITHER window; it runs free/late
577
+ instead (liveness: a job is delayed, never permanently failed). `omnirun budget
578
+ --daily $D --weekly $W` sets both caps; `omnirun budget` (no flags) shows
579
+ `spent` and `cap` for each window.
580
+
581
+ ### The impure `Control` driver
582
+
583
+ `Control` (`control.py`) is `tick`'s impure counterpart — it does all the I/O
584
+ that `tick` cannot:
585
+
586
+ ```
587
+ run_tick(now):
588
+ 1. reconcile: poll each PLACING/RUNNING job's provider → update job state
589
+ (terminal → realize budget; LOST → requeue with attempts+1)
590
+ 2. gather slots: ask each provider.offer() for every distinct pending req
591
+ 3. load ledger: Store.load_ledger(window, cap, now)
592
+ 4. tick: pure function → decisions
593
+ 5. enact: for each decision:
594
+ hold → save HELD
595
+ place → Store.reserve (atomic, #12 guard) → PLACING stub →
596
+ provider.place → commit budget → save RUNNING
597
+ ```
598
+
599
+ `Store.reserve` (the §9 `reserve_entry` guard) flips QUEUED/HELD → PLACING
600
+ and writes a stub `Placement` in ONE transaction — only one concurrent tick
601
+ can win the race; the other sees the state already PLACING and skips it.
602
+
603
+ **Daemonless vs daemon — one tick everywhere.** Direct `omnirun submit` stores
604
+ the spec QUEUED and runs exactly ONE synchronous tick; no background process is
605
+ required (a placed job then runs on the backend while the laptop is off).
606
+ `omnirun serve` runs the same `Control.run_tick` on a background thread every
607
+ `poll_interval_s` (default 10 s). The tick logic is identical; only the cadence
608
+ differs.
609
+
610
+ ### The 8 correctness invariants (`tests/test_scheduler_invariants.py`)
611
+
612
+ A Hypothesis `RuleBasedStateMachine` drives the real `Control` + SQLite `Store`
613
+ + deterministic `FlakyProvider` doubles through random interleavings of
614
+ `submit / run_tick / provider_responds / provider_fails / cancel / advance_time`,
615
+ asserting all eight invariants after EVERY step:
616
+
617
+ 1. **budget_safety** — committed+spent ≤ cap; per-job ≤ `max_cost`; free
618
+ slots cost 0.
619
+ 2. **admission_soundness** — every live placement's provider can satisfy the
620
+ resource request (§8 hit-guard).
621
+ 3. **concurrency_safety** — non-terminal placements per provider ≤ its
622
+ `max_parallel` cap (the §9 `Store.reserve` atomic guard).
623
+ 4. **liveness_no_silent_loss** — a non-cancelled job is always in a live or
624
+ terminal state (reconcile requeues LOST, no silent drops).
625
+ 5. **cancellation_completeness** — a cancelled job has zero live placements
626
+ and is never re-placed by a later tick (§7).
627
+ 6. **deadline_defense** — no paid placement while a fitting free slot met the
628
+ deadline (4a takes priority over 4b).
629
+ 7. **crash_isolation** — a failing provider never crashes the tick nor blocks
630
+ healthy providers.
631
+ 8. **tick_convergence** — a second identical tick creates no new placements
632
+ (idempotency).
633
+
634
+ **At-least-once caveat (I2 note):** assertions are store-level only (job states,
635
+ ledger totals, `count_active_jobs ≤ cap`). The fakes do not model backend
636
+ instances, so "exactly one live backend instance per job" is NOT asserted —
637
+ that property is knowingly false across a `place`-failure boundary until
638
+ Phase 4's orphan-recovery lands.
639
+
640
+ ## 11. Queue & scheduler daemon (optional)
641
+
642
+ Direct `submit` stays daemonless (§10). For fan-out — many jobs, or spreading a
643
+ batch across backends with per-backend concurrency caps — an *optional*
644
+ long-lived scheduler daemon is available (`daemon.py`). It uses the SAME
645
+ `Control.run_tick` as the daemonless path.
646
+
647
+ - **`omnirun serve`** runs the daemon in the foreground: a localhost TCP socket
648
+ (default `127.0.0.1:8787`), newline-delimited JSON request/response, plus a
649
+ scheduler thread ticking every `poll_interval_s` (default 10s). It owns a
650
+ durable queue persisted via the SQL `Store`; a restart re-reads and resumes.
651
+ - **Scheduler tick** (§10): reconcile → gather slots → load ledger → pure tick
652
+ → enact. Placement reserves a backend slot atomically (§9 guard, state →
653
+ PLACING) before dispatching `provider.place` on a thread pool; it backfills
654
+ as jobs complete. Each backend's `max_parallel` caps its concurrent
655
+ non-terminal jobs.
656
+ - **Job lifecycle**: QUEUED → (HELD) → PLACING → RUNNING → SUCCEEDED / FAILED /
657
+ CANCELLED.
658
+ - **Commands**: `enqueue [--count N] [--backend NAME] -- CMD...`, `queue` (show the table),
659
+ `queue --wait` (poll until all terminal), `queue --cancel <qid|all>`; the
660
+ socket also speaks `ping` / `list` / `shutdown`.
661
+ - **Placement is greedy** — favors free-first, then cheapest-affordable-paid.
662
+ Assignment/least-loaded fairness and warm-worker reuse (every placement is
663
+ still a cold one-shot `provider.place`) are deferred refinements.
664
+
665
+ ## 12. CLI
385
666
 
386
667
  ```
387
668
  omnirun submit [--name N] [--gpus 1] [--gpu-type H100 | --vram 40] [--time 15h]
388
- [--backend uni] [--yes] [--max-cost 30] [--dirty] -- python train.py ...
669
+ [--backend uni] [--push] -- python train.py ...
389
670
  omnirun offers [same resource flags] # probe & table, no submit
390
671
  omnirun ps # all known jobs, refreshed statuses
391
672
  omnirun status <job> | logs [-f] <job> | cancel <job> | pull <job> [dest]
392
673
  omnirun backends check # config + connectivity sanity per backend
674
+ omnirun backends discover # probe live capability/health; cache facts
393
675
  omnirun gc # reap finished job dirs, leaked instances
394
- omnirun serve [--host H] [--port P] # run the scheduler daemon (optional, §9)
395
- omnirun enqueue [resource flags] [--count N] [--backend NAME] -- CMD... # queue a job
676
+ omnirun serve [--host H] [--port P] # run the scheduler daemon (optional, §11)
677
+ omnirun enqueue [resource flags] [--count N] [--backend NAME] -- CMD...
396
678
  omnirun queue [--wait] [--cancel qid|all] # inspect / wait on / cancel the daemon queue
397
679
  ```
398
680
 
399
- ## 11. Implementation notes
681
+ ## 13. Implementation notes
400
682
 
401
683
  - Python ≥3.12, deps: `typer`, `rich`, `httpx`, `pydantic`, `kaggle` (thin API
402
684
  client), `google-colab-cli` (optional extra). Marketplaces via plain REST
@@ -424,16 +706,22 @@ omnirun queue [--wait] [--cancel qid|all] # inspect / wait on / cancel the daemo
424
706
  ```
425
707
  src/omnirun/
426
708
  models.py # JobSpec, ResourceSpec, EnvSpec/EnvKind (auto|uv|pip|conda|system|none),
427
- # Offer, JobStatus, JobHandle
709
+ # Offer, JobStatus, JobHandle, JobState, Placement, Slot,
710
+ # Cost, Availability, Deadline, JobPolicy, Decision
428
711
  config.py # TOML load, backend registry (project_root, gpu_map, max_parallel),
429
712
  # PolicyConfig, DaemonConfig
430
713
  repo.py # git state, clean/pushed checks, RepoRef, env_file, bundle creation,
431
714
  # remote_clone_plan/worker_clone_url/remote_is_public (public-repo direct clone)
432
715
  bootstrap.py # bootstrap.sh generation (shared payload), notebook_env_spec
433
- store.py # $OMNIRUN_STATE_DIR/jobs/<id>/meta.json
434
- queue.py # durable QueueStore/QueueEntry backing the scheduler
716
+ scheduler.py # pure tick(jobs, slots, now) -> decisions; SchedPolicy
717
+ control.py # impure Control driver: reconcile → gather → tick → enact
718
+ providers/ # base.py (Provider protocol, CancelMode),
719
+ # adapter.py (BackendProvider: Backend+Store → Provider seam)
720
+ state/ # store.py (Store, open_store, reserve, ledger_add/realize),
721
+ # schema.py, migrate.py
722
+ queue.py # durable QueueStore/QueueEntry backing the daemon
435
723
  daemon.py # optional localhost scheduler daemon (serve/enqueue/queue)
436
- chooser.py # parallel probing, ranking, offer table
724
+ chooser.py # parallel probing, ranking, offer table (display; tick does placement)
437
725
  execlayer/ # base.py (Exec protocol), local.py, ssh.py (ControlMaster, login_shell)
438
726
  backends/ # base.py, jobdir.py (shared job-dir/project-root/push/.env/status
439
727
  # helpers), local.py, ssh.py, slurm.py, kaggle.py, colab.py,
@@ -441,14 +729,16 @@ omnirun queue [--wait] [--cancel qid|all] # inspect / wait on / cancel the daemo
441
729
  # vast.py, thunder.py
442
730
  cli.py # typer app
443
731
  ```
444
- - Testing tiers: (1) pure unit (codegen, ranking, repo state); (2) e2e through
445
- `local` backend full submit→run→pull without network; (3) dockerized sshd and
446
- slurm cluster integration tests, opt-in; (4) live backends — needs user creds.
732
+ - Testing tiers: (1) pure unit (codegen, ranking, repo state, scheduler tick); (2) Hypothesis stateful invariant suite (`tests/test_scheduler_invariants.py`,
733
+ 8 invariants over real Control + SQLite Store + Fake/Flaky providers); (3) e2e through
734
+ `local` backend full submit→run→pull without network; (4) dockerized sshd and
735
+ slurm cluster integration tests, opt-in; (5) live backends — needs user creds.
447
736
 
448
- ## 12. Non-goals
737
+ ## 14. Non-goals
449
738
 
450
739
  Still out of scope: data syncing, artifact versioning, DAGs/pipelines, multi-node jobs,
451
740
  spot-preemption recovery, image building, a web UI. Cross-backend queueing is **no longer**
452
- a non-goal — the optional scheduler daemon (§9) provides it. Placement fairness
453
- (least-loaded/assignment rather than greedy) and warm-worker reuse are *planned but not yet
454
- built*, not non-goals.
741
+ a non-goal — the scheduler daemon (§11) provides it, running the same `Control.run_tick`.
742
+ Placement fairness (least-loaded/assignment rather than greedy) and warm-worker reuse
743
+ are *planned but not yet built*, not non-goals. Exact-once marketplace orphan-recovery
744
+ (`on_provisioning` + reconcile ADOPT) landed in Phase 4.