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:
@@ -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
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user