From 933c0e582a458e0528dab33193bde22c0ecffbf3 Mon Sep 17 00:00:00 2001 From: Damian Barabonkov Date: Sat, 25 Jul 2026 17:01:23 -0700 Subject: [PATCH 1/2] Code --- packages/prime-sandboxes/README.md | 13 ++- .../src/prime_sandboxes/__init__.py | 2 + .../src/prime_sandboxes/models.py | 32 +++++++- packages/prime-sandboxes/tests/test_models.py | 39 +++++++++ packages/prime/README.md | 3 + .../prime/src/prime_cli/commands/sandbox.py | 60 ++++++++++++-- packages/prime/tests/test_sandbox_cli.py | 79 +++++++++++++++++++ 7 files changed, 220 insertions(+), 8 deletions(-) diff --git a/packages/prime-sandboxes/README.md b/packages/prime-sandboxes/README.md index fcc96b29..de21225a 100644 --- a/packages/prime-sandboxes/README.md +++ b/packages/prime-sandboxes/README.md @@ -25,7 +25,7 @@ pip install prime-sandboxes ## Quick Start ```python -from prime_sandboxes import APIClient, SandboxClient, CreateSandboxRequest +from prime_sandboxes import APIClient, SandboxClient, CreateSandboxRequest, StartCommand # Initialize client = APIClient(api_key="your-api-key") @@ -42,6 +42,17 @@ request = CreateSandboxRequest( sandbox = sandbox_client.create(request) print(f"Created: {sandbox.id}") +# VM workloads use a structured argv contract; no shell is implied. +vm = sandbox_client.create(CreateSandboxRequest( + name="vm-workload", + docker_image="user-1/vm-image:latest", + vm=True, + start_command=StartCommand( + executable="/worker", + args=["--platform", "linux/amd64"], + ), +)) + # Wait for it to be ready sandbox_client.wait_for_creation(sandbox.id) diff --git a/packages/prime-sandboxes/src/prime_sandboxes/__init__.py b/packages/prime-sandboxes/src/prime_sandboxes/__init__.py index 91479b24..01899b34 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/__init__.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/__init__.py @@ -60,6 +60,7 @@ SandboxListResponse, SandboxStatus, SSHSession, + StartCommand, TeamImageOwner, TransferImageResult, UpdateImagesRequest, @@ -91,6 +92,7 @@ "SandboxStatus", "SandboxListResponse", "CreateSandboxRequest", + "StartCommand", "UpdateSandboxRequest", "CommandRequest", "CommandResponse", diff --git a/packages/prime-sandboxes/src/prime_sandboxes/models.py b/packages/prime-sandboxes/src/prime_sandboxes/models.py index e4827c57..f5809871 100644 --- a/packages/prime-sandboxes/src/prime_sandboxes/models.py +++ b/packages/prime-sandboxes/src/prime_sandboxes/models.py @@ -121,13 +121,30 @@ class AdvancedConfigs(BaseModel): model_config = ConfigDict(extra="allow") +class StartCommand(BaseModel): + """A process to start without invoking a shell.""" + + executable: str = Field(min_length=1) + args: List[str] = Field(default_factory=list) + + model_config = ConfigDict(extra="forbid") + + @model_validator(mode="after") + def validate_no_nul_bytes(self) -> "StartCommand": + if "\x00" in self.executable: + raise ValueError("executable must not contain NUL bytes") + if any("\x00" in arg for arg in self.args): + raise ValueError("args must not contain NUL bytes") + return self + + class Sandbox(BaseModel): """Sandbox model""" id: str name: str docker_image: str = Field(..., alias="dockerImage") - start_command: Optional[str] = Field(None, alias="startCommand") + start_command: Optional[Union[StartCommand, str]] = Field(None, alias="startCommand") cpu_cores: float = Field(..., alias="cpuCores") memory_gb: float = Field(..., alias="memoryGB") disk_size_gb: float = Field(..., alias="diskSizeGB") @@ -179,7 +196,7 @@ class CreateSandboxRequest(BaseModel): name: str docker_image: str - start_command: Optional[str] = "tail -f /dev/null" + start_command: Optional[Union[StartCommand, str]] = "tail -f /dev/null" cpu_cores: float = 1.0 memory_gb: float = 1.0 disk_size_gb: float = 5.0 @@ -216,6 +233,17 @@ def validate_guaranteed(self) -> "CreateSandboxRequest": raise ValueError("guaranteed is not supported for VM sandboxes") return self + @model_validator(mode="after") + def validate_vm_start_command(self) -> "CreateSandboxRequest": + if self.vm and "start_command" not in self.model_fields_set: + self.start_command = None + return self + if self.vm and isinstance(self.start_command, str): + raise ValueError( + "VM sandboxes require start_command as StartCommand(executable=..., args=[...])" + ) + return self + @model_validator(mode="after") def validate_network_lists(self) -> "CreateSandboxRequest": if not self.vm and ( diff --git a/packages/prime-sandboxes/tests/test_models.py b/packages/prime-sandboxes/tests/test_models.py index 42b82b36..53ec1784 100644 --- a/packages/prime-sandboxes/tests/test_models.py +++ b/packages/prime-sandboxes/tests/test_models.py @@ -7,6 +7,7 @@ CreateSandboxRequest, Sandbox, SandboxStatus, + StartCommand, ) @@ -28,6 +29,44 @@ def test_create_sandbox_request_defaults(): assert request.timeout_minutes == 60 assert request.region is None assert request.labels == [] + assert request.start_command == "tail -f /dev/null" + + +def test_vm_start_command_preserves_argv(): + request = CreateSandboxRequest( + name="vm-workload", + docker_image="team/image:v1", + vm=True, + start_command=StartCommand( + executable="/worker", + args=["--platform", "linux/amd64", "value with spaces"], + ), + ) + + assert request.model_dump(exclude_none=True)["start_command"] == { + "executable": "/worker", + "args": ["--platform", "linux/amd64", "value with spaces"], + } + + +def test_vm_without_start_command_does_not_inherit_container_default(): + request = CreateSandboxRequest( + name="interactive-vm", + docker_image="team/image:v1", + vm=True, + ) + + assert request.start_command is None + + +def test_vm_rejects_legacy_string_start_command(): + with pytest.raises(ValidationError): + CreateSandboxRequest( + name="vm-workload", + docker_image="team/image:v1", + vm=True, + start_command="/worker --platform linux/amd64", + ) def test_create_sandbox_request_accepts_region(): diff --git a/packages/prime/README.md b/packages/prime/README.md index 470f85f6..6d7f5b71 100644 --- a/packages/prime/README.md +++ b/packages/prime/README.md @@ -192,6 +192,9 @@ prime sandbox create user-1/vm-image:latest --vm --gpu-count 1 --gpu-type H100_8 # Create a CPU-only VM sandbox prime sandbox create user-1/vm-image:latest --vm +# Create a one-shot VM workload (arguments after -- are preserved exactly) +prime sandbox create user-1/vm-image:latest --vm -- /worker --platform linux/amd64 + # List sandboxes prime sandbox list diff --git a/packages/prime/src/prime_cli/commands/sandbox.py b/packages/prime/src/prime_cli/commands/sandbox.py index b9928f1d..5fbbaa60 100644 --- a/packages/prime/src/prime_cli/commands/sandbox.py +++ b/packages/prime/src/prime_cli/commands/sandbox.py @@ -24,6 +24,7 @@ Sandbox, SandboxClient, SandboxNotRunningError, + StartCommand, UnauthorizedError, ) from rich.markup import escape @@ -160,6 +161,18 @@ def _network_access_description( return "Unrestricted" +def _start_command_data(value: Any) -> Any: + if isinstance(value, StartCommand): + return value.model_dump() + return value + + +def _format_start_command(value: Any) -> str: + if isinstance(value, StartCommand): + return json.dumps([value.executable, *value.args]) + return str(value) if value else "N/A" + + def _format_sandbox_for_details(sandbox: Sandbox) -> Dict[str, Any]: """Format sandbox data for details display (both table and JSON)""" data: Dict[str, Any] = { @@ -167,7 +180,7 @@ def _format_sandbox_for_details(sandbox: Sandbox) -> Dict[str, Any]: "name": sandbox.name, "type": "VM" if sandbox.vm else "Container", "docker_image": sandbox.docker_image, - "start_command": sandbox.start_command, + "start_command": _start_command_data(sandbox.start_command), "status": sandbox.status, "cpu_cores": sandbox.cpu_cores, "memory_gb": sandbox.memory_gb, @@ -413,7 +426,7 @@ def get( table.add_row("Name", sandbox_data["name"]) table.add_row("Type", sandbox_data["type"]) table.add_row("Docker Image", sandbox_data["docker_image"]) - table.add_row("Start Command", sandbox_data["start_command"] or "N/A") + table.add_row("Start Command", _format_start_command(sandbox.start_command)) sandbox_status_color = status_color(sandbox_data["status"], SANDBOX_STATUS_COLORS) table.add_row("Status", Text(sandbox_data["status"], style=sandbox_status_color)) @@ -496,11 +509,21 @@ def create( None, help="Image to run. When using --vm, provide the VM image reference.", ), + command: Optional[List[str]] = typer.Argument( + None, + help=( + "Executable and arguments to start. Use '--' before the executable " + "when its arguments begin with '-'." + ), + metavar="[COMMAND]...", + ), name: Optional[str] = typer.Option( None, help="Name for the sandbox (auto-generated if not provided)" ), start_command: Optional[str] = typer.Option( - "tail -f /dev/null", help="Command to run in the container" + None, + "--start-command", + help="Legacy container-only command string", ), cpu_cores: float = typer.Option(1.0, help="Number of CPU cores"), memory_gb: float = typer.Option(1.0, help="Memory in GB"), @@ -691,10 +714,37 @@ def create( console.print("[red]Error:[/red] --network-allow/--network-deny require --vm") raise typer.Exit(1) + if command and start_command is not None: + console.print( + "[red]Error:[/red] provide either trailing COMMAND arguments " + "or --start-command, not both" + ) + raise typer.Exit(1) + if vm and start_command is not None: + console.print( + "[red]Error:[/red] --start-command is legacy container-only syntax. " + "For VMs, pass an executable after '--'." + ) + raise typer.Exit(1) + + resolved_start_command: StartCommand | str | None + if command: + resolved_start_command = StartCommand( + executable=command[0], + args=command[1:], + ) + elif start_command is not None: + resolved_start_command = start_command + elif vm: + resolved_start_command = None + else: + # Preserve the existing long-running default for container callers. + resolved_start_command = "tail -f /dev/null" + request = CreateSandboxRequest( name=name, docker_image=docker_image, - start_command=start_command, + start_command=resolved_start_command, cpu_cores=cpu_cores, memory_gb=memory_gb, disk_size_gb=disk_size_gb, @@ -718,7 +768,7 @@ def create( console.print("\n[bold]Sandbox Configuration:[/bold]") console.print(f"Name: {name}") console.print(f"Docker Image: {docker_image}") - console.print(f"Start Command: {start_command or 'N/A'}") + console.print(f"Start Command: {_format_start_command(resolved_start_command)}") console.print(f"Resources: {cpu_cores} CPU, {memory_gb}GB RAM, {disk_size_gb}GB disk") if guaranteed: console.print("Scheduling: [green]Guaranteed QoS[/green]") diff --git a/packages/prime/tests/test_sandbox_cli.py b/packages/prime/tests/test_sandbox_cli.py index 3282b5df..ac2f7a3c 100644 --- a/packages/prime/tests/test_sandbox_cli.py +++ b/packages/prime/tests/test_sandbox_cli.py @@ -347,6 +347,85 @@ def mock_create(self: Any, request: Any) -> Any: assert captured["request"].vm is True +def test_sandbox_create_vm_start_command_preserves_argv( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _configure_cli(monkeypatch) + captured: dict[str, Any] = {} + + def mock_create(self: Any, request: Any) -> Any: + captured["request"] = request + return SimpleNamespace(id="sbx-vm-command") + + monkeypatch.setattr("prime_cli.commands.sandbox.SandboxClient.create", mock_create) + + result = runner.invoke( + app, + [ + "sandbox", + "create", + "team-1/worker:v1", + "--vm", + "--yes", + "--", + "/worker", + "--platform", + "linux/amd64", + "value with spaces", + ], + ) + + assert result.exit_code == 0, result.output + command = captured["request"].start_command + assert command.executable == "/worker" + assert command.args == ["--platform", "linux/amd64", "value with spaces"] + assert '["/worker", "--platform", "linux/amd64", "value with spaces"]' in strip_ansi( + result.output + ) + + +def test_sandbox_create_container_keeps_legacy_default( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _configure_cli(monkeypatch) + captured: dict[str, Any] = {} + + def mock_create(self: Any, request: Any) -> Any: + captured["request"] = request + return SimpleNamespace(id="sbx-container") + + monkeypatch.setattr("prime_cli.commands.sandbox.SandboxClient.create", mock_create) + + result = runner.invoke( + app, + ["sandbox", "create", "python:3.12", "--yes"], + ) + + assert result.exit_code == 0, result.output + assert captured["request"].start_command == "tail -f /dev/null" + + +def test_sandbox_create_container_keeps_explicit_empty_legacy_command( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _configure_cli(monkeypatch) + captured: dict[str, Any] = {} + + def mock_create(self: Any, request: Any) -> Any: + captured["request"] = request + return SimpleNamespace(id="sbx-container") + + monkeypatch.setattr("prime_cli.commands.sandbox.SandboxClient.create", mock_create) + + result = runner.invoke( + app, + ["sandbox", "create", "python:3.12", "--start-command", "", "--yes"], + ) + + assert result.exit_code == 0, result.output + assert captured["request"].start_command == "" + + def test_sandbox_create_gpu_without_docker_image(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("PRIME_API_KEY", "dummy") monkeypatch.setenv("PRIME_DISABLE_VERSION_CHECK", "1") From b35c0cdcc831d1879c138d24ae649d6e2411173c Mon Sep 17 00:00:00 2001 From: Damian Barabonkov Date: Sat, 25 Jul 2026 17:16:27 -0700 Subject: [PATCH 2/2] Code --- .github/workflows/ci.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 54d38953..da3b510a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -72,6 +72,8 @@ jobs: - name: Run Ruff uses: chartboost/ruff-action@v1 + with: + args: check --no-preview lock-check: # Advisory check: warns (does not fail) if uv.lock has drifted from