omnirun 0.2.0__tar.gz → 0.2.2__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.0 → omnirun-0.2.2}/.github/workflows/publish.yml +3 -0
  2. {omnirun-0.2.0 → omnirun-0.2.2}/DESIGN.md +7 -4
  3. {omnirun-0.2.0 → omnirun-0.2.2}/PKG-INFO +5 -2
  4. {omnirun-0.2.0 → omnirun-0.2.2}/README.md +4 -1
  5. {omnirun-0.2.0 → omnirun-0.2.2}/TESTING.md +23 -10
  6. {omnirun-0.2.0 → omnirun-0.2.2}/flake.nix +3 -0
  7. {omnirun-0.2.0 → omnirun-0.2.2}/pyproject.toml +1 -1
  8. omnirun-0.2.2/scripts/check-version.sh +32 -0
  9. omnirun-0.2.2/scripts/pre-push +39 -0
  10. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/__init__.py +1 -1
  11. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/backends/colab.py +5 -3
  12. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/backends/kaggle.py +41 -17
  13. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/backends/local.py +4 -1
  14. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/backends/marketplace.py +1 -1
  15. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/backends/slurm.py +4 -2
  16. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/backends/ssh.py +1 -1
  17. omnirun-0.2.2/src/omnirun/backends/tarsafe.py +38 -0
  18. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/backends/vast.py +15 -9
  19. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/bootstrap.py +28 -2
  20. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/config.py +13 -2
  21. {omnirun-0.2.0 → omnirun-0.2.2}/tests/test_bootstrap.py +32 -0
  22. omnirun-0.2.2/tests/test_config.py +37 -0
  23. {omnirun-0.2.0 → omnirun-0.2.2}/tests/test_kaggle.py +43 -0
  24. {omnirun-0.2.0 → omnirun-0.2.2}/tests/test_marketplaces.py +1 -1
  25. {omnirun-0.2.0 → omnirun-0.2.2}/tests/test_slurm.py +26 -0
  26. omnirun-0.2.2/tests/test_tarsafe.py +53 -0
  27. {omnirun-0.2.0 → omnirun-0.2.2}/.envrc +0 -0
  28. {omnirun-0.2.0 → omnirun-0.2.2}/.github/workflows/checks.yml +0 -0
  29. {omnirun-0.2.0 → omnirun-0.2.2}/.github/workflows/ci.yml +0 -0
  30. {omnirun-0.2.0 → omnirun-0.2.2}/.gitignore +0 -0
  31. {omnirun-0.2.0 → omnirun-0.2.2}/.python-version +0 -0
  32. {omnirun-0.2.0 → omnirun-0.2.2}/AGENTS.md +0 -0
  33. {omnirun-0.2.0 → omnirun-0.2.2}/CLAUDE.md +0 -0
  34. {omnirun-0.2.0 → omnirun-0.2.2}/flake.lock +0 -0
  35. {omnirun-0.2.0 → omnirun-0.2.2}/main.py +0 -0
  36. {omnirun-0.2.0 → omnirun-0.2.2}/research/colab-kaggle.md +0 -0
  37. {omnirun-0.2.0 → omnirun-0.2.2}/research/gpu-marketplaces.md +0 -0
  38. {omnirun-0.2.0 → omnirun-0.2.2}/research/launcher-landscape.md +0 -0
  39. {omnirun-0.2.0 → omnirun-0.2.2}/research/ml-job-queues.md +0 -0
  40. {omnirun-0.2.0 → omnirun-0.2.2}/research/slurm-ssh.md +0 -0
  41. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/backends/__init__.py +0 -0
  42. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/backends/base.py +0 -0
  43. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/backends/jobdir.py +0 -0
  44. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/backends/runpod.py +0 -0
  45. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/backends/thunder.py +0 -0
  46. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/chooser.py +0 -0
  47. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/cli.py +0 -0
  48. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/daemon.py +0 -0
  49. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/execlayer/__init__.py +0 -0
  50. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/execlayer/base.py +0 -0
  51. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/execlayer/local.py +0 -0
  52. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/execlayer/ssh.py +0 -0
  53. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/models.py +0 -0
  54. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/queue.py +0 -0
  55. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/repo.py +0 -0
  56. {omnirun-0.2.0 → omnirun-0.2.2}/src/omnirun/store.py +0 -0
  57. {omnirun-0.2.0 → omnirun-0.2.2}/tests/__init__.py +0 -0
  58. {omnirun-0.2.0 → omnirun-0.2.2}/tests/conftest.py +0 -0
  59. {omnirun-0.2.0 → omnirun-0.2.2}/tests/test_chooser.py +0 -0
  60. {omnirun-0.2.0 → omnirun-0.2.2}/tests/test_cli.py +0 -0
  61. {omnirun-0.2.0 → omnirun-0.2.2}/tests/test_colab.py +0 -0
  62. {omnirun-0.2.0 → omnirun-0.2.2}/tests/test_local_backend.py +0 -0
  63. {omnirun-0.2.0 → omnirun-0.2.2}/tests/test_queue.py +0 -0
  64. {omnirun-0.2.0 → omnirun-0.2.2}/tests/test_repo.py +0 -0
  65. {omnirun-0.2.0 → omnirun-0.2.2}/tests/test_ssh_backend.py +0 -0
  66. {omnirun-0.2.0 → omnirun-0.2.2}/tests/test_ssh_exec.py +0 -0
  67. {omnirun-0.2.0 → omnirun-0.2.2}/tests/test_store.py +0 -0
  68. {omnirun-0.2.0 → omnirun-0.2.2}/uv.lock +0 -0
@@ -22,6 +22,9 @@ jobs:
22
22
  - name: Checkout
23
23
  uses: actions/checkout@v7
24
24
 
25
+ - name: Check tag matches package version
26
+ run: ./scripts/check-version.sh "${GITHUB_REF_NAME}"
27
+
25
28
  - name: Install uv
26
29
  uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0
27
30
  with:
@@ -282,10 +282,13 @@ Worker access to private repos — **no git credentials ever leave the laptop**:
282
282
  bundle base64 inside the kernel's `run.py`** (next to the base64 bootstrap) — *not* a
283
283
  dataset. The old dataset approach 409'd every kernel that referenced a still-processing
