Merge pull request #26 from lukemartinlogan/master

Add real_kwargs. Improve sudo based on Ares

jarvis_util/introspect/system_info.py

@@ -775,6 +775,7 @@ def find_storage(self,
                      min_cap=None,
                      min_avail=None,
                      mount_res=None,
+                     shared=None,
                      df=None):
         """
         Find a set of storage devices.
@@ -790,6 +791,7 @@ def find_storage(self,
         :param min_cap: Remove devices with too little overall capacity
         :param min_avail: Remove devices with too little available space
         :param mount_res: A regex or list of regexes to match mount points
+        :param shared: Whether to search for devices which are shared
         :param df: The data frame to run this query
         :return: Dataframe
         """
@@ -829,6 +831,8 @@ def find_storage(self,
         if common and condense:
             df = df.groupby(['mount']).first().reset_index()
         #     df = df.drop_columns('host')
+        if shared is not None:
+            df = df[lambda r: r['shared'] == shared]
         return df
 
     @staticmethod

jarvis_util/jutil_manager.py

@@ -25,4 +25,5 @@ def __init__(self):
         self.hide_output = False
         self.debug_mpi_exec = False
         self.debug_local_exec = False
+        self.debug_scp = False
 

jarvis_util/shell/exec_info.py

@@ -6,6 +6,7 @@
 
 from enum import Enum
 from jarvis_util.util.hostfile import Hostfile
+from jarvis_util.jutil_manager import JutilManager
 import os
 from abc import ABC, abstractmethod
 
@@ -33,7 +34,7 @@ class ExecInfo:
     def __init__(self,  exec_type=ExecType.LOCAL, nprocs=None, ppn=None,
                  user=None, pkey=None, port=None,
                  hostfile=None, hosts=None, env=None,
-                 sleep_ms=0, sudo=False, cwd=None,
+                 sleep_ms=0, sudo=False, sudoenv=True, cwd=None,
                  collect_output=None, pipe_stdout=None, pipe_stderr=None,
                  hide_output=None, exec_async=False, stdin=None):
         """
@@ -49,6 +50,7 @@ def __init__(self,  exec_type=ExecType.LOCAL, nprocs=None, ppn=None,
         :param env: The environment variables to use for command.
         :param sleep_ms: Sleep for a period of time AFTER executing
         :param sudo: Execute command with root privilege. E.g., SSH, PSSH
+        :param sudoenv: Support environment preservation in sudo
         :param cwd: Set current working directory. E.g., SSH, PSSH
         :param collect_output: Collect program output in python buffer
         :param pipe_stdout: Pipe STDOUT into a file. (path string)
@@ -71,6 +73,7 @@ def __init__(self,  exec_type=ExecType.LOCAL, nprocs=None, ppn=None,
         self._set_env(env)
         self.cwd = cwd
         self.sudo = sudo
+        self.sudoenv = sudoenv
         self.sleep_ms = sleep_ms
         self.collect_output = collect_output
         self.pipe_stdout = pipe_stdout
@@ -152,6 +155,7 @@ def __init__(self):
         self.exit_code = None
         self.stdout = ''
         self.stderr = ''
+        self.jutil = JutilManager.get_instance()
 
     def failed(self):
         return self.exit_code != 0
@@ -164,20 +168,31 @@ def set_exit_code(self):
     def wait(self):
         pass
 
-    def smash_cmd(self, cmds):
+    def smash_cmd(self, cmds, sudo, basic_env, sudoenv):
         """
         Convert a list of commands into a single command for the shell
         to execute.
 
         :param cmds: A list of commands or a single command string
+        :param prefix: A prefix for each command
+        :param sudo: Whether or not root is required
+        :param basic_env: The environment to forward to the command
+        :param sudoenv: Whether sudo supports environment forwarding
         :return:
         """
-        if isinstance(cmds, list):
-            return ' && '.join(cmds)
-        elif isinstance(cmds, str):
-            return cmds
-        else:
-            raise Exception('Command must be either list or string')
+        env = None
+        if sudo:
+            env = ''
+            if sudoenv:
+                env = [f'-E {key}=\"{val}\"' for key, val in
+                       basic_env.items()]
+                env = ' '.join(env)
+            env = f'sudo {env}'
+        if not isinstance(cmds, (list, tuple)):
+            cmds = [cmds]
+        if env is not None:
+            cmds = [f'{env} {cmd}' for cmd in cmds]
+        return ' && '.join(cmds)
 
     def wait_list(self, nodes):
         for node in nodes:

jarvis_util/shell/local_exec.py

@@ -27,8 +27,6 @@ def __init__(self, cmd, exec_info):
         """
 
         super().__init__()
-        jutil = JutilManager.get_instance()
-        cmd = self.smash_cmd(cmd)
 
         # Managing console output and collection
         self.collect_output = exec_info.collect_output
@@ -39,13 +37,13 @@ def __init__(self, cmd, exec_info):
         self.hide_output = exec_info.hide_output
         # pylint: disable=R1732
         if self.collect_output is None:
-            self.collect_output = jutil.collect_output
+            self.collect_output = self.jutil.collect_output
         if self.pipe_stdout is not None:
             self.pipe_stdout_fp = open(self.pipe_stdout, 'wb')
         if self.pipe_stderr is not None:
             self.pipe_stderr_fp = open(self.pipe_stderr, 'wb')
         if self.hide_output is None:
-            self.hide_output = jutil.hide_output
+            self.hide_output = self.jutil.hide_output
         # pylint: enable=R1732
         self.stdout = io.StringIO()
         self.stderr = io.StringIO()
@@ -57,13 +55,13 @@ def __init__(self, cmd, exec_info):
         self.exit_code = 0
 
         # Copy ENV
+        self.basic_env = exec_info.basic_env.copy()
         self.env = exec_info.env.copy()
         for key, val in os.environ.items():
             if key not in self.env:
                 self.env[key] = val
 
         # Managing command execution
-        self.cmd = cmd
         self.sudo = exec_info.sudo
         self.stdin = exec_info.stdin
         self.exec_async = exec_info.exec_async
@@ -72,13 +70,15 @@ def __init__(self, cmd, exec_info):
             self.cwd = os.getcwd()
         else:
             self.cwd = exec_info.cwd
-        if jutil.debug_local_exec:
+
+        # Create the command
+        cmd = self.smash_cmd(cmd, self.sudo, self.basic_env, exec_info.sudoenv)
+        self.cmd = cmd
+        if self.jutil.debug_local_exec:
             print(cmd)
         self._start_bash_processes()
 
     def _start_bash_processes(self):
-        if self.sudo:
-            self.cmd = f'sudo {self.cmd}'
         time.sleep(self.sleep_ms)
         # pylint: disable=R1732
         self.proc = subprocess.Popen(self.cmd,

jarvis_util/shell/pssh_exec.py

@@ -22,7 +22,6 @@ def __init__(self, cmd, exec_info):
         :param exec_info: Info needed to execute command with SSH
         """
         super().__init__()
-        self.cmd = self.smash_cmd(cmd)
         self.exec_async = exec_info.exec_async
         self.hosts = exec_info.hostfile.hosts
         self.execs_ = []

jarvis_util/shell/scp.py

@@ -4,6 +4,7 @@
 """
 from .local_exec import LocalExec
 from .exec_info import Executable
+from jarvis_util.jutil_manager import JutilManager
 
 
 class _Scp(LocalExec):
@@ -22,12 +23,15 @@ def __init__(self, src_path, dst_path, exec_info):
         """
 
         self.addr = exec_info.hostfile.hosts[0]
+        if self.addr == 'localhost' or self.addr == '127.0.0.1':
+            return
         self.src_path = src_path
         self.dst_path = dst_path
         self.user = exec_info.user
         self.pkey = exec_info.pkey
         self.port = exec_info.port
         self.sudo = exec_info.sudo
+        self.jutil = JutilManager.get_instance()
         super().__init__(self.rsync_cmd(src_path, dst_path),
                          exec_info.mod(env=exec_info.basic_env))
 
@@ -47,6 +51,8 @@ def rsync_cmd(self, src_path, dst_path):
         else:
             lines.append(f'{self.addr}:{dst_path}')
         rsync_cmd = ' '.join(lines)
+        if self.jutil.debug_scp:
+            print(rsync_cmd)
         return rsync_cmd
 
 

jarvis_util/shell/ssh_exec.py

@@ -19,13 +19,14 @@ def __init__(self, cmd, exec_info):
         :param exec_info: Info needed to execute command with SSH
         """
 
-        cmd = self.smash_cmd(cmd)
         self.addr = exec_info.hostfile.hosts[0]
         self.user = exec_info.user
         self.pkey = exec_info.pkey
         self.port = exec_info.port
         self.sudo = exec_info.sudo
         self.ssh_env = exec_info.env
+        self.basic_env = exec_info.env
+        cmd = self.smash_cmd(cmd, self.sudo, self.basic_env, exec_info.sudoenv)
         if not exec_info.hostfile.is_local():
             super().__init__(self.ssh_cmd(cmd),
                              exec_info.mod(env=exec_info.basic_env))
@@ -50,7 +51,7 @@ def ssh_cmd(self, cmd):
                 cmd_lines.append(f'{key}=\"{val}\"')
         cmd_lines.append(cmd)
         env_cmd = ' '.join(cmd_lines)
-        real_cmd = f'{ssh_cmd} \"{env_cmd}\"'
+        real_cmd = f'{ssh_cmd} {env_cmd}'
         return real_cmd
 
 

jarvis_util/util/argparse.py

@@ -45,6 +45,7 @@ def __init__(self, args=None, exit_on_fail=True, **custom_info):
         self.menu = None
         self.menu_name = None
         self.kwargs = {}
+        self.real_kwargs = {}
         self.define_options()
         self._parse()
 
@@ -190,6 +191,7 @@ def _parse(self):
         default_args = self.default_kwargs(
             list(self.menu['kw_opts'].values()) + self.menu['pos_opts'])
         default_args.update(self.kwargs)
+        self.real_kwargs = self.kwargs
         self.kwargs = default_args
 
     @staticmethod

jarvis_util/util/hostfile.py

@@ -193,14 +193,11 @@ def save(self, path):
             fp.write('\n'.join(self.all_hosts))
         return self
 
-    def ip_list(self):
-        return self.hosts_ip
-
-    def hostname_list(self):
-        return self.hosts
+    def list(self):
+        return [Hostfile(all_hosts=[host]) for host in self.hosts]
 
     def enumerate(self):
-        return enumerate(self.hosts)
+        return enumerate(self.list())
 
     def host_str(self, sep=','):
         return sep.join(self.hosts)