GPU support for DockerCommandLineCodeExecutor

[millerh1]

It seems like DockerCommandLineCodeExecutor doesn't use GPUs even when they are available. I think you would need to update _docker_code_executor.py to something like this:

# Create a device request asking for all GPUs (equivalent to --gpus all)
gpu_request = [DeviceRequest(count=-1, capabilities=[['gpu']])]

self._container = await asyncio.to_thread(
    client.containers.create,
    self._image,
    name=self.container_name,
    entrypoint=shell_command,
    command=command,
    tty=True,
    detach=True,
    auto_remove=self._auto_remove,
    volumes={str(self._bind_dir.resolve()): {"bind": "/workspace", "mode": "rw"}, **self._extra_volumes},
    working_dir="/workspace",
    extra_hosts=self._extra_hosts,
    # Pass the GPU request to the container creation call
    device_requests=gpu_request,
)
await asyncio.to_thread(self._container.start)

I tried this locally and it worked nicely. Is this of interest?

[ekzhu]

I think we will likely need to update parameter list to include device_requests.

Would you like to help with this by submitting a PR? Issue: #6302

[millerh1]

Sure! I'll take a look.