more udocker fixes

ratt-ru · Jun 18, 2019 · 049857c · 049857c
1 parent 4c4f0d4
commit 049857c
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 16 deletions.
diff --git a/examples/simulation_pipeline.py b/examples/simulation_pipeline.py
@@ -7,7 +7,11 @@
 OUTPUT = "output"
 MSDIR = "msdir"
 PREFIX = "stimela-example"  # Prefix for output images
-SINGULARTITY_IMAGE_DIR = os.environ["STIMELA_SINGULARTITY_IMAGE_DIR"]
+try:
+    SINGULARTITY_IMAGE_DIR = os.environ["STIMELA_SINGULARTITY_IMAGE_DIR"]
+except KeyError:
+    pass
+
 
 # MS name
 MS = "meerkat_simulation_example.ms"
@@ -19,7 +23,7 @@
 # Start stimela Recipe instance
 pipeline = stimela.Recipe("Simulation Example",     # Recipe name
                   ms_dir=MSDIR,
-                  singularity_image_dir=SINGULARTITY_IMAGE_DIR,
+#                  singularity_image_dir=SINGULARTITY_IMAGE_DIR,
                   )
 
 pipeline.JOB_TYPE = "udocker"

diff --git a/stimela/udocker.py b/stimela/udocker.py
@@ -32,7 +32,8 @@ def __init__(self, image, name,
                  shared_memory="1gb",
                  time_out=-1,
                  log_container=None,
-                 COMMAND=""):
+                 COMMAND="",
+                 use_graphics=False):
         """
         Python wrapper to docker engine tools for managing containers.
         """
@@ -49,6 +50,7 @@ def __init__(self, image, name,
         self.PID = os.getpid()
         self.uptime = "00:00:00"
         self.time_out = time_out
+        self.use_graphics = use_graphics
         self.cont_logger = utils.logger.StimelaLogger(log_container or stimela.LOG_FILE)
 
 
@@ -85,7 +87,9 @@ def run(self, *args):
         tstart = time.time()
         utils.xrun("udocker run", ["=".join(args)] + [volumes, environs,
                         "--workdir=%s"%(self.WORKDIR) if self.WORKDIR else "",
-                        "--rm", self.name, self.COMMAND or ""], timeout=self.time_out)
+                        "--rm",
+                        "--dri" if self.use_graphics else "",
+                        self.name, self.COMMAND or ""], timeout=self.time_out)
 
         self.status = "running"
         uptime = seconds_hms(time.time() - tstart)

diff --git a/stimela/utils/logger.py b/stimela/utils/logger.py
@@ -9,7 +9,7 @@
 from datetime import datetime
 
 class StimelaLogger(object):
-    def __init__(self, lfile):
+    def __init__(self, lfile, jtype="udocker"):
 
         self.lfile = lfile
         # Create file if it does not exist
@@ -28,26 +28,44 @@ def __init__(self, lfile):
         if changed:
             self.write()
 
+
+        self.jtype = jtype
+
     def _inspect(self, name):
-        output = subprocess.check_output("docker inspect {}".format(name), shell=True).decode()
-        output_file = StringIO(output[3:-3])
+
+        output = subprocess.check_output("{0:s} inspect {1:s}".format(self.jtype, name), shell=True).decode()
+        if self.jtype == "docker":
+            output_file = StringIO(output[3:-3])
+        else:
+            output_file = StringIO(output)
         jdict = yaml.safe_load(output_file)
         output_file.close()
 
         return jdict
 
     def log_image(self, name, image_dir, replace=False, cab=False):
         info = self._inspect(name)
-
-        if name not in self.info['images'].keys() or replace:
-            self.info['images'][name] = {
-                'TIME'      :   info['Created'].split('.')[0].replace('Z', '0'),
-                'ID'        :   info['Id'].split(':')[1],
-                'CAB'       :   cab,
-                'DIR'       :   image_dir, 
-            }
+        if self.jtype == "docker":
+            if name not in self.info['images'].keys() or replace:
+                self.info['images'][name] = {
+                    'TIME'      :   info['Created'].split('.')[0].replace('Z', '0'),
+                    'ID'        :   info['Id'].split(':')[1],
+                    'CAB'       :   cab,
+                    'DIR'       :   image_dir, 
+                }
+            else:
+                print('Image {0} has already been logged.'.format(name))
         else:
-            print('Image {0} has already been logged.'.format(name))
+            if name not in self.info['images'].keys() or replace:
+                self.info['images'][name] = {
+                    'TIME'      :   info['created'].split('.')[0].replace('Z', '0'),
+                    'ID'        :   info['id'],
+                    'CAB'       :   cab,
+                    'DIR'       :   image_dir, 
+                }
+            else:
+                print('Image {0} has already been logged.'.format(name))
+
 
     def log_container(self, name):
         info = self._inspect(name)