omnirun 0.2.2__tar.gz → 0.2.3__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 (68) hide show
  1. {omnirun-0.2.2 → omnirun-0.2.3}/PKG-INFO +1 -1
  2. {omnirun-0.2.2 → omnirun-0.2.3}/pyproject.toml +1 -1
  3. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/__init__.py +1 -1
  4. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/backends/base.py +18 -1
  5. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/backends/colab.py +7 -2
  6. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/backends/kaggle.py +7 -2
  7. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/backends/local.py +7 -2
  8. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/backends/marketplace.py +62 -12
  9. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/backends/slurm.py +7 -2
  10. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/backends/ssh.py +7 -2
  11. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/cli.py +18 -8
  12. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/daemon.py +16 -9
  13. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/repo.py +62 -0
  14. {omnirun-0.2.2 → omnirun-0.2.3}/tests/test_cli.py +77 -3
  15. {omnirun-0.2.2 → omnirun-0.2.3}/tests/test_marketplaces.py +113 -0
  16. {omnirun-0.2.2 → omnirun-0.2.3}/tests/test_queue.py +7 -2
  17. {omnirun-0.2.2 → omnirun-0.2.3}/tests/test_repo.py +52 -1
  18. {omnirun-0.2.2 → omnirun-0.2.3}/uv.lock +1 -1
  19. {omnirun-0.2.2 → omnirun-0.2.3}/.envrc +0 -0
  20. {omnirun-0.2.2 → omnirun-0.2.3}/.github/workflows/checks.yml +0 -0
  21. {omnirun-0.2.2 → omnirun-0.2.3}/.github/workflows/ci.yml +0 -0
  22. {omnirun-0.2.2 → omnirun-0.2.3}/.github/workflows/publish.yml +0 -0
  23. {omnirun-0.2.2 → omnirun-0.2.3}/.gitignore +0 -0
  24. {omnirun-0.2.2 → omnirun-0.2.3}/.python-version +0 -0
  25. {omnirun-0.2.2 → omnirun-0.2.3}/AGENTS.md +0 -0
  26. {omnirun-0.2.2 → omnirun-0.2.3}/CLAUDE.md +0 -0
  27. {omnirun-0.2.2 → omnirun-0.2.3}/DESIGN.md +0 -0
  28. {omnirun-0.2.2 → omnirun-0.2.3}/README.md +0 -0
  29. {omnirun-0.2.2 → omnirun-0.2.3}/TESTING.md +0 -0
  30. {omnirun-0.2.2 → omnirun-0.2.3}/flake.lock +0 -0
  31. {omnirun-0.2.2 → omnirun-0.2.3}/flake.nix +0 -0
  32. {omnirun-0.2.2 → omnirun-0.2.3}/main.py +0 -0
  33. {omnirun-0.2.2 → omnirun-0.2.3}/research/colab-kaggle.md +0 -0
  34. {omnirun-0.2.2 → omnirun-0.2.3}/research/gpu-marketplaces.md +0 -0
  35. {omnirun-0.2.2 → omnirun-0.2.3}/research/launcher-landscape.md +0 -0
  36. {omnirun-0.2.2 → omnirun-0.2.3}/research/ml-job-queues.md +0 -0
  37. {omnirun-0.2.2 → omnirun-0.2.3}/research/slurm-ssh.md +0 -0
  38. {omnirun-0.2.2 → omnirun-0.2.3}/scripts/check-version.sh +0 -0
  39. {omnirun-0.2.2 → omnirun-0.2.3}/scripts/pre-push +0 -0
  40. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/backends/__init__.py +0 -0
  41. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/backends/jobdir.py +0 -0
  42. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/backends/runpod.py +0 -0
  43. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/backends/tarsafe.py +0 -0
  44. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/backends/thunder.py +0 -0
  45. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/backends/vast.py +0 -0
  46. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/bootstrap.py +0 -0
  47. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/chooser.py +0 -0
  48. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/config.py +0 -0
  49. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/execlayer/__init__.py +0 -0
  50. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/execlayer/base.py +0 -0
  51. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/execlayer/local.py +0 -0
  52. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/execlayer/ssh.py +0 -0
  53. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/models.py +0 -0
  54. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/queue.py +0 -0
  55. {omnirun-0.2.2 → omnirun-0.2.3}/src/omnirun/store.py +0 -0
  56. {omnirun-0.2.2 → omnirun-0.2.3}/tests/__init__.py +0 -0
  57. {omnirun-0.2.2 → omnirun-0.2.3}/tests/conftest.py +0 -0
  58. {omnirun-0.2.2 → omnirun-0.2.3}/tests/test_bootstrap.py +0 -0
  59. {omnirun-0.2.2 → omnirun-0.2.3}/tests/test_chooser.py +0 -0
  60. {omnirun-0.2.2 → omnirun-0.2.3}/tests/test_colab.py +0 -0
  61. {omnirun-0.2.2 → omnirun-0.2.3}/tests/test_config.py +0 -0
  62. {omnirun-0.2.2 → omnirun-0.2.3}/tests/test_kaggle.py +0 -0
  63. {omnirun-0.2.2 → omnirun-0.2.3}/tests/test_local_backend.py +0 -0
  64. {omnirun-0.2.2 → omnirun-0.2.3}/tests/test_slurm.py +0 -0
  65. {omnirun-0.2.2 → omnirun-0.2.3}/tests/test_ssh_backend.py +0 -0
  66. {omnirun-0.2.2 → omnirun-0.2.3}/tests/test_ssh_exec.py +0 -0
  67. {omnirun-0.2.2 → omnirun-0.2.3}/tests/test_store.py +0 -0
  68. {omnirun-0.2.2 → omnirun-0.2.3}/tests/test_tarsafe.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: omnirun