284
284
  dataset (a systematic race) and needed a create/delete lifecycle; embedding removes
285
- both. A size guard (`MAX_EMBED_B64`, 40 MiB, covering bundle + env together) rejects
286
- oversized repos, so this suits code-sized repos (data is never shipped jobs fetch
287
- their own). The bootstrap clones/fetches from the bundle into the object store. In the
288
- public-repo case no bundle is embedded, so the ceiling is a non-issue.
285
+ both. A size guard (`KAGGLE_MAX_SOURCE_BYTES`, ~1 MiB, measuring the whole `run.py`
286
+ kernel source embedded bootstrap + bundle + env against Kaggle's real push limit,
287
+ overridable per backend via `max_source_bytes`) rejects oversized pushes early with a
288
+ size-naming error instead of an opaque Kaggle-side failure, so this suits code-sized
289
+ repos (data is never shipped — jobs fetch their own). The bootstrap clones/fetches from
290
+ the bundle into the object store. In the public-repo case no bundle is embedded, so the
291
+ ceiling is a non-issue.
289
292
 
290
293
  This keeps the invariant intact: a public clone needs no credentials, and private repos
291
294
  still never touch origin from the worker. What changed is that public-repo workers now
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: omnirun
3
- Version: 0.2.0
3
+ Version: 0.2.2
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
@@ -140,7 +140,9 @@ section name is yours (you can have several of the same type). All backends
140
140
  accept `enabled = false` to keep the config without probing it. Every backend
141
141
  also accepts `project_root` (point it at an existing checkout to reuse its
142
142
  `.git`/`.venv`) and `max_parallel` (queue-daemon cap on concurrent jobs, default
143
- `1`).
143
+ `1`). `project_root` may be a single path (applied to every repo) or a table
144
+ mapping repo slug -> path, with an optional `default` key as fallback, so one
145
+ backend can serve several repos (each getting the right checkout).
144
146
 
145
147
  ```toml
146
148
  [policy]
@@ -159,6 +161,7 @@ poll_interval_s = 10.0 # scheduler tick: refresh running + place pending
159
161
  type = "local"
160
162
  # root = "$HOME/.omnirun" # OMNIRUN_ROOT: where jobs live
161
163
  # project_root = "$HOME/proj" # reuse an existing checkout's .git and .venv
164
+ # project_root = { my-repo = "$HOME/proj" } # or per-repo (by slug), else default path
162
165
  # max_parallel = 2 # queue daemon: run up to 2 jobs here at once
163
166
 
164
167
  # ---- any machine you can ssh into ----
@@ -120,7 +120,9 @@ section name is yours (you can have several of the same type). All backends
120
120
  accept `enabled = false` to keep the config without probing it. Every backend
121
121
  also accepts `project_root` (point it at an existing checkout to reuse its
122
122
  `.git`/`.venv`) and `max_parallel` (queue-daemon cap on concurrent jobs, default
123
- `1`).
123
+ `1`). `project_root` may be a single path (applied to every repo) or a table
124
+ mapping repo slug -> path, with an optional `default` key as fallback, so one
125
+ backend can serve several repos (each getting the right checkout).
124
126
 
125
127
  ```toml
126
128
  [policy]
@@ -139,6 +141,7 @@ poll_interval_s = 10.0 # scheduler tick: refresh running + place pending
139
141
  type = "local"
140
142
  # root = "$HOME/.omnirun" # OMNIRUN_ROOT: where jobs live
141
143
  # project_root = "$HOME/proj" # reuse an existing checkout's .git and .venv
144
+ # project_root = { my-repo = "$HOME/proj" } # or per-repo (by slug), else default path
142
145
  # max_parallel = 2 # queue daemon: run up to 2 jobs here at once
143
146
 
144
147
  # ---- any machine you can ssh into ----
@@ -116,15 +116,27 @@ bootstrap script — there is **no per-job dataset** anymore. This replaced a
116
116
  dataset-based design that hit a systematic **409** (Kaggle rejects a kernel
117
117
  referencing a still-processing dataset). Consequence: kaggle no longer creates or
118
118
  deletes any datasets; nothing worker-side to reap. Only caveat is bundle size — a
119
- code-only repo bundle is well under the embed cap; a repo with large committed
120
- blobs will be rejected client-side before push.
121
-
122
- **Newly added (unit-tested, NOT yet live-verified):** for a **public** repo the
123
- kernel now clones the repo directly over its own internet connection and no bundle
124
- is embedded (`repo.remote_clone_plan` gates on public + reachable). Kaggle also now
125
- injects a gitignored `<repo>/.env` (base64-embedded as `ENV_B64`, decoded to a 0600
126
- file and sourced) parity with Colab. Both paths pass unit tests; a live Kaggle
127
- (and Colab) job against a public repo has not been re-run yet.
119
+ code-only repo bundle is well under the source cap; a repo with large committed
120
+ blobs is rejected client-side before push (see the measured cap below).
121
+
122
+ **Public-repo direct clone — LIVE-VERIFIED (this session).** A CPU kernel was
123
+ submitted against `github.com/flyingleafe/omnirun` (public) at the pushed master
124
+ tip with `env=none`: the worker ran `git clone --bare` over its own connection,
125
+ checked out the exact sha (`Preparing worktree (detached HEAD 6fcedf0)`
126
+ `CLONE OK cwd=…/.trees/6fcedf04a860`, `has_pyproject=True`), no bundle embedded,
127
+ and `proof.txt` was pulled back. So a **public repo of any size** runs on Kaggle,
128
+ unconstrained by the source cap. Kaggle also injects a gitignored `<repo>/.env`
129
+ (base64 `ENV_B64` → 0600 file → sourced) — unit-tested, parity with Colab.
130
+
131
+ **Kernel-source cap — MEASURED (this session).** Pushing `run.py` payloads of
132
+ increasing size to the kernels API: `<=1 MiB` (1,048,570 B) **accepted**,
133
+ `>=1.1 MiB` **rejected with HTTP 400** — so Kaggle's limit is **1 MiB**. The
134
+ pre-submit guard (`KAGGLE_MAX_SOURCE_BYTES = 1 MiB`, override via
135
+ `max_source_bytes`) now measures the full `run.py` (bootstrap + any bundle + any
136
+ `.env`) against this and fails early naming size, instead of the old 40 MiB
137
+ `MAX_EMBED_B64` guard that let 1–40 MiB pushes fail opaquely on Kaggle's side.
138
+ (Aside: the current kaggle client exposes `kernels_delete(ref, no_confirm=True)`,
139
+ used here to clean up the probe kernels — a future `gc`/`cancel` could use it.)
128
140
 
