feature: add GPU capability to local mode

teticio · teticio · commit 5b887e7b0218 · 2022-05-25T17:27:31.000+01:00
diff --git a/src/sagemaker/local/image.py b/src/sagemaker/local/image.py
@@ -746,6 +746,7 @@ def _create_docker_host(self, host, environment, optml_subdirs, command, volumes
             "volumes": [v.map for v in optml_volumes],
             "environment": environment,
             "networks": {"sagemaker-local": {"aliases": [host]}},
+            "deploy": {"resources": {"reservations": {"devices": [{"capabilities": ["gpu"]}]}}},
         }
 
         if command != "process":

Original file line number	Diff line number	Diff line change
`@@ -746,6 +746,7 @@ def _create_docker_host(self, host, environment, optml_subdirs, command, volumes`
`746`	`746`	`"volumes": [v.map for v in optml_volumes],`
`747`	`747`	`"environment": environment,`
`748`	`748`	`"networks": {"sagemaker-local": {"aliases": [host]}},`
	`749`	`+ "deploy": {"resources": {"reservations": {"devices": [{"capabilities": ["gpu"]}]}}},`
`749`	`750`	`}`
`750`	`751`
`751`	`752`	`if command != "process":`