3
- Version: 0.2.2
3
+ Version: 0.2.3
4
4
  Summary: Run jobs from your repo anywhere: Slurm over SSH, any SSH box, Kaggle, Colab, or marketplace GPUs
5
5
  Author-email: Dmitrii Mukhutdinov <flyingleafe@gmail.com>
6
6
  License: MIT
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "omnirun"
3
- version = "0.2.2"
3
+ version = "0.2.3"
4
4
  description = "Run jobs from your repo anywhere: Slurm over SSH, any SSH box, Kaggle, Colab, or marketplace GPUs"
5
5
  readme = "README.md"
6
6
  requires-python = ">=3.12"
@@ -1,3 +1,3 @@
1
1
  """omnirun — run jobs from your repo anywhere."""
2
2
 
3
- __version__ = "0.2.2"
3
+ __version__ = "0.2.3"
@@ -22,6 +22,14 @@ class BackendError(RuntimeError):
22
22
  """Raised for backend-level failures with a user-actionable message."""
23
23
 
24
24
 
25
+ #: Optional hook submit() calls the instant it creates a *billable* resource,
26
+ #: before the (possibly minutes-long) wait for it to become usable. It hands the
27
+ #: client a partial JobHandle so a recovery stub can be persisted immediately —
28
+ #: an interrupted submit then stays visible to `ps`/`gc` instead of orphaning a
29
+ #: running instance with no local record.
30
+ ProvisioningSink = Callable[[JobHandle], None]
31
+
32
+
25
33
  class Backend(ABC):
