Custom docker

gpu_cluster/controllers/cpu_container_controller.py

@@ -5,23 +5,43 @@
 from .container_controller import ContainerController
 import docker
 
+
 class CPUContainerController(ContainerController):
 
     def __init__(self, config):
         super().__init__(config)
         self.client = docker.from_env(version='auto')
-    
+
     def create_container(self, image, user="", token_required=False, budget=-1):
-        uport = self.get_port()
-        mport = self.get_port()
+        uport = super().get_port()
+        mport = super().get_port()
         while uport == mport:
-            mport = self.get_port()
+            mport = super().get_port()
 
-        ports = {'8888/tcp':uport,
-                 '6006/tcp':mport}
+        ports = {'8888/tcp': uport,
+                 '6006/tcp': mport}
 
         print(image)
-        c_id = self.client.containers.run(image, "", auto_remove=True, detach=True, ports=ports).id
+        container_list = self.client.containers.list(filters={'name': image})
+        if container_list:
+            c_id = self.client.containers.run(image, "", auto_remove=True, detach=True, ports=ports).id
+
+        else:
+            # Add a client.images.search to check if the path to the container exists on docker hub. If not, error out
+            has_result = self.client.images.search(image)
+            if not has_result:
+                print('No image in DockerHub')
+                return 'No image in DockerHub' , '', ''
+
+            image_tag = image.split(':')
+            docker_image = self.client.images.pull(image_tag[0], image_tag[1])
+
+            # If pull returns more than one image, get the first one in the list
+            if hasattr(docker_image, '__len__'):
+                docker_image = docker_image[0]
+
+            # Do you have to build the image after you pull it from Docker Hub?
+            c_id = self.client.containers.run(docker_image, '', auto_remove=True, detach=True, ports=ports).id
         print(c_id)
 
         uurl = ""
@@ -36,13 +56,11 @@ def create_container(self, image, user="", token_required=False, budget=-1):
             uurl = base_url + str(uport)
             murl = base_url + str(mport)
         print(image)
-
-        #TODO insert budget
-        db_session.add(Instance(c_id, uport, mport, uurl, murl, user, budget, token))   
+        # TODO insert budget
+        db_session.add(Instance(c_id, uport, mport, uurl, murl, user, budget, token))
         db_session.commit()
         return c_id, uurl, murl
 
     def kill_container(self, c_id):
         c = self.client.containers.get(c_id)
-        c.stop()
-
+        c.stop()

gpu_cluster/controllers/gpu_container_controller.py

@@ -3,17 +3,14 @@
 from .container_controller import ContainerController
 from nvdocker import NVDockerClient
 
+
 class GPUContainerController(ContainerController):
 
     def __init__(self, config):
         super().__init__(config)
         self.docker_client = NVDockerClient()
-    
+
     def create_container(self, image, user="", token_required=False, budget=-1, num_gpus=1):
-        if NVDockerClient.least_used_gpu() == None :
-            #TODO Add handle multi node functionality here
-            pass
-
         # Get 2 open ports for UI and Monitor
         uport = super().get_port()
         mport = super().get_port()
@@ -29,8 +26,8 @@ def create_container(self, image, user="", token_required=False, budget=-1, num_
         for g in range(num_gpus):
             if NVDockerClient.gpu_memory_usage(g)["free_mb"] > 0:
                 gpus.append(g)
-                
-        # Assemble config for container 
+
+        # Assemble config for container
         container_config = {
             "ports": {
                 '8888/tcp': uport,
@@ -42,10 +39,31 @@ def create_container(self, image, user="", token_required=False, budget=-1, num_
             "auto_remove": True
         }
 
-        #create container
-        c_id = self.docker_client.create_container(image, **container_config).id
+        # create container
+        container_list = self.docker_client.docker_image_list(filters={'name': image})
+        print(image)
+        if container_list:
+            c_id = self.docker_client.create_container(image, **container_config).id
+
+        else:
+            # Add a client.images.search to check if the path to the container exists on docker hub. If not, error out
+            has_result = self.docker_client.docker_image_search(image)
+            if not has_result:
+                print('No image in DockerHub')
+                return 'No image in DockerHub' , '', ''
+
+            image_tag = image.split(':')
+            docker_image = self.docker_client.docker_image_pull(image_tag[0], image_tag[1])
+
+            # If pull returns more than one image, get the first one in the list
+            if hasattr(docker_image, '__len__'):
+                docker_image = docker_image[0]
+            print(docker_image)
+
+            # Do you have to build the image after you pull it from Docker Hub?
+            c_id = self.docker_client.create_container(docker_image, **container_config).id
 
-        #assemble endpoints for UI, monitor and get the access token if needed
+        # assemble endpoints for UI, monitor and get the access token if needed
         uurl = ""
         murl = ""
         token = ""

gpu_cluster/routes/cluster_api.py

@@ -21,6 +21,9 @@ def create_container(self):
                 abort(400)
 
         cid, ui_url, murl = self.controller.create_container(request.json['image'],  token_required=request.json['token_required'])#, user=request.json['user'], budget=request.json['budget'] )
+        if ui_url == '' or murl == '':
+            abort(400)
+
         return jsonify({'cid': cid, 'ui_url' : ui_url, 'monitor_url': murl})
 
     def confirm_launch(self):