1
0
mirror of https://github.com/newnius/YAO-agent.git synced 2025-12-15 13:56:44 +00:00

add support for gpu specify

This commit is contained in:
2019-04-12 16:54:33 +08:00
parent 5c90e5a94f
commit 21b1d2edd7
3 changed files with 10 additions and 6 deletions

View File

@@ -75,14 +75,18 @@ class MyHandler(BaseHTTPRequestHandler):
'CONTENT_TYPE': self.headers['Content-Type'],
})
docker_image = form["image"].value
docker_name = form["name"].value
docker_cmd = form["cmd"].value
docker_workspace = form["workspace"].value
docker_gpus = form["gpus"].value
try:
client = docker.from_env()
container = client.containers.run(
image=docker_image,
hostname=docker_name,
command=docker_cmd,
environment={"key": "value"},
environment={"repo": docker_workspace, "NVIDIA_VISIBLE_DEVICES": docker_gpus},
runtime="nvidia",
detach=True
)