26
34
  """One configured execution target (a cluster, a machine, a provider account).
27
35
 
@@ -46,7 +54,16 @@ class Backend(ABC):
46
54
  def probe(self, res: ResourceSpec) -> list[Offer]: ...
47
55
 
48
56
  @abstractmethod
49
- def submit(self, spec: JobSpec, offer: Offer) -> JobHandle: ...
57
+ def submit(
58
+ self,
59
+ spec: JobSpec,
60
+ offer: Offer,
61
+ on_provisioning: ProvisioningSink | None = None,
62
+ ) -> JobHandle:
63
+ """Run the job and return a handle. If a billable resource is created
64
+ before the handle is ready, call on_provisioning with a partial handle
65
+ first so the client can persist a recovery stub (see ProvisioningSink)."""
66
+ ...
50
67
 
51
68
  @abstractmethod
52
69
  def status(self, handle: JobHandle) -> StatusReport: ...
@@ -33,7 +33,7 @@ from pathlib import Path
33
33
  from typing import Any
34
34
 
35
35
  from omnirun.backends import jobdir, tarsafe
36
- from omnirun.backends.base import Backend, BackendError, register
36
+ from omnirun.backends.base import Backend, BackendError, ProvisioningSink, register
37
37
  from omnirun.bootstrap import (
38
38
  HEARTBEAT_STALE_S,
39
39
  BootstrapParams,
@@ -329,7 +329,12 @@ class ColabBackend(Backend):
329
329
 
330
330
  # ---- submit -----------------------------------------------------------------
331
331
 
332
- def submit(self, spec: JobSpec, offer: Offer) -> JobHandle:
332
+ def submit(
333
+ self,
334
+ spec: JobSpec,
335
+ offer: Offer,
336
+ on_provisioning: ProvisioningSink | None = None,
337
+ ) -> JobHandle:
333
338
  session = self._session(spec.job_id)
334
339
  job_dir = f"{COLAB_ROOT}/jobs/{spec.job_id}"
335
340
 
@@ -33,7 +33,7 @@ from datetime import datetime, timedelta, timezone
33
33
  from pathlib import Path
34
34
  from typing import Any
35
35
 
36
- from omnirun.backends.base import Backend, BackendError, register
36
+ from omnirun.backends.base import Backend, BackendError, ProvisioningSink, register
37
37
  from omnirun.backends import jobdir, tarsafe
38
38
  from omnirun.bootstrap import (
39
39
  BootstrapParams,
@@ -393,7 +393,12 @@ class KaggleBackend(Backend):
393
393
  "the backend's `max_source_bytes` if Kaggle accepts more"
394
394
  )
395
395
 
396
- def submit(self, spec: JobSpec, offer: Offer) -> JobHandle:
396
+ def submit(
397
+ self,
398
+ spec: JobSpec,
399
+ offer: Offer,
400
+ on_provisioning: ProvisioningSink | None = None,
401
+ ) -> JobHandle:
397
402
  api = self._api()
398
403
  user = self._username(api)
399
404
  job_id = spec.job_id
@@ -13,7 +13,7 @@ from pathlib import Path
13
13
  from typing import TYPE_CHECKING
14
14
 
15
15
  from omnirun.backends import jobdir
16
- from omnirun.backends.base import Backend, register
16
+ from omnirun.backends.base import Backend, ProvisioningSink, register
17
17
  from omnirun.bootstrap import BootstrapParams
18
18
  from omnirun.execlayer.base import shell_quote
19
19
  from omnirun.execlayer.local import LocalExec
@@ -110,7 +110,12 @@ class LocalBackend(Backend):
110
110
 
111
111
  # --- lifecycle ---------------------------------------------------------
112
112
 
113
- def submit(self, spec: JobSpec, offer: Offer) -> JobHandle:
113
+ def submit(
114
+ self,
115
+ spec: JobSpec,
116
+ offer: Offer,
117
+ on_provisioning: ProvisioningSink | None = None,
118
+ ) -> JobHandle:
114
119
  root = jobdir.remote_root(self.exec, self.config.root)
115
120
  project_root = jobdir.resolve_project_root(
116
121
  self.exec,
@@ -7,6 +7,13 @@ the job (client-side git push + bootstrap.sh) -> detached launch (setsid+nohup,
7
7
  pidfile), exactly like the plain ssh backend.
8
8
 
9
9
  Termination policy (billable resources must not leak):
10
+ - **Provisioning stub** (issue #7): the instant ``submit`` rents an instance —
11
+ before the minutes-long provisioning wait — it calls the ``on_provisioning``
12
+ hook with a partial handle carrying the instance id, so the client persists a
13
+ recovery record *first*. A submit killed mid-wait then stays visible to
14
+ ``ps``/``status`` (reported as PROVISIONING with a billing warning) and
15
+ reclaimable by ``omnirun gc --all``, instead of orphaning a running instance
16
+ with no local trace. The final handle overwrites the stub on success.
10
17
  - **Client-side auto-terminate** (``auto_terminate``, default true): the instance
11
18
  is destroyed in ``pull_outputs`` (after a successful pull), in ``cancel``, and
12
19
  in ``gc``. ``status`` never terminates — when the job is terminal it appends a
@@ -52,7 +59,7 @@ import httpx
52
59
  from pydantic import BaseModel, Field
53
60
 
54
61
  from omnirun.backends import jobdir
55
- from omnirun.backends.base import Backend, BackendError
62
+ from omnirun.backends.base import Backend, BackendError, ProvisioningSink
56
63
  from omnirun.bootstrap import BootstrapParams
57
64
  from omnirun.execlayer.base import Exec, shell_quote
58
65
  from omnirun.repo import local_root_of
@@ -270,9 +277,26 @@ class MarketplaceBackend(Backend, ABC):
270
277
  unfit_reasons=[reason],
271
278
  )
272
279
 
273
- def submit(self, spec: JobSpec, offer: Offer) -> JobHandle:
280
+ def submit(
281
+ self,
282
+ spec: JobSpec,
283
+ offer: Offer,
284
+ on_provisioning: ProvisioningSink | None = None,
285
+ ) -> JobHandle:
274
286
  inst = self._create_instance(spec, offer)
275
287
  instance_id = inst.instance_id
288
+ # The instance is now billing but the handle isn't ready yet. Hand the
289
+ # client a stub carrying the instance id *before* the provisioning wait,
290
+ # so an interrupted submit (a kill mid-wait) stays visible to ps/gc
291
+ # instead of orphaning a running instance with no local record (#7).
292
+ if on_provisioning is not None:
293
+ on_provisioning(
294
+ JobHandle(
295
+ backend=self.name,
296
+ job_id=spec.job_id,
297
+ data={"instance_id": instance_id, "provisioning": True},
298
+ )
299
+ )
276
300
  try:
277
301
  inst = self._wait_provisioned(instance_id)
278
302
  assert inst.ssh_target is not None
@@ -379,6 +403,31 @@ class MarketplaceBackend(Backend, ABC):
379
403
 
380
404
  def status(self, handle: JobHandle) -> StatusReport:
381
405
  instance_id = handle.data["instance_id"]
406
+ if not handle.data.get("job_dir"):
407
+ # A provisioning stub: submit was interrupted before it staged the
408
+ # job, so only the instance id was persisted. Report whether the
409
+ # rented instance is still billing so the user knows to reclaim it.
410
+ try:
411
+ inst = self._get_instance(instance_id)
412
+ except BackendError as e:
413
+ return StatusReport(
414
+ status=JobStatus.PROVISIONING,
415
+ detail=f"submit was interrupted while provisioning instance "
416
+ f"{instance_id}; could not reach the {self.name} API to check "
417
+ f"it ({e}) — verify in the provider console, run `omnirun gc`",
418
+ )
419
+ if inst is None:
420
+ return StatusReport(
421
+ status=JobStatus.LOST,
422
+ detail=f"instance {instance_id} never finished provisioning and "
423
+ "no longer exists",
424
+ )
425
+ return StatusReport(
426
+ status=JobStatus.PROVISIONING,
427
+ detail=f"submit was interrupted after renting instance {instance_id} "
428
+ f"(now {inst.status or 'unknown'}) — it is still billing; run "
429
+ f"`omnirun gc --all` (or cancel {handle.job_id}) to destroy it",
430
+ )
382
431
  inst: Instance | None = None
383
432
  api_error: str | None = None
384
433
  try:
@@ -436,16 +485,17 @@ class MarketplaceBackend(Backend, ABC):
436
485
  return paths
437
486
 
438
487
  def cancel(self, handle: JobHandle) -> None:
439
- try: # best-effort remote kill; instance dies right after anyway
440
- ex = self._exec_from_handle(handle)
441
- q = shell_quote(handle.data["job_dir"])
442
- ex.run(
443
- f"if [ -f {q}/pid ]; then p=$(cat {q}/pid); "
444
- f'kill -TERM -- "-$p" 2>/dev/null || kill -TERM "$p" 2>/dev/null; fi; true',
445
- timeout=30,
446
- )
447
- except Exception:
448
- pass
488
+ if handle.data.get("job_dir"): # a provisioning stub has nothing to kill
489
+ try: # best-effort remote kill; instance dies right after anyway
490
+ ex = self._exec_from_handle(handle)
491
+ q = shell_quote(handle.data["job_dir"])
492
+ ex.run(
493
+ f"if [ -f {q}/pid ]; then p=$(cat {q}/pid); "
494
+ f'kill -TERM -- "-$p" 2>/dev/null || kill -TERM "$p" 2>/dev/null; fi; true',
495
+ timeout=30,
496
+ )
497
+ except Exception:
498
+ pass
449
499
  instance_id = handle.data["instance_id"]
450
500
  if self._get_instance(instance_id) is not None:
451
501
  self._terminate(instance_id)
@@ -17,7 +17,7 @@ from datetime import datetime, timedelta
17
17
  from pathlib import Path
18
18
 
19
19
  from omnirun.backends import jobdir
20
- from omnirun.backends.base import Backend, BackendError, register
20
+ from omnirun.backends.base import Backend, BackendError, ProvisioningSink, register
21
21
  from omnirun.bootstrap import BootstrapParams, generate_bootstrap
22
22
  from omnirun.config import BackendConfig
23
23
  from omnirun.execlayer.base import Exec, ExecError, shell_quote
@@ -279,7 +279,12 @@ class SlurmBackend(Backend):
279
279
  )
280
280
  return f"{sbatch}{sep}{bootstrap}"
281
281
 
282
- def submit(self, spec: JobSpec, offer: Offer | None = None) -> JobHandle:
282
+ def submit(
283
+ self,
284
+ spec: JobSpec,
285
+ offer: Offer | None = None,
286
+ on_provisioning: ProvisioningSink | None = None,
287
+ ) -> JobHandle:
283
288
  ex = self.exec_
284
289
  root = jobdir.remote_root(ex, self.config.root)
285
290
  project_root = jobdir.resolve_project_root(
@@ -10,7 +10,7 @@ from collections.abc import Iterator
10
10
  from pathlib import Path
11
11
 
12
12
  from omnirun.backends import jobdir
13
- from omnirun.backends.base import Backend, BackendError, register
13
+ from omnirun.backends.base import Backend, BackendError, ProvisioningSink, register
14
14
  from omnirun.bootstrap import BootstrapParams
15
15
  from omnirun.execlayer.base import Exec, ExecError, shell_quote
16
16
  from omnirun.execlayer.ssh import RECONNECT_HINT, SSHExec
@@ -213,7 +213,12 @@ class SshBackend(Backend):
213
213
 
214
214
  # --- submit / lifecycle ---------------------------------------------------
215
215
 
216
- def submit(self, spec: JobSpec, offer: Offer | None = None) -> JobHandle:
216
+ def submit(
217
+ self,
218
+ spec: JobSpec,
219
+ offer: Offer | None = None,
220
+ on_provisioning: ProvisioningSink | None = None,
221
+ ) -> JobHandle:
217
222
  ex = self.exec_
218
223
  root = jobdir.remote_root(ex, self.config.root)
219
224
  project_root = jobdir.resolve_project_root(
@@ -27,6 +27,7 @@ from omnirun.config import (
27
27
  )
28
28
  from omnirun.models import (
29
29
  EnvSpec,
30
+ JobHandle,
30
31
  JobRecord,
31
32
  JobSpec,
32
33
  JobStatus,
@@ -451,15 +452,24 @@ def submit(
451
452
  _render_payload(backend_obj, spec, picked.offer)
452
453
  return
453
454
 
454
- handle = backend_obj.submit(spec, picked.offer)
455
- JobStore().save(
456
- JobRecord(
457
- spec=spec,
458
- handle=handle,
459
- offer=picked.offer,
460
- submitted_at=datetime.now(timezone.utc),
455
+ store = JobStore()
456
+ picked_offer = picked.offer
457
+
458
+ def _persist(h: JobHandle) -> None:
459
+ # Called once with a provisioning stub (if the backend rents something
460
+ # before the handle is ready) and again with the final handle, so an
461
+ # interrupted submit still leaves a reclaimable record (#7).
462
+ store.save(
463
+ JobRecord(
464
+ spec=spec,
465
+ handle=h,
466
+ offer=picked_offer,
467
+ submitted_at=datetime.now(timezone.utc),
468
+ )
461
469
  )
462
- )
470
+
471
+ handle = backend_obj.submit(spec, picked_offer, on_provisioning=_persist)
472
+ _persist(handle)
463
473
  console.print(f"[green]submitted[/green] {spec.job_id} -> {picked.offer.label}")
464
474
  console.print(f"follow logs with: omnirun logs -f {spec.job_id}")
465
475
 
@@ -26,6 +26,7 @@ from omnirun import chooser
26
26
  from omnirun.backends.base import Backend, make_backend
27
27
  from omnirun.config import Config
28
28
  from omnirun.models import (
29
+ JobHandle,
29
30
  JobRecord,
30
31
  JobSpec,
31
32
  JobStatus,
@@ -433,19 +434,25 @@ class Daemon:
433
434
  if be is None:
434
435
  self._fail_submit(qid, f"backend {backend_name!r} unavailable")
435
436
  return
437
+
438
+ def _persist(h: JobHandle) -> None:
439
+ # Persist a stub the instant a billable resource is created, then
440
+ # the final handle — an interrupted submit stays reclaimable (#7).
441
+ self._jobstore.save(
442
+ JobRecord(
443
+ spec=spec,
444
+ handle=h,
445
+ offer=offer,
446
+ submitted_at=datetime.now(timezone.utc),
447
+ )
448
+ )
449
+
436
450
  try:
437
- handle = be.submit(spec, offer)
451
+ handle = be.submit(spec, offer, on_provisioning=_persist)
438
452
  except Exception as e:
439
453
  self._retry_or_fail(qid, str(e))
440
454
  return
441
- self._jobstore.save(
442
- JobRecord(
443
- spec=spec,
444
- handle=handle,
445
- offer=offer,
446
- submitted_at=datetime.now(timezone.utc),
447
- )
448
- )
455
+ _persist(handle)
449
456
  with self._lock:
450
457
  entry = self._store.get(qid)
451
458
  if entry is None:
@@ -11,6 +11,7 @@ import os
11
11
  import re
12
12
  import shutil
13
13
  import subprocess
14
+ import tempfile
14
15
  from pathlib import Path
15
16
 
16
17
  from omnirun.models import RepoRef
@@ -32,6 +33,60 @@ def _git(root: Path, *args: str) -> subprocess.CompletedProcess[str]:
32
33
  )
33
34
 
34
35
 
36
+ # Identity stamped on wip commits so `git commit-tree` never fails on a repo
37
+ # without a configured user.name/email; the commit is ephemeral anyway.
38
+ _WIP_IDENTITY = {
39
+ "GIT_AUTHOR_NAME": "omnirun",
40
+ "GIT_AUTHOR_EMAIL": "wip@omnirun.invalid",
41
+ "GIT_COMMITTER_NAME": "omnirun",
42
+ "GIT_COMMITTER_EMAIL": "wip@omnirun.invalid",
43
+ }
44
+
45
+
46
+ def _wip_commit(root: Path, parent_sha: str) -> str:
47
+ """Create a dangling commit capturing the working tree so a `--dirty` submit
48
+ runs exactly what's on disk: tracked modifications, staged changes,
49
+ deletions, and untracked non-ignored files. The user's index, HEAD, and
50
+ working tree are left untouched — a scratch index is used — and gitignored
51
+ files (e.g. `.env`) stay out, matching a normal commit. Its parent is
52
+ `parent_sha` (HEAD); returns the new commit's sha."""
53
+ with tempfile.TemporaryDirectory() as td:
54
+ env = {**os.environ, **_WIP_IDENTITY, "GIT_INDEX_FILE": str(Path(td) / "index")}
55
+
56
+ def run(*args: str) -> subprocess.CompletedProcess[str]:
57
+ return subprocess.run(
58
+ ["git", "-C", str(root), *args],
59
+ capture_output=True,
60
+ text=True,
61
+ env=env,
62
+ timeout=600,
63
+ )
64
+
65
+ # Seed the scratch index from HEAD, then stage every working-tree change.
66
+ for step in (("read-tree", parent_sha), ("add", "-A")):
67
+ r = run(*step)
68
+ if r.returncode != 0:
69
+ raise RepoError(
70
+ f"capturing --dirty working tree failed at `git {step[0]}`:\n"
71
+ f"{r.stderr.strip()}"
72
+ )
73
+ tree = run("write-tree")
74
+ if tree.returncode != 0:
75
+ raise RepoError(
76
+ "capturing --dirty working tree failed at `git write-tree`:\n"
77
+ f"{tree.stderr.strip()}"
78
+ )
79
+ commit = run(
80
+ "commit-tree", tree.stdout.strip(), "-p", parent_sha, "-m", "omnirun wip"
81
+ )
82
+ if commit.returncode != 0:
83
+ raise RepoError(
84
+ "capturing --dirty working tree failed at `git commit-tree`:\n"
85
+ f"{commit.stderr.strip()}"
86
+ )
87
+ return commit.stdout.strip()
88
+
89
+
35
90
  def find_repo_root(start: Path | None = None) -> Path:
36
91
  """Toplevel of the git repo containing `start` (default: cwd)."""
37
92
  start = start or Path.cwd()
@@ -111,6 +166,13 @@ def capture_repo_state(
111
166
  f"run `git push origin {branch}` first or pass --push"
112
167
  )
113
168
 
169
+ # A --dirty submit must run the working tree, not plain HEAD: snapshot it
170
+ # into a wip commit and ship that sha (pushed to refs/omnirun/<sha12> for the
171
+ # SSH family, bundled for notebooks). `dirty` is only true here when
172
+ # allow_dirty was set, else we raised above.
173
+ if dirty:
174
+ sha = _wip_commit(root, sha)
175
+
114
176
  return RepoRef(
115
177
  remote_url=remote_url,
116
178
  sha=sha,
@@ -11,7 +11,12 @@ from typing import ClassVar
11
11
  import pytest
12
12
  from typer.testing import CliRunner
13
13
 
14
- from omnirun.backends.base import Backend, BackendError, register # noqa: E402
14
+ from omnirun.backends.base import ( # noqa: E402
15
+ Backend,
16
+ BackendError,
17
+ ProvisioningSink,
18
+ register,
19
+ )
15
20
  from omnirun.cli import app # noqa: E402
16
21
  from omnirun.models import ( # noqa: E402
17
22
  JobHandle,
@@ -47,7 +52,12 @@ class StubBackend(Backend):
47
52
  )
48
53
  ]
