update

executor.py

@@ -8,8 +8,7 @@ from urllib import parse
 PORT_NUMBER = 8000
 
 
-# This class will handles any incoming request from
-# the browser
+# This class will handles any incoming request from the browser
 class MyHandler(BaseHTTPRequestHandler):
 	# Handler for the GET requests
 	def do_GET(self):
@@ -17,7 +16,6 @@ class MyHandler(BaseHTTPRequestHandler):
 		query = parse.parse_qs(req.query)
 
 		if req.path == "/ping":
-			# Open the static file requested and send it
 			self.send_response(200)
 			self.send_header('Content-type', 'application/json')
 			self.end_headers()
@@ -25,37 +23,42 @@ class MyHandler(BaseHTTPRequestHandler):
 
 		elif req.path == "/logs":
 			try:
-				container_id = query['id'][0]
+				container_id = query.get('id')[0]
 				client = docker.from_env()
 				container = client.containers.get(container_id)
 				msg = {'code': 0, 'logs': str(container.logs().decode())}
 			except Exception as e:
-				msg = {'code': 0, 'error': str(e)}
+				msg = {'code': 1, 'error': str(e)}
 			self.send_response(200)
 			self.send_header('Content-type', 'application/json')
 			self.end_headers()
 			self.wfile.write(bytes(json.dumps(msg), "utf-8"))
 
 		elif req.path == "/status":
-			container_id = query['id'][0]
-			client = docker.from_env()
-			container = client.containers.list(all=True, filters={'id': container_id})
-			if len(container) > 0:
-				container = container[0]
-				status = {
-					'id': container.short_id,
-					'image': container.attrs['Config']['Image'],
-					'image_digest': container.attrs['Image'],
-					'command': container.attrs['Config']['Cmd'],
-					'created_at': container.attrs['Created'],
-					'finished_at': container.attrs['State']['FinishedAt'],
-					'status': container.status
-				}
-				if status['command'] is not None:
-					status['command'] = ' '.join(container.attrs['Config']['Cmd'])
-				msg = {'code': 0, 'status': status}
-			else:
-				msg = {'code': 1, 'error': "container not exist"}
+			try:
+				container_id = query.get('id')[0]
+				client = docker.from_env()
+				container = client.containers.list(all=True, filters={'id': container_id})
+				if len(container) > 0:
+					container = container[0]
+					status = {
+						'id': container.short_id,
+						'image': container.attrs['Config']['Image'],
+						'image_digest': container.attrs['Image'],
+						'command': container.attrs['Config']['Cmd'],
+						'created_at': container.attrs['Created'],
+						'finished_at': container.attrs['State']['FinishedAt'],
+						'status': container.status,
+						'hostname': container.attrs['Config']['Hostname'],
+						'state': container.attrs['State']
+					}
+					if status['command'] is not None:
+						status['command'] = ' '.join(container.attrs['Config']['Cmd'])
+					msg = {'code': 0, 'status': status}
+				else:
+					msg = {'code': 1, 'error': "container not exist"}
+			except Exception as e:
+				msg = {'code': 2, 'error': str(e)}
 			self.send_response(200)
 			self.send_header('Content-type', 'application/json')
 			self.end_headers()
@@ -74,18 +77,25 @@ class MyHandler(BaseHTTPRequestHandler):
 					'REQUEST_METHOD': 'POST',
 					'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
+			docker_image = form.getvalue('image')
+			docker_name = form.getvalue('name')
+			docker_cmd = form.getvalue('cmd')
+			docker_workspace = form.getvalue('workspace')
+			docker_gpus = form.getvalue('gpus')
+			docker_mem_limit = form.getvalue('mem_limit')
+			docker_cpu_limit = form.getvalue('cpu_limit')
+			docker_network = form.getvalue('network')
 
 			try:
 				client = docker.from_env()
 				container = client.containers.run(
-					image=docker_image,
+					name=docker_name,
 					hostname=docker_name,
+					image=docker_image,
 					command=docker_cmd,
+					mem_limit=docker_mem_limit,
+					nano_cpus=docker_cpu_limit,
+					network=docker_network,
 					environment={"repo": docker_workspace, "NVIDIA_VISIBLE_DEVICES": docker_gpus},
 					runtime="nvidia",
 					detach=True
@@ -107,7 +117,7 @@ class MyHandler(BaseHTTPRequestHandler):
 					'REQUEST_METHOD': 'POST',
 					'CONTENT_TYPE': self.headers['Content-Type'],
 				})
-			container_id = form["id"].value
+			container_id = form.getvalue('id')
 
 			client = docker.from_env()
 			container = client.containers.get(container_id)
@@ -127,7 +137,7 @@ class MyHandler(BaseHTTPRequestHandler):
 					'REQUEST_METHOD': 'POST',
 					'CONTENT_TYPE': self.headers['Content-Type'],
 				})
-			container_id = form["id"].value
+			container_id = form.getvalue('id')
 
 			client = docker.from_env()
 			container = client.containers.get(container_id)