129
141
  - Provide: API token at `~/.config/kaggle/kaggle.json` (kaggle.com → Settings →
130
142
  Create New Token) or `KAGGLE_USERNAME`/`KAGGLE_KEY` env vars. The account
@@ -202,7 +214,8 @@ module.
202
214
 
203
215
  **Kaggle** (`src/omnirun/backends/kaggle.py`) — CONFIRMED LIVE (P100 job ran; output pulled)
204
216
  - [x] Bundle delivery (private/unpushed repo): **embedded base64 in `run.py`**, no dataset — sidesteps the 409 the old dataset design hit. `dataset_create_new`/`dataset_delete` are no longer used.
205
- - [ ] Public-repo **direct clone** (worker `git clone` over its own connection, no bundle) — unit-tested via `repo.remote_clone_plan`; not yet run live on Kaggle.
217
+ - [x] Public-repo **direct clone** (worker `git clone` over its own connection, no bundle) — **LIVE-VERIFIED**: CPU job cloned `flyingleafe/omnirun` at master tip, ran in the worktree, `proof.txt` pulled.
218
+ - [x] Kernel-source push cap — **MEASURED**: 1 MiB (≤1 MiB accepted, ≥1.1 MiB → HTTP 400); `KAGGLE_MAX_SOURCE_BYTES` guard set to match, overridable via `max_source_bytes`.
206
219
  - [ ] `.env` injection (`ENV_B64` embedded → 0600 file → sourced) — unit-tested; not yet run live on Kaggle.
207
220
  - [x] `kernels_status` response shape + status strings (`queued/running/complete/error/cancelAcknowledged`).
208
221
  - [x] `kernels_output` kwarg + downloaded log format (`<slug>.log`).
@@ -71,6 +71,9 @@
71
71
 
72
72
  shellHook = ''
73
73
  ${pre-commit-check.shellHook}
74
+ # Guard against tagging a release without bumping the version.
75
+ # pre-commit owns only the pre-commit hook, so pre-push is ours.
76
+ ln -sf "$PWD/scripts/pre-push" "$(git rev-parse --git-path hooks)/pre-push"
74
77
  export LD_LIBRARY_PATH=${pkgs.lib.makeLibraryPath [ pkgs.stdenv.cc.cc.lib pkgs.zlib ]}:$LD_LIBRARY_PATH
75
78
  export UV_PYTHON=${python}/bin/python
76
79
  uv sync --quiet
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "omnirun"
3
- version = "0.2.0"
3
+ version = "0.2.2"
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"
@@ -0,0 +1,32 @@
1
+ #!/usr/bin/env bash
2
+ # Verify the package version is internally consistent, and optionally that it
3
+ # matches an expected version (e.g. a release tag).
4
+ #
5
+ # scripts/check-version.sh # pyproject.toml == src/omnirun/__init__.py
6
+ # scripts/check-version.sh v0.2.2 # ...and both equal 0.2.2
7
+ #
8
+ # Used by both the publish workflow and the pre-push hook so the rule lives in
9
+ # one place.
10
+ set -euo pipefail
11
+
12
+ root="$(cd "$(dirname "$0")/.." && pwd)"
13
+
14
+ extract() { grep -m1 -E "$2" "$root/$1" | sed -E 's/.*"([^"]+)".*/\1/'; }
15
+
16
+ pyproject="$(extract pyproject.toml '^version *= *"')"
17
+ init="$(extract src/omnirun/__init__.py '^__version__ *= *"')"
18
+
19
+ if [ "$pyproject" != "$init" ]; then
20
+ echo "version mismatch: pyproject.toml=$pyproject src/omnirun/__init__.py=$init" >&2
21
+ exit 1
22
+ fi
23
+
24
+ if [ "$#" -ge 1 ]; then
25
+ expected="${1#v}"
26
+ if [ "$pyproject" != "$expected" ]; then
27
+ echo "version mismatch: expected $expected but package version is $pyproject" >&2
28
+ exit 1
29
+ fi
30
+ fi
31
+
32
+ echo "version OK: $pyproject"
@@ -0,0 +1,39 @@
1
+ #!/usr/bin/env bash
2
+ # Refuse to push a release tag whose commit does not declare a matching version.
3
+ #
4
+ # Installed as .git/hooks/pre-push by the devShell (see flake.nix). It inspects
5
+ # the *pushed commit's* content (via `git show <sha>:path`), so it catches a
6
+ # stale worktree too. CI (publish.yml) enforces the same rule server-side.
7
+ set -euo pipefail
8
+
9
+ version_at() { git show "$1:$2" | grep -m1 -E "$3" | sed -E 's/.*"([^"]+)".*/\1/'; }
10
+
11
+ status=0
12
+ while read -r _local_ref local_sha remote_ref _remote_sha; do
13
+ case "$remote_ref" in
14
+ refs/tags/v[0-9]*) ;;
15
+ *) continue ;;
16
+ esac
17
+ # Deleting a tag pushes the zero sha; nothing to check.
18
+ case "$local_sha" in
19
+ 0000000000000000000000000000000000000000) continue ;;
20
+ esac
21
+
22
+ tag="${remote_ref#refs/tags/}"
23
+ pyproject="$(version_at "$local_sha" pyproject.toml '^version *= *"')"
24
+ init="$(version_at "$local_sha" src/omnirun/__init__.py '^__version__ *= *"')"
25
+
26
+ if [ "$tag" != "v$pyproject" ]; then
27
+ echo "pre-push: tag $tag != pyproject version $pyproject (commit ${local_sha:0:12})" >&2
28
+ status=1
29
+ fi
30
+ if [ "$pyproject" != "$init" ]; then
31
+ echo "pre-push: pyproject $pyproject != __init__ $init (commit ${local_sha:0:12})" >&2
32
+ status=1
33
+ fi
34
+ done
35
+
36
+ if [ "$status" -ne 0 ]; then
37
+ echo "pre-push: bump the version before tagging (see scripts/check-version.sh)" >&2
38
+ fi
39
+ exit "$status"
@@ -1,3 +1,3 @@
1
1
  """omnirun — run jobs from your repo anywhere."""
