No longer use runtime=nvidia to expose GPUs to the container.

Instead use `DeviceRequest`, which is equivalent to running with `--gpu all`.

PiperOrigin-RevId: 431933575
Change-Id: Iee9f70ae8f291d16850e4f58349559d9468ef4d8

docker/run_docker.py

@@ -227,10 +227,14 @@ def main(argv):
   ])
 
   client = docker.from_env()
+  device_requests = [
+      docker.types.DeviceRequest(driver='nvidia', capabilities=[['gpu']])
+  ] if FLAGS.use_gpu else None
+
   container = client.containers.run(
       image=FLAGS.docker_image_name,
       command=command_args,
-      runtime='nvidia' if FLAGS.use_gpu else None,
+      device_requests=device_requests,
       remove=True,
       detach=True,
       mounts=mounts,