fix: adjust unit test

gverkes · akrishna1995 · commit bbe53dc807e2 · 2024-02-29T15:58:05.000-08:00
diff --git 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"]}]}}
     }
 
 

Original file line number	Diff line number	Diff line change
`@@ -871,7 +871,7 @@ def test_container_has_gpu_support(tmpdir, sagemaker_session):`
`871`	`871`	`docker_host = sagemaker_container._create_docker_host("host-1", {}, set(), "train", [])`
`872`	`872`	`assert "deploy" in docker_host`
`873`	`873`	`assert docker_host["deploy"] == {`
`874`		`- "resources": {"reservations": {"devices": [{"capabilities": ["gpu"]}]}}`
	`874`	`+ "resources": {"reservations": {"devices": [{"count": "all", "capabilities": ["gpu"]}]}}`
`875`	`875`	`}`
`876`	`876`
`877`	`877`