49
54
 
50
- def submit(self, spec: JobSpec, offer: Offer) -> JobHandle:
55
+ def submit(
56
+ self,
57
+ spec: JobSpec,
58
+ offer: Offer,
59
+ on_provisioning: ProvisioningSink | None = None,
60
+ ) -> JobHandle:
51
61
  type(self).submitted[spec.job_id] = (spec, offer)
52
62
  return JobHandle(
53
63
  backend=self.name, job_id=spec.job_id, data={"token": f"t-{spec.job_id}"}
@@ -96,7 +106,12 @@ class UnreachableBackend(Backend):
96
106
  )
97
107
  ]
98
108
 
99
- def submit(self, spec: JobSpec, offer: Offer) -> JobHandle: # pragma: no cover
109
+ def submit(
110
+ self,
111
+ spec: JobSpec,
112
+ offer: Offer,
113
+ on_provisioning: ProvisioningSink | None = None,
114
+ ) -> JobHandle: # pragma: no cover
100
115
  raise AssertionError("dry-run must never submit")
101
116
 
102
117
  def status(self, handle: JobHandle) -> StatusReport: # pragma: no cover
@@ -116,6 +131,47 @@ class UnreachableBackend(Backend):
116
131
  return []
117
132
 
118
133
 
134
+ @register("provfail")
135
+ class ProvisionThenFailBackend(Backend):
136
+ """Rents a resource (emits a provisioning stub) then dies before returning a
137
+ handle — models an interrupted marketplace submit (issue #7)."""
138
+
139
+ def probe(self, res: ResourceSpec) -> list[Offer]:
140
+ return [Offer(backend=self.name, label=f"{self.name}: ok", fits=True)]
141
+
142
+ def submit(
143
+ self,
144
+ spec: JobSpec,
145
+ offer: Offer,
146
+ on_provisioning: ProvisioningSink | None = None,
147
+ ) -> JobHandle:
148
+ if on_provisioning is not None:
149
+ on_provisioning(
150
+ JobHandle(
151
+ backend=self.name,
152
+ job_id=spec.job_id,
153
+ data={"instance_id": "inst-42", "provisioning": True},
154
+ )
155
+ )
156
+ raise BackendError("killed mid-provision")
157
+
158
+ def status(self, handle: JobHandle) -> StatusReport: # pragma: no cover
159
+ raise BackendError("n/a")
160
+
161
+ def logs(
162
+ self, handle: JobHandle, follow: bool = False
163
+ ) -> Iterator[str]: # pragma: no cover
164
+ yield ""
165
+
166
+ def cancel(self, handle: JobHandle) -> None: # pragma: no cover
167
+ pass
168
+
169
+ def pull_outputs(
170
+ self, handle: JobHandle, dest: Path
171
+ ) -> list[Path]: # pragma: no cover
172
+ return []
173
+
174
+
119
175
  BASE_CONFIG = """\
120
176
  [policy]
121
177
  auto_wait_threshold = "15m"
@@ -126,6 +182,9 @@ type = "stub"
126
182
 
127
183
  [backends.offline]
128
184
  type = "unreachable"
185
+
186
+ [backends.provfail]
187
+ type = "provfail"
129
188
  """
