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))