如何在Python脚本中集成Nvidia-Docker API实现GPU映射？

Hey there! I’ve been in your exact situation before—trying to integrate NVIDIA GPU support into Python scripts that use the Docker API can feel a bit obscure since the docs don’t spell it out super clearly. Let me break down exactly how to make this work seamlessly, no special NVIDIA API required.

NVIDIA Docker doesn’t have a separate API; instead, it hooks into Docker’s native API through either runtime configuration or modern device requests. Here are the two reliable methods:

Method 1: Use the NVIDIA Runtime (For nvidia-docker2) #

If you’ve installed the nvidia-docker2 package (the older but still widely used setup), you just need to specify the nvidia runtime when creating your container, plus set a couple of environment variables to control GPU access.

Here’s how to modify your existing Python code:

import docker

# Initialize Docker client
client = docker.from_env()

# Build your container parameters with GPU support
container_config = {
    'image': 'your-nvidia-base-image:tag',  # e.g., nvidia/cuda:11.8.0-runtime-ubuntu22.04
    'command': 'your-workload-command',
    'detach': True,
    'runtime': 'nvidia',  # Critical: tells Docker to use the NVIDIA runtime
    'environment': [
        'NVIDIA_VISIBLE_DEVICES=all',  # Use all GPUs, or specify IDs like "0,1"
        'NVIDIA_DRIVER_CAPABILITIES=compute,utility'  # Enables core GPU features
    ],
    # Add your existing parameters here (ports, volumes, etc.)
}

# Create and start the container
container = client.containers.create(**container_config)
container.start()

Method 2: Use Device Requests (Docker 19.03+ Recommended) #

For newer Docker versions (19.03+), you don’t even need nvidia-docker2—just the nvidia-container-toolkit. This method uses Docker’s native device_requests API to request GPU resources directly, which is more aligned with Docker’s modern design.

Here’s the updated code for this approach:

import docker

client = docker.from_env()

container_config = {
    'image': 'your-nvidia-base-image:tag',
    'command': 'your-workload-command',
    'detach': True,
    'device_requests': [
        docker.types.DeviceRequest(
            count=-1,  # Request all available GPUs; set to a number (e.g., 1) for specific count
            capabilities=[['gpu']]
        )
    ],
    # Optional: Fine-tune GPU visibility with env vars
    'environment': [
        'NVIDIA_VISIBLE_DEVICES=all'
    ],
    # Add your existing parameters here
}

container = client.containers.create(**container_config)
container.start()

• First, verify your Docker environment is set up correctly: Run docker run --gpus all nvidia/cuda:11.8.0-runtime-ubuntu22.04 nvidia-smi manually. If this outputs GPU info, your setup is good.
• For batch container creation, adjust NVIDIA_VISIBLE_DEVICES in the loop to assign specific GPUs to each container (e.g., "0" for the first container, "1" for the second) to avoid resource conflicts.
• If you want all containers to use NVIDIA by default, you can set "default-runtime": "nvidia" in /etc/docker/daemon.json, then restart Docker—this lets you skip specifying runtime or device_requests in every API call.

内容的提问来源于stack exchange，提问作者C. Berger