diff --git a/src/sagemaker/local/image.py b/src/sagemaker/local/image.py index 7893ee9260..39c879ef6d 100644 --- a/src/sagemaker/local/image.py +++ b/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": diff --git a/tests/unit/sagemaker/local/test_local_image.py b/tests/unit/sagemaker/local/test_local_image.py index ebca91a9f9..08c55fa0b4 100644 --- a/tests/unit/sagemaker/local/test_local_image.py +++ b/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"]}]}} }