Skip to main content

Goals

  • Understand when an environment needs GPU access
  • Configure machine_size to provision a GPU-enabled sandbox

Prerequisites

Introduction

Some tasks require GPU access - for example, evaluating CUDA kernels or training models inside the sandbox. By default, sandboxes use CPU-only machine sizes like "0.5:1". To give the agent access to a GPU, set machine_size to "nvidia-l4".

Configuration

Specify machine_size="nvidia-l4" in your SandboxSettings:
self.sandbox_settings = SandboxSettings(
    environment="YourUsername/YourEnv",
    image="nvidia/cuda:12.4.0-devel-ubuntu22.04",
    machine_size="nvidia-l4",
    block_network=False,
    bucket_config=SandboxBucketConfig(
        mount_path="/tmp/sandbox/",
        read_only=True,
    )
)
Your Docker image must include CUDA drivers and any GPU libraries your tasks depend on. The official nvidia/cuda images are a good starting - pick a tag that matches your CUDA version and build on top of it with your own dependencies. This gives the agent access to an NVIDIA L4 GPU inside the sandbox.
GPU environments have limited capacity and are more expensive than CPU-only environments.

Managing GPU cost with sandbox lifetime

Because GPU sandboxes are expensive, consider scoping the sandbox lifetime so the GPU is only allocated while it’s actually needed. For example, in a kernel-optimization environment the agent spends most of its time reasoning and writing code, it only needs the GPU when executing the kernel. You can design the sandbox so it is created at the start of the tool call that runs the kernel and destroyed when that call finishes, rather than keeping it alive for the entire session. See the Sandbox Lifecycle section for more on controlling when sandboxes are created and destroyed.