diff --git a/docs/tongs/approval.md b/docs/tongs/approval.md index 2b19248..2d62873 100644 --- a/docs/tongs/approval.md +++ b/docs/tongs/approval.md @@ -1,9 +1,9 @@ # First-run approval The user, org, and repo layers are installed deliberately and are trusted; they skip the gate. -A **workspace**-sourced tong (from a repo you cloned) could otherwise request your secrets, host mounts, or the docker socket simply by being present, so the launcher gates it: +A **workspace**-sourced tong (from a repo you cloned) could otherwise request your secrets, host mounts, the docker socket, or the host's GPUs simply by being present, so the launcher gates it: -- Before starting, it prints exactly what the tong requests — image, secret references, mounts, networks, and docker-socket access — and asks you to approve. +- Before starting, it prints exactly what the tong requests — image, secret references, mounts, networks, docker-socket access, and GPU access — and asks you to approve. - Approval is keyed by workspace path + tong name + a hash of the merged definition, stored in `~/.swarmforge/approvals.json`. Any change to the definition re-prompts. - The gate defaults to **No**, and a non-interactive stdin reads as No. A scripted `--no-prompt` run **fails closed** rather than auto-approving. - Approving `image: foo:latest` approves a moving target; **pinned digests are the recommended convention** for workspace tongs. diff --git a/docs/tongs/definitions.md b/docs/tongs/definitions.md index bb1cafb..5d4687c 100644 --- a/docs/tongs/definitions.md +++ b/docs/tongs/definitions.md @@ -18,6 +18,7 @@ mounts: # opt-in magic words only, never raw host paths - workspace:ro # or workspace:/code:ro to bind it somewhere else resources: memory: 512m # string or number + # gpus: all # optional host GPU access (see Resources) networks: # optional extra pre-existing networks to also join - some-existing-net # entrypoint: [...] # optional argv override of the image ENTRYPOINT @@ -68,3 +69,21 @@ Mounts are opt-in **magic words**, never raw host paths. Only two are recognized - `workspace[:/target][:mode]` — bind-mounts the session workspace, at `/workspace` unless an absolute `target` says otherwise (e.g. `workspace:ro`, `workspace:/code`, `workspace:/code:ro`). A custom target lets an image that expects its sources elsewhere be used unmodified (it does not set the working directory — the process still starts in the image's own `WORKDIR`). A target is refused unless it is an absolute path free of whitespace, and refused if it resolves to `/`, overlaps another of the tong's mounts, or overlaps a path the tong's own wiring occupies — the secret-delivery tmpfs at `/run/swarmforge` and the `/bin/sh` its wrapper execs (for a tong with secret references), or the docker socket (for a tong that mounts it). The workspace bind is paired with the same git-dir mounts the anvil gets from `swarmforge/gitguard.py`: read-only guards over the config and hooks the host's git obeys, and — when the workspace is a linked worktree or another checkout whose git dir lives outside it — that git dir at its own absolute path, which is where the checkout's `.git` pointer file says to look (without it, git inside the tong fails with "not a git repository"). The one mount the anvil gets and a tong does not is the restated worktree registration — a tong keeps the host's copy, so `git worktree list` inside one names the checkout by its host path (the [git guard](../git-guard.md) has the rest). When every `workspace` mount is `ro`, the ride-along git-dir mounts are forced read-only too. - `docker-socket[:mode]` — bind-mounts the host docker socket onto the same path inside the container (so it takes no target). This is full host docker control and is always called out explicitly in the workspace approval prompt; it is the grant a broker tong needs. + +## Resources + +```yaml +resources: + memory: 8g # docker --memory: a string or number + gpus: all # docker --gpus: all, a count, or devices +``` + +Both are optional and passed to `docker run` as written. `gpus` grants the tong the host's GPUs and accepts exactly these forms: + +- `all` — every GPU on the host. +- A positive count, e.g. `2` (at most 2147483647). +- `device=` for one device, or `"device=,,..."` — double quotes included — for several, where an id is an index or a `GPU-`/`MIG-` UUID: a letter or digit, then letters, digits, and `._:-`. Docker reads the `--gpus` value as CSV, so a list needs those literal double quotes inside the value, which in YAML means wrapping it in single quotes: `gpus: '"device=0,1"'`. An unquoted `device=0,1` is refused with that hint. + +Docker's other `--gpus` keys are deliberately refused. `driver=`, `capabilities=`, and `options=` can grant more than GPU access — `driver=cdi`, for one, injects any device registered on the host — and `count=` is redundant with a plain count. +GPU access needs a docker host set up for GPU containers — the NVIDIA Container Toolkit on Linux; Docker Desktop's WSL 2 backend provides GPU support itself. Without it `docker run` fails and the launch stops. +A workspace tong's GPU request is called out in its [approval prompt](approval.md). diff --git a/swarmforge/anvil/approval.py b/swarmforge/anvil/approval.py index 5e562bd..833a35f 100644 --- a/swarmforge/anvil/approval.py +++ b/swarmforge/anvil/approval.py @@ -2,10 +2,11 @@ The user, org, and Swarmforge-repo layers are installed deliberately and are trusted. The workspace is any repo you happened to clone, so a workspace-sourced -tong -- which may request secrets, host mounts, or the docker socket -- is gated: -its privilege summary is rendered and the user is asked to approve it before the -anvil starts. Approval is keyed by workspace path, tong name, and a hash of the -merged definition, so any change re-prompts, and persists in the user layer. +tong -- which may request secrets, host mounts, the docker socket, or host +GPUs -- is gated: its privilege summary is rendered and the user is asked to +approve it before the anvil starts. Approval is keyed by workspace path, tong +name, and a hash of the merged definition, so any change re-prompts, and +persists in the user layer. Deciding what a definition asks for and remembering the answer are the pure core's job; rendering the question and reading the reply are this module's. @@ -26,7 +27,8 @@ def render_privilege_summary(name, summary): `summary` is the structured output of `tongs.privilege_summary`. Only the privileges actually requested are shown, and docker-socket access -- the broadest grant, since it is full control of the host's docker -- is always - called out explicitly so it cannot be approved unseen. + called out explicitly so it cannot be approved unseen, as is a GPU request, + which grants the host's GPU devices. """ lines = ["Workspace tong %r requests approval:" % name] lines.append(" image: %s" % (summary.get("image") or "(none declared)")) @@ -42,6 +44,8 @@ def render_privilege_summary(name, summary): lines.append(" networks: %s" % ", ".join(str(n) for n in networks)) if summary.get("socket"): lines.append(" docker socket: full host docker control") + if summary.get("gpus") is not None: + lines.append(" gpus: %s (host GPU access)" % summary["gpus"]) return "\n".join(lines) diff --git a/swarmforge/tongs/__init__.py b/swarmforge/tongs/__init__.py index 07c6bd4..5572266 100644 --- a/swarmforge/tongs/__init__.py +++ b/swarmforge/tongs/__init__.py @@ -11,7 +11,7 @@ This package is the pure core of the tongs launcher, one module per concern: - model the vocabulary of a definition, and readiness resolution + model the vocabulary of a definition, readiness and GPU-request resolution discovery reading the layer directories and merging by name validate schema validation secrets secret references, the provider table, FIFO delivery @@ -87,6 +87,7 @@ from .model import ( DEFAULT_READINESS_TIMEOUT_S, ENV_PREFIX, + GPU_COUNT_MAX, INTERFACE_KINDS, LABEL_CONFIG_HASH, LABEL_TONG_NAME, @@ -102,6 +103,7 @@ WORKSPACE, WORKSPACE_HOST_ENV, parse_duration, + parse_gpus, readiness_settings, warn, ) @@ -187,6 +189,7 @@ "tong_env_var", "DEFAULT_READINESS_TIMEOUT_S", "ENV_PREFIX", + "GPU_COUNT_MAX", "INTERFACE_KINDS", "LABEL_CONFIG_HASH", "LABEL_TONG_NAME", @@ -202,6 +205,7 @@ "WORKSPACE", "WORKSPACE_HOST_ENV", "parse_duration", + "parse_gpus", "readiness_settings", "warn", "DEFAULT_DOCKER_SOCKET", diff --git a/swarmforge/tongs/approvals.py b/swarmforge/tongs/approvals.py index bc20388..37d9a4d 100644 --- a/swarmforge/tongs/approvals.py +++ b/swarmforge/tongs/approvals.py @@ -32,9 +32,9 @@ def privilege_summary(defn): """Structured summary of what a definition asks for, for the approval gate. Gathers the privileges a reviewer must see before approving a - workspace-sourced tong: image, secret references, mounts, networks, and - docker-socket access. Rendering and prompting are the caller's job; this - just assembles the facts. + workspace-sourced tong: image, secret references, mounts, networks, + docker-socket access, and host GPU access. Rendering and prompting are the + caller's job; this just assembles the facts. """ return { "image": defn.get("image"), @@ -42,9 +42,20 @@ def privilege_summary(defn): "mounts": list(defn.get("mounts") or []), "networks": list(defn.get("networks") or []), "socket": _has_socket_mount(defn), + "gpus": _requested_gpus(defn), } +def _requested_gpus(defn): + """The raw `resources.gpus` a definition declares, or None. + + The gate runs before validation, so the value is reported as declared + rather than parsed: anything that might reach docker is shown. + """ + resources = defn.get("resources") + return resources.get("gpus") if isinstance(resources, dict) else None + + def is_workspace_sourced(source_layer): """True if a tong's winning layer is the (untrusted) workspace and so gates.""" return source_layer == WORKSPACE diff --git a/swarmforge/tongs/argv.py b/swarmforge/tongs/argv.py index d714d67..4020bf6 100644 --- a/swarmforge/tongs/argv.py +++ b/swarmforge/tongs/argv.py @@ -12,7 +12,7 @@ from swarmforge.names import canonical_path, dir_hint, path_token, sanitize_token from .mcp import _is_network_facing, _ordered_aliases -from .model import LABEL_CONFIG_HASH, LABEL_TONG_NAME, WORKSPACE_HOST_ENV +from .model import LABEL_CONFIG_HASH, LABEL_TONG_NAME, WORKSPACE_HOST_ENV, parse_gpus from .mounts import DEFAULT_DOCKER_SOCKET, _has_socket_mount, tong_mount_specs from .secrets import SECRET_FIFO_TMPFS, declared_run_override @@ -90,9 +90,11 @@ def session_container_name(session_id, name): def tong_resource_flags(defn): """docker resource flags from a tong's `resources:` block. - v1 understands `memory` (mapped to `--memory`). Unknown keys are ignored for - forward compatibility. Raises `ValueError` if `resources` is present but not a - mapping. + v1 understands `memory` (mapped to `--memory`) and `gpus` (mapped to `--gpus` + verbatim once `parse_gpus` accepts it: `all`, a count, or a device + selector). Unknown keys are ignored for forward compatibility. Raises + `ValueError` if `resources` is present but not a mapping, or if `gpus` is + not a request `parse_gpus` accepts. """ resources = defn.get("resources") if resources is None: @@ -103,6 +105,9 @@ def tong_resource_flags(defn): memory = resources.get("memory") if memory is not None: flags += ["--memory", str(memory)] + gpus = parse_gpus(resources.get("gpus")) + if gpus is not None: + flags += ["--gpus", gpus] return flags diff --git a/swarmforge/tongs/model.py b/swarmforge/tongs/model.py index 0649876..56a14ce 100644 --- a/swarmforge/tongs/model.py +++ b/swarmforge/tongs/model.py @@ -3,7 +3,8 @@ The leaf of the package: it imports nothing else from `swarmforge.tongs`, so every other module can depend on it. It holds the constant sets the launcher dispatches on (layers, lifecycles, interface kinds, readiness modes), the docker -labels it stamps, and the pure resolution of a definition's `readiness:` block. +labels it stamps, and the pure resolution of a definition's `readiness:` block +and `resources.gpus` request. `warn` lives here too, so the one `tongs: ` stderr prefix has a single home. """ @@ -70,6 +71,48 @@ def parse_duration(value, default=None): return seconds +# Of docker's --gpus keys only `device`: driver/capabilities/options reach beyond GPUs. +_GPU_ID = r"[A-Za-z0-9][A-Za-z0-9._:-]*" +_GPU_COUNT_RE = re.compile(r"[0-9]+") +# Docker reads a count with Go's Atoi; a 32-bit int fits on every platform it builds for. +GPU_COUNT_MAX = 2**31 - 1 +_GPU_DEVICE_RE = re.compile(r'device=%s|"device=%s(?:,%s)*"' % (_GPU_ID, _GPU_ID, _GPU_ID)) +_GPU_UNQUOTED_LIST_RE = re.compile(r"device=%s(?:,%s)+" % (_GPU_ID, _GPU_ID)) +_GPU_FORMS = ( + "'all', a positive GPU count, 'device=', or '\"device=,,...\"'" +) + + +def parse_gpus(value): + """Resolve a `resources.gpus` declaration to its `docker run --gpus` value. + + `None` (no GPUs requested) yields None. Accepted, and returned verbatim as + a string: `all`; a count from 1 to `GPU_COUNT_MAX`, as an int or a string + of ASCII digits; or a device selector, `device=` or the docker-quoted + list `"device=,,..."` (literal double quotes, since docker reads + the value as CSV), each id starting with a letter or digit followed by + letters, digits, and `._:-`. Raises `ValueError` + for anything else -- including docker's other `--gpus` keys, which can + reach devices beyond GPUs -- so validation and argv assembly share one + definition of a usable request. + """ + if value is None: + return None + if _is_int(value) and 0 < value <= GPU_COUNT_MAX: + return str(value) + if isinstance(value, str): + if value == "all" or _GPU_DEVICE_RE.fullmatch(value): + return value + if _GPU_COUNT_RE.fullmatch(value) and 0 < int(value) <= GPU_COUNT_MAX: + return value + if _GPU_UNQUOTED_LIST_RE.fullmatch(value): + raise ValueError( + "a device list needs literal double quotes because docker reads " + "--gpus as CSV; write gpus: '\"%s\"', got %r" % (value, value) + ) + raise ValueError("must be %s, got %r" % (_GPU_FORMS, value)) + + def readiness_settings(defn): """Resolve a tong's readiness declaration to `(mode, command, timeout_s)`. diff --git a/swarmforge/tongs/validate.py b/swarmforge/tongs/validate.py index 525a213..02dfe99 100644 --- a/swarmforge/tongs/validate.py +++ b/swarmforge/tongs/validate.py @@ -4,8 +4,8 @@ fields the launcher dispatches on, so a malformed definition becomes a list of errors here rather than a docker failure mid-launch. Every check that would otherwise restate launcher logic -- the mount grammar and its targets, secret -env names, readiness durations -- calls the same function the launcher calls, so -validation and assembly can never disagree. +env names, readiness durations, GPU requests -- calls the same function the +launcher calls, so validation and assembly can never disagree. """ import re @@ -18,6 +18,7 @@ TRANSPORTS, _is_int, parse_duration, + parse_gpus, ) from .mounts import ( mount_destination, @@ -183,5 +184,9 @@ def err(msg): isinstance(memory, str) or _is_int(memory) or isinstance(memory, float) ): err("resources.memory must be a string or number") + try: + parse_gpus(resources.get("gpus")) + except ValueError as exc: + err("resources.gpus: %s" % exc) return errors diff --git a/tests/test_anvil_approval.py b/tests/test_anvil_approval.py index 412297c..615af6d 100644 --- a/tests/test_anvil_approval.py +++ b/tests/test_anvil_approval.py @@ -46,11 +46,24 @@ def test_renders_requested_privileges(self): self.assertIn("workspace:ro", text) self.assertIn("docker socket", text) + def test_calls_out_requested_gpus(self): + defn = dict(WORKSPACE_TONG, resources={"gpus": '"device=0,1"'}) + text = launcher.render_privilege_summary("github", tongs.privilege_summary(defn)) + self.assertIn(' gpus: "device=0,1" (host GPU access)', text.splitlines()) + + def test_non_mapping_resources_renders_no_gpus_line(self): + for resources in ("8g", ["gpus"]): + with self.subTest(resources=resources): + defn = dict(WORKSPACE_TONG, resources=resources) + text = launcher.render_privilege_summary("github", tongs.privilege_summary(defn)) + self.assertNotIn("gpus", text) + def test_omits_unrequested_sections(self): - defn = {"image": "x", "interface": {"kind": "none"}} + defn = {"image": "x", "interface": {"kind": "none"}, "resources": {"memory": "1g"}} text = launcher.render_privilege_summary("x", tongs.privilege_summary(defn)) self.assertNotIn("secrets:", text) self.assertNotIn("docker socket", text) + self.assertNotIn("gpus", text) class GateTests(unittest.TestCase): @@ -124,6 +137,13 @@ def test_a_changed_definition_is_unapproved(self): with self.assertRaises(launcher.ApprovalDenied): self._gate(_merged("gh", changed), prompt=False) + def test_adding_gpus_to_an_approved_tong_reprompts_and_shows_them(self): + self._gate(_merged("gh", WORKSPACE_TONG), answer="y\n") + changed = dict(WORKSPACE_TONG, resources={"gpus": "all"}) + out = self._gate(_merged("gh", changed), answer="y\n") + self.assertIn("gpus: all (host GPU access)", out) + self.assertIn("Approve workspace tong 'gh'?", out) + def test_missing_workspace_path_fails_closed(self): merged = _merged("gh", WORKSPACE_TONG) with self.assertRaises(launcher.ApprovalDenied): diff --git a/tests/test_anvil_orchestrate.py b/tests/test_anvil_orchestrate.py index 8e7a64a..20d627e 100644 --- a/tests/test_anvil_orchestrate.py +++ b/tests/test_anvil_orchestrate.py @@ -368,6 +368,19 @@ def test_shared_tong_recreated_when_hash_differs(self): self.assertIn(("rm_force", "swarmforge-shared-ollama"), docker.calls) self.assertEqual(len(docker.run_argvs), 1) + def test_shared_tong_recreated_with_gpus_when_a_request_is_added(self): + states = {"swarmforge-shared-ollama": + {"running": True, "label": tongs.config_hash(SHARED_OLLAMA)}} + docker = FakeDocker(states=states) + defn = dict(SHARED_OLLAMA, resources={"gpus": "all"}) + self._run(docker, _merged("ollama", defn, source=tongs.REPO)) + self.assertIn(("rm_force", "swarmforge-shared-ollama"), docker.calls) + self.assertEqual(len(docker.run_argvs), 1) + started = docker.run_argvs[0] + flag = started.index("--gpus") + self.assertEqual(started[flag + 1], "all") + self.assertLess(flag, started.index("ollama/ollama")) + def test_shared_tong_recreated_when_absent(self): # rm_force clears any stopped leftover before the fresh start. docker = FakeDocker() diff --git a/tests/test_tongs_approvals.py b/tests/test_tongs_approvals.py index 1860f04..c109c89 100644 --- a/tests/test_tongs_approvals.py +++ b/tests/test_tongs_approvals.py @@ -49,6 +49,30 @@ def test_summary_reports_requested_privileges(self): def test_no_socket_without_mount(self): self.assertFalse(tongs.privilege_summary(def_of(GITHUB_TONG))["socket"]) + def test_summary_reports_requested_gpus(self): + defn = def_of(GITHUB_TONG) + defn["resources"] = {"memory": "8g", "gpus": "all"} + self.assertEqual(tongs.privilege_summary(defn)["gpus"], "all") + + def test_no_gpus_without_request(self): + defn = def_of(GITHUB_TONG) + self.assertIsNone(tongs.privilege_summary(defn)["gpus"]) + defn["resources"] = {"memory": "8g"} + self.assertIsNone(tongs.privilege_summary(defn)["gpus"]) + + def test_non_mapping_resources_reports_no_gpus(self): + for resources in ("8g", ["gpus"]): + with self.subTest(resources=resources): + defn = def_of(GITHUB_TONG) + defn["resources"] = resources + self.assertIsNone(tongs.privilege_summary(defn)["gpus"]) + + def test_summary_reports_an_unvalidated_gpus_value_as_declared(self): + # The gate runs before validation, so a value is shown rather than parsed. + defn = def_of(GITHUB_TONG) + defn["resources"] = {"gpus": 0} + self.assertEqual(tongs.privilege_summary(defn)["gpus"], 0) + class ApprovalKeyingTests(unittest.TestCase): def setUp(self): @@ -68,6 +92,13 @@ def test_definition_change_reprompts(self): changed["image"] = "ghcr.io/crypticswarm/github-tong@sha256:MOVED" self.assertFalse(tongs.is_approved(approvals, self.ws, "github", changed)) + def test_gpu_request_reprompts(self): + approvals = {} + tongs.record_approval(approvals, self.ws, "github", self.defn) + changed = def_of(GITHUB_TONG) + changed["resources"] = {"gpus": "all"} + self.assertFalse(tongs.is_approved(approvals, self.ws, "github", changed)) + def test_keyed_by_workspace_path(self): approvals = {} tongs.record_approval(approvals, self.ws, "github", self.defn) diff --git a/tests/test_tongs_argv.py b/tests/test_tongs_argv.py index c983044..140d503 100644 --- a/tests/test_tongs_argv.py +++ b/tests/test_tongs_argv.py @@ -80,6 +80,22 @@ def test_session_container_name_empty_token_has_no_trailing_dash(self): def test_resource_flags_memory(self): self.assertEqual(tongs.tong_resource_flags({"resources": {"memory": "512m"}}), ["--memory", "512m"]) + def test_resource_flags_gpus_passed_through_verbatim(self): + for declared, flag in (("all", "all"), (2, "2"), ('"device=0,1"', '"device=0,1"')): + with self.subTest(gpus=declared): + self.assertEqual( + tongs.tong_resource_flags({"resources": {"gpus": declared}}), + ["--gpus", flag], + ) + + def test_resource_flags_memory_and_gpus_together(self): + flags = tongs.tong_resource_flags({"resources": {"memory": "8g", "gpus": "all"}}) + self.assertEqual(flags, ["--memory", "8g", "--gpus", "all"]) + + def test_resource_flags_unusable_gpus_raises(self): + with self.assertRaises(ValueError): + tongs.tong_resource_flags({"resources": {"gpus": 0}}) + def test_resource_flags_absent_is_empty(self): self.assertEqual(tongs.tong_resource_flags({}), []) @@ -148,6 +164,17 @@ def test_run_argv_mounts_and_resources(self): self.assertIn("--memory", argv) self.assertEqual(argv[argv.index("--memory") + 1], "256m") + def test_run_argv_gpus_precede_the_image(self): + defn = def_of(PORT_TONG) + defn["resources"] = {"gpus": "all"} + argv = tongs.tong_run_argv( + "pg", defn, container_name="ctr-pg", network="net", alias="pg", + ) + flag = argv.index("--gpus") + self.assertEqual(argv[flag + 1], "all") + self.assertLess(flag, argv.index("postgres:16")) + self.assertEqual(argv[-1], "postgres:16") + def test_run_argv_mount_target_is_workspace_unless_declared(self): def argv_for(mounts): defn = def_of(NONE_TONG) diff --git a/tests/test_tongs_model.py b/tests/test_tongs_model.py index 6b0a6c8..a332a52 100644 --- a/tests/test_tongs_model.py +++ b/tests/test_tongs_model.py @@ -64,5 +64,63 @@ def test_readiness_portless_without_mode_is_none(self): self.assertEqual(mode, "none") +class GpuRequestTests(unittest.TestCase): + def test_parse_gpus_none_is_not_requested(self): + self.assertIsNone(tongs.parse_gpus(None)) + + def test_parse_gpus_accepts_all_and_positive_counts(self): + self.assertEqual(tongs.parse_gpus("all"), "all") + self.assertEqual(tongs.parse_gpus(2), "2") + self.assertEqual(tongs.parse_gpus("2"), "2") + + def test_parse_gpus_accepts_device_selectors_verbatim(self): + for selector in ("device=0", "device=GPU-3a8f-11ee", '"device=0"', + '"device=0,1"', '"device=GPU-3a8f,MIG-7c2e"', "device=GPU-1:0"): + with self.subTest(selector=selector): + self.assertEqual(tongs.parse_gpus(selector), selector) + + def test_parse_gpus_rejects_non_positive_and_non_count_types(self): + for bad in (True, False, 0, -1, 1.5, ["all"], {"count": 1}): + with self.subTest(value=bad), self.assertRaises(ValueError): + tongs.parse_gpus(bad) + + def test_parse_gpus_rejects_counts_docker_would_not_read_as_positive(self): + # int() accepts several of these; docker's Atoi does not. + for bad in ("0", "00", "-1", "+1", "1_0", "\u0663", " 2", ""): + with self.subTest(value=bad), self.assertRaises(ValueError): + tongs.parse_gpus(bad) + + def test_parse_gpus_rejects_docker_keys_beyond_device(self): + for bad in ("count=2", "driver=cdi,device=vendor.com/class=x", + "capabilities=compute", "options=x", '"device=0",driver=nvidia', + "ALL", " all", "device=", '"device=0,"', "device=-0"): + with self.subTest(value=bad), self.assertRaises(ValueError): + tongs.parse_gpus(bad) + + def test_parse_gpus_rejects_control_characters(self): + for bad in ("\x1b[1A\x1b[2K", "all\x1b[2K", "device=0\r"): + with self.subTest(value=bad), self.assertRaises(ValueError) as caught: + tongs.parse_gpus(bad) + self.assertNotIn("\x1b", str(caught.exception)) + + def test_parse_gpus_unquoted_device_list_hints_at_the_quoting(self): + with self.assertRaisesRegex(ValueError, "double quotes") as caught: + tongs.parse_gpus("device=0,1") + self.assertIn("""gpus: '"device=0,1"'""", str(caught.exception)) + + def test_parse_gpus_caps_the_count(self): + limit = tongs.GPU_COUNT_MAX + self.assertEqual(tongs.parse_gpus(limit), str(limit)) + self.assertEqual(tongs.parse_gpus(str(limit)), str(limit)) + for bad in (limit + 1, str(limit + 1), 10**30, "99999999999999999999"): + with self.subTest(value=bad), self.assertRaises(ValueError): + tongs.parse_gpus(bad) + + def test_parse_gpus_matches_the_whole_value(self): + for bad in ("all\n", "2\n", "device=0\n", '"device=0,1"\n', "device=0,1\n"): + with self.subTest(value=bad), self.assertRaises(ValueError): + tongs.parse_gpus(bad) + + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/tests/test_tongs_validate.py b/tests/test_tongs_validate.py index 426b4be..b3ed75c 100644 --- a/tests/test_tongs_validate.py +++ b/tests/test_tongs_validate.py @@ -229,6 +229,20 @@ def test_resources_memory_type_checked(self): self.assertTrue(any("memory" in e for e in errors)) self.assertEqual(tongs.validate_tong("t", self._base(resources={"memory": "512m"})), []) + def test_resources_gpus_accepts_all_a_count_or_a_device_selector(self): + for gpus in ("all", 2, "2", "device=0", '"device=0,1"'): + with self.subTest(gpus=gpus): + self.assertEqual( + tongs.validate_tong("t", self._base(resources={"gpus": gpus})), [] + ) + + def test_resources_gpus_rejects_unusable_requests(self): + for gpus in (True, 0, -2, 1.5, "", " ", "0", ["all"], "count=2", + "driver=cdi,device=x", "device=0,1", "all\x1b[2K"): + with self.subTest(gpus=gpus): + errors = tongs.validate_tong("t", self._base(resources={"gpus": gpus})) + self.assertTrue(any("resources.gpus" in e for e in errors), errors) + def test_volume_requires_volume_mountpoint_and_readiness_mode(self): errors = tongs.validate_tong("t", {"lifecycle": "session", "image": "x", "interface": {"kind": "volume"}}) joined = " ".join(errors)