nvidia.py 文件源码

def add_nvidia_docker_to_config(container_config):
    if not container_config.get('HostConfig', None):
        container_config['HostConfig'] = {}

    nvidia_config = get_nvidia_configuration()

    # Setup the Volumes
    container_config['HostConfig'].setdefault('VolumeDriver', nvidia_config['VolumeDriver'])
    container_config['HostConfig'].setdefault('Binds', [])
    container_config['HostConfig']['Binds'].extend(nvidia_config['Volumes'])

    # Get nvidia control devices
    devices = container_config['HostConfig'].get('Devices', [])
    # suport both '0 1' and '0, 1' formats, just like nvidia-docker
    gpu_isolation = os.getenv('NV_GPU', '').replace(',', ' ').split()
    pattern = re.compile(r'/nvidia([0-9]+)$')
    for device in nvidia_config['Devices']:
        if gpu_isolation:
            card_number = pattern.search(device)
            if card_number and card_number.group(1) not in gpu_isolation:
                continue
        devices.extend(parse_devices([device]))

    container_config['HostConfig']['Devices'] = devices