fix: make sure gpus are found in local_gpu run (aws#4384)

* fix: make sure gpus are found in local_gpu run

* fix: black formatting

* fix: adjust unit test

Author: gverkes

src/sagemaker/local/image.py

@@ -860,7 +860,9 @@ def _create_docker_host(
         # to setting --runtime=nvidia in the docker commandline.
         if self.instance_type == "local_gpu":
             host_config["deploy"] = {
-                "resources": {"reservations": {"devices": [{"capabilities": ["gpu"]}]}}
+                "resources": {
+                    "reservations": {"devices": [{"count": "all", "capabilities": ["gpu"]}]}
+                }
             }
 
         if not self.is_studio and command == "serve":

tests/unit/sagemaker/local/test_local_image.py

@@ -871,7 +871,7 @@ def test_container_has_gpu_support(tmpdir, sagemaker_session):
     docker_host = sagemaker_container._create_docker_host("host-1", {}, set(), "train", [])
     assert "deploy" in docker_host
     assert docker_host["deploy"] == {
-        "resources": {"reservations": {"devices": [{"capabilities": ["gpu"]}]}}
+        "resources": {"reservations": {"devices": [{"count": "all", "capabilities": ["gpu"]}]}}
     }