snsd0805/gpu-provider
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: remove container automatically

Browse Source
This commit is contained in:
Ting-Jun Wang 2024-06-01 13:06:05 +08:00
parent 0438f9f7b1
commit bc476e398f
Signed by: snsd0805
GPG Key ID: 48D331A3D6160354
1 changed files with 11 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
src/communication.py
View File

@ -52,6 +52,8 @@ class ClusterCommunicationModule():
if self.node_manager.docker_client.swarm.attrs == {}:
print("Build new docker swarm...")
self.node_manager.docker_client.swarm.init(advertise_addr=self.host, listen_addr=f"{self.host}:2377", force_new_cluster=True)
print("Create new overlay network")
self.node_manager.docker_client.networks.create(name='train-net', driver='overlay', attachable=True)
# send docker swarm token to the worker
token = self.node_manager.docker_client.swarm.attrs['JoinTokens']['Worker']
@ -122,7 +124,7 @@ class ClusterCommunicationModule():
print(f"[RUN_CONTAINER] {image}")
self.run_container(image, train, args)
print(f"[RUN_CONTAINER] {image}")
print(f"[RUN_CONTAINER SUCCESS] {image}")
self.client_sock.send('[RUN_CONTAINER_SUCCESS] {}'.encode())
return True
@ -178,8 +180,10 @@ class ClusterCommunicationModule():
def exit(self):
if self.node_manager.status == 'master':
for conn in self.worker_conns:
try:
conn.send('[STOP] {}'.encode())
check, args = conn.recv(1024).decode().split(' ')
except:
print(f'{args} has stopped.')
self.node_manager.docker_client.swarm.leave(force=True)
@ -239,7 +243,9 @@ class ClusterCommunicationModule():
result = container.wait()
status_code = result['StatusCode']
print(status_code, type(status_code))
if status_code != 0:
print(f'[ERROR] some error occur in the docker container, error_code={status_code}')
container.remove()
def scatter_container(self, image_name, train):
def master_run(image_name):