2
2
 
3
- __version__ = "0.2.0"
3
+ __version__ = "0.2.2"
@@ -32,7 +32,7 @@ from datetime import datetime, timedelta, timezone
32
32
  from pathlib import Path
33
33
  from typing import Any
34
34
 
35
- from omnirun.backends import jobdir
35
+ from omnirun.backends import jobdir, tarsafe
36
36
  from omnirun.backends.base import Backend, BackendError, register
37
37
  from omnirun.bootstrap import (
38
38
  HEARTBEAT_STALE_S,
@@ -360,7 +360,9 @@ class ColabBackend(Backend):
360
360
  BootstrapParams(
361
361
  omnirun_root=COLAB_ROOT,
362
362
  project_root=jobdir.project_root_of(
363
- COLAB_ROOT, spec.repo.slug, self.config.project_root
363
+ COLAB_ROOT,
364
+ spec.repo.slug,
365
+ self.config.project_root_for(spec.repo.slug),
364
366
  ),
365
367
  code=code,
366
368
  ),
@@ -583,7 +585,7 @@ class ColabBackend(Backend):
583
585
  if not local_tar.exists():
584
586
  raise BackendError("colab download produced no local file")
585
587
  with tarfile.open(local_tar) as tf:
586
- tf.extractall(td, filter="data")
588
+ tarsafe.extract_all(tf, Path(td))
587
589
  outputs = Path(td) / "outputs"
588
590
  if outputs.is_dir():
589
591
  shutil.copytree(outputs, dest, dirs_exist_ok=True)
@@ -34,7 +34,7 @@ from pathlib import Path
34
34
  from typing import Any
35
35
 
36
36
  from omnirun.backends.base import Backend, BackendError, register
37
- from omnirun.backends import jobdir
37
+ from omnirun.backends import jobdir, tarsafe
38
38
  from omnirun.bootstrap import (
39
39
  BootstrapParams,
40
40
  CodeSource,
@@ -76,10 +76,15 @@ MAX_MEM_GB = 30
76
76
  MAX_DISK_GB = 55
77
77
  KAGGLE_ROOT = "/kaggle/tmp/omnirun" # venv + tree scratch; /kaggle/working stays small
78
78
  RESULT_TAR = "omnirun-job.tar.gz"
79
- # Cap on the base64 git bundle embedded in the kernel source. Kaggle rejects
80
- # oversized kernel pushes; fail early with a clear message instead. Generous
81
- # a code-only repo bundle is normally well under a megabyte.
82
- MAX_EMBED_B64 = 40 * 1024 * 1024
79
+ # Cap on the kernel source (run.py) omnirun pushes: the embedded bootstrap +
80
+ # any git bundle (private/unpushed repos) + any out-of-band .env all count
81
+ # against it. Kaggle rejects an oversized push with an opaque HTTP 400, so we
82
+ # guard at the real threshold and fail early naming size as the cause. Measured
83
+ # live against the kernels API: <=1 MiB accepted, >=1.1 MiB rejected — i.e. the
84
+ # limit is 1 MiB. A private/unpushed repo therefore only fits when its bundle is
85
+ # small; a *public* repo is cloned by the worker directly (no bundle shipped) and
86
+ # is unaffected by this cap. Override per backend with `max_source_bytes`.
87
+ KAGGLE_MAX_SOURCE_BYTES = 1024 * 1024
83
88
  LOG_POLL_INTERVAL_S = 30.0 # poll etiquette: >=30s against the kaggle API
84
89
 
85
90
 
@@ -365,6 +370,29 @@ class KaggleBackend(Backend):
365
370
 
366
371
  # ---- submit ---------------------------------------------------------------
367
372
 
373
+ def _guard_source_size(self, harness: str, *, shipped_bundle: bool) -> None:
374
+ """Fail early, with a size-naming message, when the kernel source is over
375
+ Kaggle's real push limit — instead of letting the push fail opaquely."""
376
+ limit = int(self.config.extra("max_source_bytes", KAGGLE_MAX_SOURCE_BYTES))
377
+ size = len(harness.encode())
378
+ if size <= limit:
379
+ return
380
+ # the git bundle (full history + tree at the sha) usually dominates when
381
+ # a private/unpushed repo forces one; a public repo ships no bundle.
382
+ remedy = (
383
+ "push the repo to a public remote so the worker clones it directly "
384
+ "(no bundle shipped), slim the git history (shallow/squash) or exclude "
385
+ "heavy dirs from the tree, "
386
+ if shipped_bundle
387
+ else "slim the tree or "
388
+ )
389
+ raise BackendError(
390
+ f"kaggle kernel source is {size / (1024 * 1024):.1f} MB, over the "
391
+ f"~{limit / (1024 * 1024):.1f} MB Kaggle push limit (the push would "
392
+ f"fail opaquely) — {remedy}run this job on another backend, or raise "
393
+ "the backend's `max_source_bytes` if Kaggle accepts more"
394
+ )
395
+
368
396
  def submit(self, spec: JobSpec, offer: Offer) -> JobHandle:
369
397
  api = self._api()
370
398
  user = self._username(api)
@@ -401,31 +429,27 @@ class KaggleBackend(Backend):
401
429
  env_b64 = (
402
430
  base64.b64encode(envf.read_text().encode()).decode() if envf else ""
403
431
  )
404
- if len(bundle_b64) + len(env_b64) > MAX_EMBED_B64:
405
- raise BackendError(
406
- f"embedded payload is too large for a kaggle kernel "
407
- f"(~{(len(bundle_b64) + len(env_b64)) // (1024 * 1024)}MB base64) "
408
- "— slim the git history (shallow/squash), push a public repo, or "
409
- "run this job on another backend"
410
- )
411
432
 
412
433
  script = generate_bootstrap(
413
434
  notebook_env_spec(spec),
414
435
  BootstrapParams(
415
436
  omnirun_root=KAGGLE_ROOT,
416
437
  project_root=jobdir.project_root_of(
417
- KAGGLE_ROOT, spec.repo.slug, self.config.project_root
438
+ KAGGLE_ROOT,
439
+ spec.repo.slug,
440
+ self.config.project_root_for(spec.repo.slug),
418
441
  ),
419
442
  code=code,
420
443
  ),
421
444
  )
422
445
  boot_b64 = base64.b64encode(script.encode()).decode()
423
446
 
447
+ harness = _render_harness(job_id, boot_b64, bundle_b64, env_b64)
448
+ self._guard_source_size(harness, shipped_bundle=bool(bundle_b64))
449
+
424
450
  k_dir = stage / "kernel"
425
451
  k_dir.mkdir()
426
- (k_dir / "run.py").write_text(
427
- _render_harness(job_id, boot_b64, bundle_b64, env_b64)
428
- )
452
+ (k_dir / "run.py").write_text(harness)
429
453
  meta: dict[str, Any] = {
430
454
  "id": kernel_ref,
431
455
  "title": slug, # title must slugify to the id slug
@@ -610,7 +634,7 @@ class KaggleBackend(Backend):
610
634
  extract = Path(td) / "extract"
611
635
  extract.mkdir()
612
636
  with tarfile.open(tar) as tf:
613
- tf.extractall(extract, filter="data")
637
+ tarsafe.extract_all(tf, extract)
614
638
  outputs = extract / "outputs"
615
639
  if outputs.is_dir():
616
640
  shutil.copytree(outputs, dest, dirs_exist_ok=True)
@@ -113,7 +113,10 @@ class LocalBackend(Backend):
113
113
  def submit(self, spec: JobSpec, offer: Offer) -> JobHandle:
114
114
  root = jobdir.remote_root(self.exec, self.config.root)
115
115
  project_root = jobdir.resolve_project_root(
116
- self.exec, root, spec.repo.slug, self.config.project_root
116
+ self.exec,
117
+ root,
118
+ spec.repo.slug,
119
+ self.config.project_root_for(spec.repo.slug),
117
120
  )
118
121
  params = BootstrapParams(
119
122
  omnirun_root=root,
@@ -286,7 +286,7 @@ class MarketplaceBackend(Backend, ABC):
286
286
  self._wait_ssh(ex)
287
287
  root = jobdir.remote_root(ex, self.config.root)
288
288
  project_root = jobdir.resolve_project_root(
289
- ex, root, spec.repo.slug, self.config.project_root
289
+ ex, root, spec.repo.slug, self.config.project_root_for(spec.repo.slug)
290
290
  )
291
291
  params = BootstrapParams(
292
292
  omnirun_root=root,
@@ -181,6 +181,8 @@ class SlurmBackend(Backend):
181
181
 
182
182
  reasons: list[str] = []
183
183
  notes: list[str] = []
184
+ if res.wants_gpu() and self.config.extra("has_gpus", True) is False:
185
+ reasons.append(f"backend {self.name} is CPU-only (has_gpus = false)")
184
186
  if res.wants_gpu() and res.gpu_type and self.config.gpu_map:
185
187
  if _gpu_map_lookup(self.config.gpu_map, res.gpu_type) is None:
186
188
  reasons.append(
@@ -258,7 +260,7 @@ class SlurmBackend(Backend):
258
260
  root = self.config.root
259
261
  job_dir = jobdir.job_dir_of(root, spec.job_id)
260
262
  project_root = jobdir.project_root_of(
261
- root, spec.repo.slug, self.config.project_root
263
+ root, spec.repo.slug, self.config.project_root_for(spec.repo.slug)
262
264
  )
263
265
  sbatch = render_sbatch(spec, self.config, job_dir, root)
264
266
  bootstrap = generate_bootstrap(
@@ -281,7 +283,7 @@ class SlurmBackend(Backend):
281
283
  ex = self.exec_
282
284
  root = jobdir.remote_root(ex, self.config.root)
283
285
  project_root = jobdir.resolve_project_root(
284
- ex, root, spec.repo.slug, self.config.project_root
286
+ ex, root, spec.repo.slug, self.config.project_root_for(spec.repo.slug)
285
287
  )
286
288
  params = BootstrapParams(
287
289
  omnirun_root=root,
@@ -217,7 +217,7 @@ class SshBackend(Backend):
217
217
  ex = self.exec_
218
218
  root = jobdir.remote_root(ex, self.config.root)
219
219
  project_root = jobdir.resolve_project_root(
220
- ex, root, spec.repo.slug, self.config.project_root
220
+ ex, root, spec.repo.slug, self.config.project_root_for(spec.repo.slug)
221
221
  )
222
222
  params = BootstrapParams(
223
223
  omnirun_root=root,
@@ -0,0 +1,38 @@
1
+ """Tolerant tar extraction for the notebook backends (kaggle, colab), which
2
+ return a job's outputs as a tarball the client must unpack.
3
+
4
+ Python >=3.12 defaults tar extraction to the PEP-706 ``data`` filter. That
5
+ filter is right (it refuses path traversal, device nodes and links pointing
6
+ outside the destination), but ``extractall`` aborts the *whole* archive the
7
+ first time the filter rejects a member. The common offender is a symlink whose
8
+ target is an absolute path — every Weights & Biases run leaves one at
9
+ ``wandb/*/logs/debug-core.log`` — so a single such link would strand every
10
+ output of an otherwise-successful job.
11
+
12
+ ``extract_all`` applies the same ``data`` filter per member but skips (with a
13
+ warning) any member the filter rejects, so the rest of the tree still lands.
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ import sys
19
+ import tarfile
20
+ from pathlib import Path
21
+
22
+
23
+ def extract_all(tf: tarfile.TarFile, dest: Path) -> list[str]:
24
+ """Extract every member of ``tf`` under ``dest`` with the safe ``data``
25
+ filter, skipping (not aborting on) any member the filter rejects. Returns
26
+ the names of the skipped members (also warned to stderr)."""
27
+ skipped: list[str] = []
28
+ for member in tf.getmembers():
29
+ try:
30
+ tf.extract(member, dest, filter="data")
31
+ except tarfile.FilterError as e:
32
+ skipped.append(member.name)
33
+ print(
34
+ f"omnirun: warning: skipped unsafe archive entry "
35
+ f"{member.name!r} ({type(e).__name__})",
36
+ file=sys.stderr,
37
+ )
38
+ return skipped
@@ -37,23 +37,25 @@ OFFER_LIMIT = 5
37
37
  MIN_RELIABILITY = 0.95
38
38
 
39
39
  # normalized name -> vast gpu_name values (their catalog naming)
40
+ # The REST /bundles/ endpoint matches gpu_name with spaces ("A100 SXM4");
41
+ # underscore forms match nothing (the vast CLI translates them, the API doesn't).
40
42
  VAST_GPU_NAMES: dict[str, list[str]] = {
41
43
  "H200": ["H200"],
42
- "H100": ["H100_SXM", "H100_PCIE", "H100_NVL"],
43
- "A100": ["A100_SXM4", "A100_PCIE", "A100X"],
44
- "A100-80": ["A100_SXM4", "A100_PCIE"],
45
- "A6000": ["RTX_A6000"],
44
+ "H100": ["H100 SXM", "H100 PCIE", "H100 NVL"],
45
+ "A100": ["A100 SXM4", "A100 PCIE", "A100X"],
46
+ "A100-80": ["A100 SXM4", "A100 PCIE"],
47
+ "A6000": ["RTX A6000"],
46
48
  "L40": ["L40", "L40S"],
47
49
  "L4": ["L4"],
48
- "V100": ["Tesla_V100"],
49
- "4090": ["RTX_4090"],
50
- "3090": ["RTX_3090"],
51
- "5090": ["RTX_5090"],
50
+ "V100": ["Tesla V100"],
51
+ "4090": ["RTX 4090"],
52
+ "3090": ["RTX 3090"],
53
+ "5090": ["RTX 5090"],
52
54
  }
53
55
 
54
56
 
55
57
  def vast_gpu_names(gpu_type: str) -> list[str]:
56
- return VAST_GPU_NAMES.get(gpu_type, [gpu_type.replace(" ", "_").replace("-", "_")])
58
+ return VAST_GPU_NAMES.get(gpu_type, [gpu_type.replace("_", " ").replace("-", " ")])
57
59
 
58
60
 
59
61
  def normalize_vast_gpu(name: str, gpu_ram_mb: float | None = None) -> str:
@@ -84,6 +86,10 @@ class VastBackend(MarketplaceBackend):
84
86
  }
85
87
  if res.gpu_type is not None:
86
88
  filt["gpu_name"] = {"in": vast_gpu_names(res.gpu_type)}
89
+ if res.gpu_type == "A100-80":
90
+ # cheapest-first + OFFER_LIMIT would otherwise fill the page
91
+ # with 40GB A100s that normalization then discards
92
+ filt["gpu_ram"] = {"gte": 70_000}
87
93
  elif res.min_vram_gb is not None:
88
94
  # vast reports gpu_ram in MB (verify live)
89
95
  filt["gpu_ram"] = {"gte": res.min_vram_gb * 1024}
@@ -188,6 +188,33 @@ def _indent(block: str, n: int) -> str:
188
188
  return "".join(pad + line + "\n" for line in block.splitlines())
189
189
 
190
190
 
191
+ def _heredoc_delim(body: str) -> str:
192
+ """A heredoc terminator that provably does not occur as a line in `body`,
193
+ so embedding an arbitrary command in a heredoc can never terminate early."""
194
+ lines = body.splitlines()
195
+ base = "OMNIRUN_CMD_EOF"
196
+ delim = base
197
+ i = 0
198
+ while delim in lines:
199
+ i += 1
200
+ delim = f"{base}_{i}"
201
+ return delim
202
+
203
+
204
+ def _command_block(command: str) -> str:
205
+ """Body of run_cmd(): the user command carried byte-exact via a single-quoted
206
+ heredoc and run with `eval`.
207
+
208
+ A single-quoted heredoc disables every expansion and preserves the body
209
+ literally — newlines, embedded heredocs, quotes, backslashes all survive — so
210
+ multi-line scripts reach the worker unmangled. `eval` runs it in the current
211
+ shell, so env activation, forwarded exports and pre-run lines all still apply.
212
+ Crucially the command is NOT indented: indenting would shift the user's own
213
+ heredoc terminators off column 0 and break them."""
214
+ delim = _heredoc_delim(command)
215
+ return f"eval \"$(cat <<'{delim}'\n{command}\n{delim}\n)\"\n"
216
+
217
+
191
218
  def _default_project_root(slug: str) -> str:
192
219
  return f"$OMNIRUN_ROOT/projects/{slug}"
193
220
 
@@ -307,8 +334,7 @@ status running
307
334
  STARTED_AT=$(now)
308
335
  set +e
309
336
  run_cmd() {{
310
- {_indent(spec.command, 2)}\
311
- }}
337
+ {_command_block(spec.command)}}}
312
338
  # run in a subshell: a bare `exit N` in the user command must not kill
313
339
  # bootstrap.sh itself (result.json would never be written -> job shows LOST)
314
340
  ( run_cmd ) > >(tee "$JOB_DIR/logs/stdout.log") 2> >(tee "$JOB_DIR/logs/stderr.log" >&2)
@@ -48,8 +48,10 @@ class BackendConfig(BaseModel, extra="allow"):
48
48
  env_setup: list[str] = Field(default_factory=list) # module loads, exports
49
49
  # Where a project's shared checkout + .venv live on the worker. Default
50
50
  # "$OMNIRUN_ROOT/projects/<slug>"; point at an existing clone to reuse it
51
- # (its .git becomes the object store, its .venv the shared env).
52
- project_root: str | None = None
51
+ # (its .git becomes the object store, its .venv the shared env). A str
52
+ # applies to every repo; a dict maps repo slug -> path (with an optional
53
+ # "default" key as fallback) so one backend can serve several repos.
54
+ project_root: str | dict[str, str] | None = None
53
55
 
54
56
  # static capability declaration (ssh backend; probe checks live via nvidia-smi)
55
57
  gpus: list[GpuDecl] = Field(default_factory=list)
@@ -74,6 +76,15 @@ class BackendConfig(BaseModel, extra="allow"):
74
76
  def extra(self, key: str, default: Any = None) -> Any:
75
77
  return (self.model_extra or {}).get(key, default)
76
78
 
79
+ def project_root_for(self, slug: str) -> str | None:
80
+ """Configured project root for a repo. A str applies to every repo; a
81
+ dict is keyed by slug, falling back to a "default" key, then to None
82
+ (the built-in "$OMNIRUN_ROOT/projects/<slug>" path)."""
83
+ pr = self.project_root
84
+ if isinstance(pr, dict):
85
+ return pr.get(slug) or pr.get("default")
86
+ return pr
87
+
77
88
 
78
89
  class DaemonConfig(BaseModel):
79
90
  host: str = "127.0.0.1" # localhost socket; bind 0.0.0.0 + auth to go remote
@@ -133,6 +133,38 @@ def test_failing_command_propagates_exit_code(
133
133
  assert "about to fail" in (job_dir / "logs" / "stdout.log").read_text()
134
134
 
135
135
 
136
+ def test_multiline_heredoc_command_runs_byte_exact(
137
+ sample_repo: Path, job_spec: JobSpec, tmp_path: Path
138
+ ) -> None:
139
+ """Regression (#3): a multi-line command with a heredoc must reach the worker
140
+ byte-identical. The old indent-per-line embedding shifted the heredoc
141
+ terminator off column 0 (breaking it) and prefixed every body line — this
142
+ proves the body is written verbatim, indentation and all."""
143
+ command = (
144
+ 'cat <<EOF > "$OMNIRUN_OUTPUT/cfg.yaml"\n'
145
+ "key: value\n"
146
+ " two-space-indent: kept\n"
147
+ "EOF\n"
148
+ 'echo "MULTILINE OK"'
149
+ )
150
+ spec = job_spec.model_copy(
151
+ update={
152
+ "job_id": JobSpec.make_job_id("heredoc"),
153
+ "command": command,
154
+ "outputs": [],
155
+ }
156
+ )
157
+ root = tmp_path / "omnirun_root"
158
+ script = stage(sample_repo, spec, root)
159
+ proc = run_bootstrap(script)
160
+
161
+ job_dir = root / "jobs" / spec.job_id
162
+ assert proc.returncode == 0, (job_dir / "logs" / "bootstrap.log").read_text()
163
+ cfg = job_dir / "outputs" / "cfg.yaml"
164
+ assert cfg.read_text() == "key: value\n two-space-indent: kept\n"
165
+ assert "MULTILINE OK" in (job_dir / "logs" / "stdout.log").read_text()
166
+
167
+
136
168
  def test_bare_exit_in_command_still_writes_result(
137
169
  sample_repo: Path, job_spec: JobSpec, tmp_path: Path
138
170
  ) -> None:
@@ -0,0 +1,37 @@
1
+ """BackendConfig helpers: project_root_for resolution (str / dict / None)."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from omnirun.config import BackendConfig
6
+
7
+
8
+ def make_config(**kw) -> BackendConfig:
9
+ return BackendConfig.model_validate({"type": "slurm", **kw})
10
+
11
+
12
+ def test_project_root_for_none_is_default():
13
+ assert make_config().project_root_for("proj") is None
14
+
15
+
16
+ def test_project_root_for_str_applies_to_every_repo():
17
+ c = make_config(project_root="$HOME/shared")
18
+ assert c.project_root_for("proj") == "$HOME/shared"
19
+ assert c.project_root_for("other") == "$HOME/shared"
20
+
21
+
22
+ def test_project_root_for_dict_slug_hit():
23
+ c = make_config(project_root={"hns": "$HOME/projects/hns"})
24
+ assert c.project_root_for("hns") == "$HOME/projects/hns"
25
+
26
+
27
+ def test_project_root_for_dict_no_match_falls_through_to_none():
28
+ c = make_config(project_root={"hns": "$HOME/projects/hns"})
29
+ assert c.project_root_for("auraflow") is None
30
+
31
+
32
+ def test_project_root_for_dict_default_fallback():
33
+ c = make_config(
34
+ project_root={"hns": "$HOME/projects/hns", "default": "$HOME/managed"}
35
+ )
36
+ assert c.project_root_for("hns") == "$HOME/projects/hns" # slug wins over default
37
+ assert c.project_root_for("auraflow") == "$HOME/managed" # fallback
@@ -496,6 +496,49 @@ def test_pull_outputs(fake_api, backend, tmp_path):
496
496
  assert files == [dest / "model.txt"]
497
497
 
498
498
 
499
+ def test_pull_outputs_tolerates_absolute_symlink(fake_api, backend, tmp_path):
500
+ """Regression (#1): every W&B run leaves a symlink to an absolute path under
501
+ wandb/*/logs/. Python's data filter rejects it; pull must skip that one entry
502
+ and still recover the real outputs rather than aborting the whole archive."""
503
+ buf = io.BytesIO()
504
+ with tarfile.open(fileobj=buf, mode="w:gz") as tf:
505
+ data = b"metric,value\n1,2\n"
506
+ good = tarfile.TarInfo("outputs/results.csv")
507
+ good.size = len(data)
508
+ tf.addfile(good, io.BytesIO(data))
509
+ link = tarfile.TarInfo("outputs/wandb/run-x/logs/debug-core.log")
510
+ link.type = tarfile.SYMTYPE
511
+ link.linkname = "/kaggle/tmp/omnirun/abs/debug-core.log" # absolute target
512
+ tf.addfile(link)
513
+ fake_api.output_files = {"omnirun-job.tar.gz": buf.getvalue()}
514
+ dest = tmp_path / "results"
515
+ files = backend.pull_outputs(make_handle(), dest)
516
+ assert (dest / "results.csv").read_bytes() == data
517
+ assert [p.name for p in files] == ["results.csv"] # link skipped, csv kept
518
+
519
+
520
+ def test_submit_guards_oversized_source(fake_api, monkeypatch):
521
+ """Regression (#2): an oversized kernel source must fail early with a
522
+ size-naming message, not sail past the guard and fail opaquely on Kaggle."""
523
+ monkeypatch.setattr(
524
+ kaggle_mod, "_create_bundle", lambda root, sha, dest: _big_bundle(dest)
525
+ )
526
+ backend = KaggleBackend(
527
+ "kaggle",
528
+ BackendConfig.model_validate({"type": "kaggle", "max_source_bytes": 4096}),
529
+ )
530
+ spec = make_spec(gpu_type="P100")
531
+ offer = backend.probe(spec.resources)[0]
532
+ with pytest.raises(BackendError, match=r"kaggle kernel source is .* over the"):
533
+ backend.submit(spec, offer)
534
+
535
+
536
+ def _big_bundle(dest: Path) -> Path:
537
+ dest = Path(dest)
538
+ dest.write_bytes(b"x" * 8192) # base64 ~11k, over the 4k test limit
539
+ return dest
540
+
541
+
499
542
  def test_cancel_without_api_support_points_at_website(fake_api, backend):
500
543
  with pytest.raises(BackendError, match="kaggle.com"):
501
544
  backend.cancel(make_handle())
@@ -474,7 +474,7 @@ def test_vast_probe_filter_and_parsing():
474
474
 
475
475
  body = json.loads(route.calls[0].request.content)
476
476
  assert body["type"] == "ondemand"
477
- assert body["gpu_name"] == {"in": ["H100_SXM", "H100_PCIE", "H100_NVL"]}
477
+ assert body["gpu_name"] == {"in": ["H100 SXM", "H100 PCIE", "H100 NVL"]}
478
478
  assert body["num_gpus"] == {"eq": 1}
479
479
  assert body["rentable"] == {"eq": True}
480
480
  assert body["verified"] == {"eq": True}
@@ -295,6 +295,22 @@ def test_render_payload_sbatch_then_bootstrap():
295
295
  assert fake.files == {}
296
296
 
297
297
 
298
+ def test_render_payload_project_root_dict_matches_slug():
299
+ # spec slug is "proj": a per-repo dict resolves to that repo's checkout...
300
+ fake = FakeExec()
301
+ b = make_backend(fake, root=ROOT, project_root={"proj": "/data/existing/proj"})
302
+ payload = b.render_payload(make_spec(), offer=None)
303
+ assert 'PROJECT_ROOT="/data/existing/proj"' in payload
304
+
305
+
306
+ def test_render_payload_project_root_dict_no_match_uses_default():
307
+ # ...and an unrelated slug falls through to the built-in "$root/projects/<slug>".
308
+ fake = FakeExec()
309
+ b = make_backend(fake, root=ROOT, project_root={"other": "/data/existing/other"})
310
+ payload = b.render_payload(make_spec(), offer=None)
311
+ assert f'PROJECT_ROOT="{ROOT}/projects/proj"' in payload
312
+
313
+
298
314
  # --- status mapping -------------------------------------------------------------------
299
315
 
300
316
 
@@ -515,6 +531,16 @@ def test_probe_warns_when_type_requested_but_no_map():
515
531
  assert "gpu_map" in offer.notes
516
532
 
517
533
 
534
+ def test_probe_unfit_for_gpu_job_when_cpu_only():
535
+ fake = FakeExec()
536
+ (offer,) = make_backend(fake, has_gpus=False).probe(ResourceSpec(gpus=1))
537
+ assert not offer.fits
538
+ assert "CPU-only" in offer.unfit_reasons[0]
539
+ # CPU jobs still fit
540
+ (offer,) = make_backend(fake, has_gpus=False).probe(ResourceSpec())
541
+ assert offer.fits
542
+
543
+
518
544
  def test_probe_wait_tier1_idle_nodes():
519
545
  fake = FakeExec()
520
546
  fake.add(r"sinfo -p gpu -t idle", stdout="node01 gpu:a100:4\nnode02 (null)\n")
@@ -0,0 +1,53 @@
1
+ """Tolerant tar extraction (shared by the kaggle + colab pull paths)."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import io
6
+ import os
7
+ import tarfile
8
+ from pathlib import Path
9
+
10
+ from omnirun.backends import tarsafe
11
+
12
+
13
+ def _tar_with_abs_symlink(dest_target: str) -> io.BytesIO:
14
+ buf = io.BytesIO()
15
+ with tarfile.open(fileobj=buf, mode="w:gz") as tf:
16
+ data = b"real output\n"
17
+ good = tarfile.TarInfo("outputs/results.txt")
18
+ good.size = len(data)
19
+ tf.addfile(good, io.BytesIO(data))
20
+ link = tarfile.TarInfo("outputs/wandb/run-x/logs/debug-core.log")
21
+ link.type = tarfile.SYMTYPE
22
+ link.linkname = dest_target
23
+ tf.addfile(link)
24
+ buf.seek(0)
25
+ return buf
26
+
27
+
28
+ def test_extract_all_skips_absolute_symlink_and_keeps_rest(tmp_path: Path) -> None:
29
+ # every W&B run leaves such a link; the default data filter would abort here.
30
+ buf = _tar_with_abs_symlink("/etc/hostname")
31
+ with tarfile.open(fileobj=buf) as tf:
32
+ skipped = tarsafe.extract_all(tf, tmp_path)
33
+ assert skipped == ["outputs/wandb/run-x/logs/debug-core.log"]
34
+ assert (tmp_path / "outputs" / "results.txt").read_bytes() == b"real output\n"
35
+ assert not (tmp_path / "outputs" / "wandb" / "run-x" / "logs").exists()
36
+
37
+
38
+ def test_extract_all_extracts_normal_tree_completely(tmp_path: Path) -> None:
39
+ src = tmp_path / "src"
40
+ (src / "sub").mkdir(parents=True)
41
+ (src / "a.txt").write_text("a")
42
+ (src / "sub" / "b.txt").write_text("b")
43
+ os.symlink("a.txt", src / "rel") # a *relative* symlink is allowed by data filter
44
+ tar = tmp_path / "t.tar.gz"
45
+ with tarfile.open(tar, "w:gz") as tf:
46
+ tf.add(src, arcname="outputs")
47
+ out = tmp_path / "out"
48
+ with tarfile.open(tar) as tf:
49
+ skipped = tarsafe.extract_all(tf, out)
50
+ assert skipped == []
51
+ assert (out / "outputs" / "a.txt").read_text() == "a"
52
+ assert (out / "outputs" / "sub" / "b.txt").read_text() == "b"
53
+ assert (out / "outputs" / "rel").is_symlink()
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