130
189
 
131
190
 
@@ -172,6 +231,21 @@ def submit_one(*extra: str) -> str:
172
231
  # ------------------------------------------------------------------ submit
173
232
 
174
233
 
234
+ def test_submit_interrupted_after_provision_leaves_reclaimable_stub(env):
235
+ """A submit that dies after renting still persists a stub record carrying the
236
+ instance id, so `ps`/`gc` can see and reclaim the orphan (issue #7)."""
237
+ result = runner.invoke(
238
+ app, ["submit", "--yes", "--backend", "provfail", "--", "python", "train.py"]
239
+ )
240
+ assert result.exit_code != 0 # the submit failed...
241
+
242
+ ids = JobStore().list_ids()
243
+ assert len(ids) == 1 # ...but a record survived
244
+ rec = JobStore().load(ids[0])
245
+ assert rec is not None and rec.handle is not None
246
+ assert rec.handle.data == {"instance_id": "inst-42", "provisioning": True}
247
+
248
+
175
249
  def test_submit_yes_happy_path(env):
176
250
  result = runner.invoke(
177
251
  app,
@@ -309,6 +309,119 @@ def test_runpod_submit_happy_path(spec, fake_ssh, fake_stage):
309
309
  assert f"sleep {24 * 3600}" in watcher
310
310
 
311
311
 
312
+ @respx.mock
313
+ def test_submit_emits_provisioning_stub_before_wait(spec, fake_ssh, fake_stage):
314
+ """on_provisioning fires with the instance id the instant it's rented —
315
+ before provisioning finishes — so an interrupted submit stays reclaimable."""
316
+ respx.post(f"{REST_BASE}/pods").mock(
317
+ return_value=httpx.Response(200, json={"id": "pod123"})
318
+ )
319
+ respx.get(f"{REST_BASE}/pods/pod123").mock(
320
+ return_value=httpx.Response(
321
+ 200,
322
+ json={
323
+ "desiredStatus": "RUNNING",
324
+ "publicIp": "1.2.3.4",
325
+ "portMappings": {"22": 40022},
326
+ },
327
+ )
328
+ )
329
+ stubs: list[JobHandle] = []
330
+ handle = runpod_backend().submit(spec, h100_offer(), on_provisioning=stubs.append)
331
+
332
+ # Exactly one stub, carrying the instance id and flagged as provisioning,
333
+ # with none of the not-yet-known connection details.
334
+ (stub,) = stubs
335
+ assert stub.backend == "runpod"
336
+ assert stub.job_id == spec.job_id
337
+ assert stub.data == {"instance_id": "pod123", "provisioning": True}
338
+ # The final handle is the full one (superset).
339
+ assert handle.data["instance_id"] == "pod123"
340
+ assert handle.data["job_dir"].endswith(spec.job_id)
341
+
342
+
343
+ @respx.mock
344
+ def test_submit_stub_persists_when_interrupted(spec, fake_ssh, monkeypatch):
345
+ """If the submit dies after renting (here: staging raises), the stub was
346
+ already handed to the client, so the orphaned instance is recorded."""
347
+ respx.post(f"{REST_BASE}/pods").mock(
348
+ return_value=httpx.Response(200, json={"id": "pod123"})
349
+ )
350
+ respx.get(f"{REST_BASE}/pods/pod123").mock(
351
+ return_value=httpx.Response(
352
+ 200,
353
+ json={
354
+ "desiredStatus": "RUNNING",
355
+ "publicIp": "1.2.3.4",
356
+ "portMappings": {"22": 40022},
357
+ },
358
+ )
359
+ )
360
+ respx.delete(f"{REST_BASE}/pods/pod123").mock(
361
+ return_value=httpx.Response(200, json={})
362
+ )
363
+
364
+ def boom(*a, **kw):
365
+ raise RuntimeError("git push exploded")
366
+
367
+ monkeypatch.setattr(jobdir, "stage_job", boom)
368
+ stubs: list[JobHandle] = []
369
+ with pytest.raises(BackendError, match="git push exploded"):
370
+ runpod_backend().submit(spec, h100_offer(), on_provisioning=stubs.append)
371
+ assert stubs and stubs[0].data["instance_id"] == "pod123"
372
+
373
+
374
+ @respx.mock
375
+ def test_status_on_provisioning_stub_flags_billing_instance(fake_ssh):
376
+ """status() of a stub whose instance is still up reports PROVISIONING and
377
+ tells the user it is billing (so ps surfaces the orphan)."""
378
+ respx.get(f"{REST_BASE}/pods/pod123").mock(
379
+ return_value=httpx.Response(200, json={"desiredStatus": "RUNNING"})
380
+ )
381
+ stub = JobHandle(
382
+ backend="runpod",
383
+ job_id="train-abc123",
384
+ data={"instance_id": "pod123", "provisioning": True},
385
+ )
386
+ report = runpod_backend().status(stub)
387
+ assert report.status is JobStatus.PROVISIONING
388
+ assert "still billing" in report.detail
389
+ assert "pod123" in report.detail
390
+
391
+
392
+ @respx.mock
393
+ def test_status_on_provisioning_stub_lost_when_instance_gone(fake_ssh):
394
+ """A stub whose instance no longer exists is LOST, not a crash."""
395
+ respx.get(f"{REST_BASE}/pods/pod123").mock(
396
+ return_value=httpx.Response(404, json={"error": "not found"})
397
+ )
398
+ stub = JobHandle(
399
+ backend="runpod",
400
+ job_id="train-abc123",
401
+ data={"instance_id": "pod123", "provisioning": True},
402
+ )
403
+ report = runpod_backend().status(stub)
404
+ assert report.status is JobStatus.LOST
405
+
406
+
407
+ @respx.mock
408
+ def test_gc_reclaims_provisioning_stub(fake_ssh):
409
+ """gc terminates the instance recorded in an interrupted-submit stub."""
410
+ respx.get(f"{REST_BASE}/pods/pod123").mock(
411
+ return_value=httpx.Response(200, json={"desiredStatus": "RUNNING"})
412
+ )
413
+ delete = respx.delete(f"{REST_BASE}/pods/pod123").mock(
414
+ return_value=httpx.Response(200, json={})
415
+ )
416
+ stub = JobHandle(
417
+ backend="runpod",
418
+ job_id="train-abc123",
419
+ data={"instance_id": "pod123", "provisioning": True},
420
+ )
421
+ runpod_backend().gc(stub)
422
+ assert delete.called
423
+
424
+
312
425
  @respx.mock
313
426
  def test_runpod_submit_provision_timeout_terminates(spec, fake_ssh, fake_stage):
314
427
  respx.post(f"{REST_BASE}/pods").mock(
@@ -7,7 +7,7 @@ import time
7
7
  from collections.abc import Iterator
8
8
  from pathlib import Path
9
9
 
10
- from omnirun.backends.base import Backend
10
+ from omnirun.backends.base import Backend, ProvisioningSink
11
11
  from omnirun.config import BackendConfig, Config, DaemonConfig
12
12
  from omnirun.daemon import Daemon, daemon_address, send_request
13
13
  from omnirun.models import (
@@ -68,7 +68,12 @@ class FakeBackend(Backend):
68
68
  )
69
69
  ]
70
70
 
71
- def submit(self, spec: JobSpec, offer: Offer) -> JobHandle:
71
+ def submit(
72
+ self,
73
+ spec: JobSpec,
74
+ offer: Offer,
75
+ on_provisioning: ProvisioningSink | None = None,
76
+ ) -> JobHandle:
72
77
  if self.fail_submit:
73
78
  raise RuntimeError("submit boom")
74
79
  self.submitted.append(spec.job_id)
@@ -80,10 +80,61 @@ def test_capture_dirty_untracked(sample_repo: Path) -> None:
80
80
  capture_repo_state(sample_repo)
81
81
 
82
82
 
83
- def test_capture_allow_dirty(sample_repo: Path) -> None:
83
+ def test_capture_allow_dirty_snapshots_working_tree(sample_repo: Path) -> None:
84
+ """--dirty ships a wip commit whose tree is the working tree, not plain HEAD,
85
+ and leaves the user's HEAD/index/working tree untouched (issue #6)."""
86
+ head = git(sample_repo, "rev-parse", "HEAD")
84
87
  (sample_repo / "job.py").write_text("changed\n")
88
+ (sample_repo / "scratch.txt").write_text("wip\n") # untracked, non-ignored
89
+
85
90
  ref = capture_repo_state(sample_repo, allow_dirty=True)
91
+
86
92
  assert ref.dirty is True
93
+ # A real, distinct commit parented on HEAD — not HEAD itself.
94
+ assert ref.sha != head
95
+ assert git(sample_repo, "rev-parse", f"{ref.sha}^") == head
96
+ # The wip tree carries the uncommitted edits and the untracked file.
97
+ assert git(sample_repo, "show", f"{ref.sha}:job.py") == "changed"
98
+ assert git(sample_repo, "show", f"{ref.sha}:scratch.txt") == "wip"
99
+ # The user's repo is not mutated: HEAD unchanged, tree still dirty on disk.
100
+ assert git(sample_repo, "rev-parse", "HEAD") == head
101
+ assert (sample_repo / "job.py").read_text() == "changed\n"
102
+ assert git(sample_repo, "status", "--porcelain") != ""
103
+
104
+
105
+ def test_capture_allow_dirty_excludes_gitignored(sample_repo: Path) -> None:
106
+ """A gitignored file (e.g. .env) rides out-of-band, never in the wip tree."""
107
+ (sample_repo / ".gitignore").write_text(".env\n")
108
+ git(sample_repo, "add", ".gitignore")
109
+ git(sample_repo, "commit", "-q", "-m", "add gitignore")
110
+ (sample_repo / ".env").write_text("SECRET=1\n")
111
+ (sample_repo / "job.py").write_text("changed\n")
112
+
113
+ ref = capture_repo_state(sample_repo, allow_dirty=True)
114
+
115
+ assert git(sample_repo, "show", f"{ref.sha}:job.py") == "changed"
116
+ missing = subprocess.run(
117
+ ["git", "-C", str(sample_repo), "cat-file", "-e", f"{ref.sha}:.env"],
118
+ capture_output=True,
119
+ )
120
+ assert missing.returncode != 0 # .env is absent from the wip tree
121
+
122
+
123
+ def test_capture_allow_dirty_captures_deletion(sample_repo: Path) -> None:
124
+ """A deleted tracked file is absent from the wip tree."""
125
+ (sample_repo / "job.py").unlink()
126
+ ref = capture_repo_state(sample_repo, allow_dirty=True)
127
+ gone = subprocess.run(
128
+ ["git", "-C", str(sample_repo), "cat-file", "-e", f"{ref.sha}:job.py"],
129
+ capture_output=True,
130
+ )
131
+ assert gone.returncode != 0
132
+
133
+
134
+ def test_capture_allow_dirty_clean_tree_is_head(sample_repo: Path) -> None:
135
+ """--dirty on a clean tree is a no-op: sha stays HEAD, no wip commit."""
136
+ ref = capture_repo_state(sample_repo, allow_dirty=True)
137
+ assert ref.dirty is False
87
138
  assert ref.sha == git(sample_repo, "rev-parse", "HEAD")
88
139
 
89
140
 
@@ -595,7 +595,7 @@ wheels = [
595
595
 
596
596
  [[package]]
597
597
  name = "omnirun"
598
- version = "0.2.0"
598
+ version = "0.2.3"
599
599
  source = { editable = "." }
600
600
  dependencies = [
601
601
  { name = "httpx" },
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