Hackillinois

nvdocker/nvdocker.py

@@ -1,12 +1,11 @@
-
-import os 
+import os
 from subprocess import check_output
 import re
 import docker
 from py3nvml.py3nvml import *
 
-class NVDockerClient:
 
+class NVDockerClient:
     nvml_initialized = False
 
     def __init__(self):
@@ -16,18 +15,19 @@ def __init__(self):
     """
     Private method to check if nvml is loaded (and load the library if it isn't loaded)
     """
+
     def __check_nvml_init():
         if not NVDockerClient.nvml_initialized:
             nvmlInit()
             print("NVIDIA Driver Version:", nvmlSystemGetDriverVersion())
             NVDockerClient.nvml_initialized = True
 
-    #TODO: Testing on MultiGPU
+    # TODO: Testing on MultiGPU
     def create_container(self, image, **kwargs):
-        #defaults
+        # defaults
         config = {
-            "auto_remove":False,
-            "detach":True
+            "auto_remove": False,
+            "detach": True
         }
         environment = {}
         for arg in kwargs:
@@ -57,29 +57,29 @@ def create_container(self, image, **kwargs):
                 if "arch" in kwargs["require"]:
                     environment["NVIDIA_REQUIRE_ARCH"] = kwargs["require"]["arch"]
             elif arg == "cuda_version":
-                print("WARNING: the CUDA_VERSION enviorment variable is a legacy variable, consider moving to NVIDIA_REQUIRE_CUDA")
+                print(
+                    "WARNING: the CUDA_VERSION enviorment variable is a legacy variable, consider moving to NVIDIA_REQUIRE_CUDA")
                 environment["CUDA_VERSION"] = kwargs["cuda_version"]
             elif arg == "environment":
                 if type(kwargs["environment"]) is dict:
-                    for k,v in kwargs["environment"]:
+                    for k, v in kwargs["environment"]:
                         environment[k] = v
                 elif type(kwargs["environment"]) is list:
                     for e in kwargs["environment"]:
                         kv = e.split("=")
-                        assert(len(kv) == 2), "Does not follow the format SOMEVAR=xxx"
+                        assert (len(kv) == 2), "Does not follow the format SOMEVAR=xxx"
                         environment[kv[0]] = kv[1]
             else:
                 config[arg] = kwargs[arg]
         config["environment"] = environment
         config["runtime"] = "nvidia"
-        
+
         c = self.docker_client.containers.run(image, "", **config)
 
         return c
 
-
     def run(self, image, cmd="", **kwargs):
-        #defaults
+        # defaults
         config = {}
         environment = {}
         for arg in kwargs:
@@ -109,16 +109,17 @@ def run(self, image, cmd="", **kwargs):
                 if "arch" in kwargs["require"]:
                     environment["NVIDIA_REQUIRE_ARCH"] = kwargs["require"]["arch"]
             elif arg == "cuda_version":
-                print("WARNING: the CUDA_VERSION enviorment variable is a legacy variable, consider moving to NVIDIA_REQUIRE_CUDA")
+                print(
+                    "WARNING: the CUDA_VERSION enviorment variable is a legacy variable, consider moving to NVIDIA_REQUIRE_CUDA")
                 environment["CUDA_VERSION"] = kwargs["cuda_version"]
             elif arg == "environment":
                 if type(kwargs["environment"]) is dict:
-                    for k,v in kwargs["environment"]:
+                    for k, v in kwargs["environment"]:
                         environment[k] = v
                 elif type(kwargs["environment"]) is list:
                     for e in kwargs["environment"]:
                         kv = e.split("=")
-                        assert(len(kv) == 2), "Does not follow the format SOMEVAR=xxx"
+                        assert (len(kv) == 2), "Does not follow the format SOMEVAR=xxx"
                         environment[kv[0]] = kv[1]
             else:
                 config[arg] = kwargs[arg]
@@ -135,34 +136,43 @@ def run(self, image, cmd="", **kwargs):
     def build_image(self, path):
         img = self.docker_client.images.build(path);
         return img
-        
+
     def get_container_logs(self, cid):
         c = self.docker_client.containers.get(cid)
         return c.logs()
 
     def get_all_container_ids(self):
         return self.docker_client.containers.list()
-    
+
     def stop_container(self, cid):
         c = self.docker_client.containers.get(cid)
         c.stop()
 
     def start_container(self, cid):
         c = self.docker_client.containers.get(cid)
         c.start()
-    
+
     def start_all_containers(self):
         for c in self.docker_client.containers.list():
             c.start()
-        
-    def stop_all_containers(self):    
+
+    def stop_all_containers(self):
         for c in self.docker_client.containers.list():
             c.stop()
 
     def exec_run(self, cid, cmd):
         c = self.docker_client.containers.get(cid)
         return c.exec_run(cmd)
 
+    def docker_image_pull(self, repo, tag):
+        return self.docker_client.images.pull(repo, tag)
+
+    def docker_image_list(self, **kwargs):
+        return self.docker_client.images.list(kwargs)
+
+    def docker_image_search(self, image):
+        return self.docker_client.images.search(image)
+
     @staticmethod
     def gpu_info():
         NVDockerClient.__check_nvml_init()
@@ -182,9 +192,9 @@ def gpu_memory_usage(id):
         gpu_handle = gpus[id]["gpu_handle"]
         gpu_memory_data = nvmlDeviceGetMemoryInfo(gpu_handle)
         rv = {}
-        #returns in megabytes
-        rv["used_mb"] = gpu_memory_data.used/1e6
-        rv["free_mb"] = gpu_memory_data.free/1e6
+        # returns in megabytes
+        rv["used_mb"] = gpu_memory_data.used / 1e6
+        rv["free_mb"] = gpu_memory_data.free / 1e6
         return rv
 
     @staticmethod