41 lines
882 B
Python
41 lines
882 B
Python
import docker
|
|
import os
|
|
|
|
|
|
if not os.path.isdir('./dataset_dir'):
|
|
os.mkdir('./dataset_dir')
|
|
|
|
docker_client = docker.from_env()
|
|
container = docker_client.containers.run(
|
|
'snsd0805/cifar100-train:v3',
|
|
volumes={
|
|
'dataset_dir': {'bind': '/dataset', 'mode': 'rw'},
|
|
'output': {'bind': '/output', 'mode': 'rw'},
|
|
},
|
|
network='train-net',
|
|
runtime='nvidia',
|
|
device_requests=[
|
|
docker.types.DeviceRequest(count=-1, capabilities=[['gpu']])
|
|
],
|
|
name='train-0',
|
|
environment={
|
|
'GPU_NUM': 1,
|
|
'NODE_NUM': 1,
|
|
'NODE_RANK': 0,
|
|
'MASTER_IP': 'train-0',
|
|
'MASTER_PORT': 21046,
|
|
},
|
|
detach=True
|
|
)
|
|
|
|
|
|
print(container.short_id)
|
|
for line in container.logs(stream=True):
|
|
print(line.strip().decode())
|
|
|
|
result = container.wait()
|
|
status_code = result['StatusCode']
|
|
print(status_code, type(status_code))
|
|
|
|
|