updates for interactive autentification etc

wetb/hawc2/simulation.py

@@ -20,6 +20,7 @@ from future import standard_library
 from wetb.hawc2 import log_file
 from wetb.hawc2.htc_file import HTCFile, fmt_path
 from wetb.hawc2.log_file import LogFile
+from wetb.utils import threadnames
 
 
 
@@ -144,19 +145,26 @@ class Simulation(object):
         self.update_status()
 
     def abort(self, update_status=True):
-        if self.status != QUEUED:
-            self.host.stop()
-            for _ in range(50):
-                if self.is_simulating is False:
-                    break
-                time.sleep(0.1)
-        if self.logFile.status not in [log_file.DONE]:
         self.status = ABORTED
         self.is_simulating = False
         self.is_done = True
+        self.host.stop()
         if update_status:
             self.update_status()
 
+#         if self.status != QUEUED:
+#             self.host.stop()
+#             for _ in range(50):
+#                 if self.is_simulating is False:
+#                     break
+#                 time.sleep(0.1)
+#         if self.logFile.status not in [log_file.DONE]:
+#             self.status = ABORTED
+#         self.is_simulating = False
+#         self.is_done = True
+#         if update_status:
+#             self.update_status()
+
     def show_status(self):
         #print ("log status:", self.logFile.status)
         if self.logFile.status == log_file.SIMULATING:
@@ -228,6 +236,8 @@ class Simulation(object):
         self.host._simulate()
         self.returncode, self.stdout = self.host.returncode, self.host.stdout
         if self.host.returncode or 'error' in self.host.stdout.lower():
+            if self.status==ABORTED:
+                return
             if "error" in self.host.stdout.lower():
                 self.errors = (list(set([l for l in self.host.stdout.split("\n") if 'error' in l.lower()])))
             self.status = ERROR
@@ -263,15 +273,12 @@ class Simulation(object):
             return dst
         turb_files = [f for f in self.htcFile.turbulence_files() if self.copy_turbulence and not os.path.isfile(os.path.join(self.exepath, f))]
         if self.ios:
-            output_patterns = [fmt(dst) for dst in (["../output/*", "../output/"] + 
-                                                    turb_files + 
-                                                    [os.path.join(self.exepath, self.stdout_filename)])]
-            output_files = set([fmt_path(f) for pattern in output_patterns for f in self.host.glob(fmt_path(os.path.join(self.tmp_exepath, pattern)), recursive=True)])
+            output_patterns = ["../output/*", "../output/"] + turb_files + [os.path.join(self.exepath, self.stdout_filename)]
         else:
-            output_patterns = [fmt(dst) for dst in (self.htcFile.output_files() + 
-                                                    turb_files + 
-                                                    [os.path.join(self.exepath, self.stdout_filename)])]
-            output_files = set([fmt_path(f) for pattern in output_patterns for f in self.host.glob(fmt_path(os.path.join(self.tmp_exepath, pattern)))])
+            output_patterns = self.htcFile.output_files() + turb_files + [os.path.join(self.exepath, self.stdout_filename)]
+        output_files = set(self.host.glob([fmt_path(os.path.join(self.tmp_exepath,fmt(p))) for p in output_patterns], recursive=self.ios))
+            
+            
         try:
             self.host._finish_simulation(output_files)
             if self.status != ERROR:
@@ -325,6 +332,7 @@ class Simulation(object):
 
 
     def simulate_distributed(self):
+        threadnames.register("Simulation %s"%self.simulation_id)
         try:
             self.prepare_simulation()
             try:
@@ -335,6 +343,7 @@ class Simulation(object):
                 raise
             finally:
                 try:
+                    if self.status!=ABORTED:
                         self.finish_simulation()
                 except:
                     print ("finish_simulation failed", str(self))

wetb/hawc2/simulation_resources.py

@@ -23,7 +23,7 @@ from wetb.utils.cluster_tools.pbsjob import SSHPBSJob, NOT_SUBMITTED, DONE
 from wetb.utils.cluster_tools.ssh_client import SSHClient
 from wetb.utils.timing import print_time
 from wetb.hawc2.htc_file import fmt_path
-
+import numpy as np
 
 class SimulationHost(object):
     def __init__(self, simulation):
@@ -60,6 +60,8 @@ class LocalSimulationHost(SimulationHost):
         return datetime.now()
 
     def glob(self, path, recursive=False):
+        if isinstance(path, list):
+            return [f for p in path for f in self.glob(p, recursive)]
         if recursive:
             return [os.path.join(root, f) for root, _, files in os.walk(path) for f in files]
         else:
@@ -168,7 +170,7 @@ class SimulationThread(Thread):
                 self.process = subprocess.Popen(" ".join([wine, hawc2exe, htcfile]), stdout=stdout, stderr=STDOUT, shell=True, cwd=exepath) #shell must be True to inwoke wine
             else:
                 self.process = subprocess.Popen([hawc2exe, htcfile], stdout=stdout, stderr=STDOUT, shell=False, cwd=exepath, creationflags=CREATE_NO_WINDOW)
-            self.process.communicate()
+            #self.process.communicate()
 
         import psutil
         try:
@@ -222,25 +224,25 @@ class PBSClusterSimulationResource(SSHPBSClusterResource):
         except:
             pass
 
-    def update_status(self):
+    def update_resource_status(self):
         try:
-            _, out, _ = self.execute("find .hawc2launcher/ -name '*.out'")
+            _, out, _ = self.ssh.execute("find .hawc2launcher/ -name '*.out'")
             self.finished = set([f.split("/")[1] for f in out.split("\n") if "/" in f])
-        except Exception:
-            #print ("resource_manager.update_status, out", str(e))
+        except Exception as e:
+            print ("resource_manager.update_status, out", str(e))
             pass
  
         try:
-            _, out, _ = self.execute("find .hawc2launcher -name 'status*' -exec cat {} \;")
+            _, out, _ = self.ssh.execute("find .hawc2launcher -name 'status*' -exec cat {} \;")
             self.loglines = {l.split(";")[0] : l.split(";")[1:] for l in out.split("\n") if ";" in l}
-        except Exception:
-            #print ("resource_manager.update_status, status file", str(e))
+        except Exception as e:
+            print ("resource_manager.update_status, status file", str(e))
             pass
         try:
-            _, out, _ = self.execute("qstat -u %s" % self.username)
+            _, out, _ = self.ssh.execute("qstat -u %s" % self.username)
             self.is_executing = set([j.split(".")[0] for j in out.split("\n")[5:] if "." in j])
-        except Exception:
-            #print ("resource_manager.update_status, qstat", str(e))
+        except Exception as e:
+            print ("resource_manager.update_status, qstat", str(e))
             pass
 
 class GormSimulationResource(PBSClusterSimulationResource):
@@ -250,49 +252,46 @@ source activate wetb_py3"""
         PBSClusterSimulationResource.__init__(self, "gorm.risoe.dk", username, password, 22, 25, 100, init_cmd, wine_cmd, "python")
 
 
-class PBSClusterSimulationHost(SimulationHost, SSHClient):
+class PBSClusterSimulationHost(SimulationHost):
     def __init__(self, simulation, resource):
         SimulationHost.__init__(self, simulation)
-        SSHClient.__init__(self, resource.host, resource.username, resource.password, resource.port)
-        self.pbsjob = SSHPBSJob(resource.shared_ssh)
+        self.ssh = resource.new_ssh_connection()
+        self.pbsjob = SSHPBSJob(resource.new_ssh_connection())
         self.resource = resource
 
     hawc2exe = property(lambda self : os.path.basename(self.sim.hawc2exe))
 
-
+    def glob(self, *args,**kwargs):
+        return self.ssh.glob(*args,**kwargs)
     def get_datetime(self):
-        v, out, err = self.execute('date "+%Y,%m,%d,%H,%M,%S"')
+        v, out, err = self.ssh.execute('date "+%Y,%m,%d,%H,%M,%S"')
         if v == 0:
             return datetime.strptime(out.strip(), "%Y,%m,%d,%H,%M,%S")
 
     #@print_time
     def _prepare_simulation(self, input_files):
-        with self:
-            self.execute(["mkdir -p .hawc2launcher/%s" % self.simulation_id], verbose=False)
-            self.execute("mkdir -p %s%s" % (self.tmp_exepath, os.path.dirname(self.log_filename)))
-            
-            self.upload_files(self.modelpath, self.tmp_modelpath, file_lst = [os.path.relpath(f, self.modelpath) for f in input_files])
+        with self.ssh:
+            self.ssh.execute(["mkdir -p .hawc2launcher/%s" % self.simulation_id], verbose=False)
+            self.ssh.execute("mkdir -p %s%s" % (self.tmp_exepath, os.path.dirname(self.log_filename)))
+            self.ssh.upload_files(self.modelpath, self.tmp_modelpath, file_lst = [os.path.relpath(f, self.modelpath) for f in input_files], callback=self.sim.progress_callback("Copy to host"))
 #             for src_file in input_files:
 #                     dst = unix_path(self.tmp_modelpath + os.path.relpath(src_file, self.modelpath))
 #                     self.execute("mkdir -p %s" % os.path.dirname(dst), verbose=False)
 #                     self.upload(src_file, dst, verbose=False)
 #                     ##assert self.ssh.file_exists(dst)
-
             f = io.StringIO(self.pbsjobfile(self.sim.ios))
             f.seek(0)
-            self.upload(f, self.tmp_exepath + "%s.in" % self.simulation_id)
-            self.execute("mkdir -p %s%s" % (self.tmp_exepath, os.path.dirname(self.stdout_filename)))
+            self.ssh.upload(f, self.tmp_exepath + "%s.in" % self.simulation_id)
+            self.ssh.execute("mkdir -p %s%s" % (self.tmp_exepath, os.path.dirname(self.stdout_filename)))
             remote_log_filename = "%s%s" % (self.tmp_exepath, self.log_filename)
-            self.execute("rm -f %s" % remote_log_filename)
-
-
+            self.ssh.execute("rm -f %s" % remote_log_filename)
             
     #@print_time
     def _finish_simulation(self, output_files):
-        with self:
+        with self.ssh:
             download_failed = []
             try:
-                self.download_files(self.tmp_modelpath, self.modelpath, file_lst = [os.path.relpath(f, self.tmp_modelpath) for f in output_files] )
+                self.ssh.download_files(self.tmp_modelpath, self.modelpath, file_lst = [os.path.relpath(f, self.tmp_modelpath) for f in output_files], callback=self.sim.progress_callback("Copy from host") )
             except:
 #                 
 #             for src_file in output_files:
@@ -303,15 +302,15 @@ class PBSClusterSimulationHost(SimulationHost, SSHClient):
 #                 except Exception as e:
 #                     download_failed.append(dst_file)
 #             if download_failed:
-                raise Warning("Failed to download %s from %s"%(",".join(download_failed), self.host))
+                raise Warning("Failed to download %s from %s"%(",".join(download_failed), self.ssh.host))
             else:
                 try:
-                    self.execute('rm -r .hawc2launcher/%s' % self.simulation_id)
+                    self.ssh.execute('rm -r .hawc2launcher/%s' % self.simulation_id)
                 finally:
                     try:
-                        self.execute('rm .hawc2launcher/status_%s' % self.simulation_id)
+                        self.ssh.execute('rm .hawc2launcher/status_%s' % self.simulation_id)
                     except:
-                        raise Warning("Fail to remove temporary files and folders on %s"%self.host)
+                        raise Warning("Fail to remove temporary files and folders on %s"%self.ssh.host)
 
 
     def _simulate(self):
@@ -324,9 +323,9 @@ class PBSClusterSimulationHost(SimulationHost, SSHClient):
             time.sleep(sleeptime)
 
         local_out_file = self.exepath + self.stdout_filename
-        with self:
+        with self.ssh:
             try:
-                self.download(self.tmp_exepath + self.stdout_filename, local_out_file)
+                self.ssh.download(self.tmp_exepath + self.stdout_filename, local_out_file)
                 with open(local_out_file) as fid:
                     _, self.stdout, returncode_str, _ = fid.read().split("---------------------")
                     self.returncode = returncode_str.strip() != "0"
@@ -334,7 +333,7 @@ class PBSClusterSimulationHost(SimulationHost, SSHClient):
                 self.returncode = 1
                 self.stdout = "error: Could not download and read stdout file"
             try:
-                self.download(self.tmp_exepath + self.log_filename, self.exepath + self.log_filename)
+                self.ssh.download(self.tmp_exepath + self.log_filename, self.exepath + self.log_filename)
             except Exception:
                 raise Warning ("Logfile not found", self.tmp_modelpath + self.log_filename)
         self.sim.logFile = LogFile.from_htcfile(self.htcFile, self.exepath)
@@ -400,7 +399,11 @@ class PBSClusterSimulationHost(SimulationHost, SSHClient):
             cp_back += "mkdir -p $PBS_O_WORKDIR/%s/. \n" % folder
             cp_back += "cp -R -f %s/. $PBS_O_WORKDIR/%s/.\n" % (folder, folder)
         rel_htcfilename = fmt_path(os.path.relpath(self.htcFile.filename, self.exepath))
-        
+        try:
+            steps = self.htcFile.simulation.time_stop[0] / self.htcFile.simulation.newmark.deltat[0]
+            walltime = "%02d:00:00"%np.ceil(steps/500/60)
+        except:
+            walltime = "04:00:00"
         init="""
 ### Standard Output
 #PBS -N h2l_%s
@@ -408,14 +411,14 @@ class PBSClusterSimulationHost(SimulationHost, SSHClient):
 #PBS -j oe
 #PBS -o %s
 ### Maximum wallclock time format HOURS:MINUTES:SECONDS
-#PBS -l walltime=04:00:00
+#PBS -l walltime=%s
 ###PBS -a 201547.53
 #PBS -lnodes=1:ppn=1
 ### Queue name
 #PBS -q workq
 ### Create scratch directory and copy data to it
 cd $PBS_O_WORKDIR
-pwd"""% (self.simulation_id, self.stdout_filename)
+pwd"""% (self.simulation_id, self.stdout_filename, walltime)
         copy_to="""
 cp -R %s /scratch/$USER/$PBS_JOBID
 ### Execute commands on scratch nodes

wetb/utils/cluster_tools/cluster_resource.py

@@ -66,13 +66,18 @@ class Resource(object):
     def __init__(self, min_cpu, min_free):
         self.min_cpu = min_cpu
         self.min_free = min_free
+        self.cpu_free=0
         self.acquired = 0
+        self.no_cpu="?"
+        self.used_by_user=0
         self.resource_lock = threading.Lock()
 
     def ok2submit(self):
         """Always ok to have min_cpu cpus and ok to have more if there are min_free free cpus"""
         try:
+            print ("ok2submit")
             total, free, user = self.check_resources()
+            print ("ok2submit", total, free, user)
         except:
             return False
         if user < self.min_cpu:
@@ -91,7 +96,7 @@ class Resource(object):
             self.acquired -= 1
 
 
-    def update_status(self):
+    def update_resource_status(self):
         try:
             self.no_cpu, self.cpu_free, self.used_by_user = self.check_resources()
         except Exception:
@@ -114,7 +119,8 @@ class SSHPBSClusterResource(Resource):
 
 
     def new_ssh_connection(self):
-        return SSHClient(self.host, self.username, self.password, self.port)
+        return SSHClient(self.host, self.ssh.username, self.ssh.password, self.ssh.port)
+        #return self.ssh
 
     def check_resources(self):
         with self.resource_lock:

wetb/utils/cluster_tools/pbsjob.py

@@ -32,8 +32,7 @@ class SSHPBSJob(object):
         if cwd != "":
             cmds.append("cd %s" % cwd)
         cmds.append("qsub %s" % job)
-        ssh = SSHClient(self.ssh.host, self.ssh.username, self.ssh.password, self.ssh.port)
-        _, out, _ = ssh.execute(";".join(cmds))
+        _, out, _ = self.ssh.execute(";".join(cmds))
         self.jobid = out.split(".")[0]
         self._status = PENDING
 

wetb/utils/cluster_tools/ssh_client.py

@@ -17,6 +17,9 @@ from wetb.utils.timing import print_time
 import glob
 import getpass
 from sshtunnel import SSHTunnelForwarder, SSH_CONFIG_FILE
+from wetb.utils.ui import UI
+from wetb.utils import threadnames
+
 
 
 
@@ -100,7 +103,7 @@ class SSHClient(object):
     "A wrapper of paramiko.SSHClient"
     TIMEOUT = 4
 
-    def __init__(self, host, username, password=None, port=22, key=None, passphrase=None, interactive_auth_handler=None, gateway=None):
+    def __init__(self, host, username, password=None, port=22, key=None, passphrase=None, interactive_auth_handler=None, gateway=None, ui=UI()):
         self.host = host
         self.username = username
         self.password = password
@@ -108,9 +111,11 @@ class SSHClient(object):
         self.key = key
         self.gateway=gateway
         self.interactive_auth_handler = interactive_auth_handler
+        self.ui = ui
         self.disconnect = 0
         self.client = None
-        self.sftp = None
+        self.ssh_lock = threading.RLock()
+        #self.sftp = None
         self.transport = None
         if key is not None:
             self.key = paramiko.RSAKey.from_private_key(StringIO(key), password=passphrase)
@@ -119,6 +124,7 @@ class SSHClient(object):
         return self.host, self.username, self.password, self.port
 
     def __enter__(self):
+        with self.ssh_lock:
             self.disconnect += 1
             if self.client is None or self.client._transport is None or self.client._transport.is_active() is False:
                 self.close()
@@ -133,7 +139,7 @@ class SSHClient(object):
 
     def connect(self):
         
-        print ("connect", self.host)
+        print ("connect", self.host, threadnames.name())
         #print (traceback.print_stack())
         if self.gateway:
             if self.gateway.interactive_auth_handler:
@@ -152,10 +158,15 @@ class SSHClient(object):
                                                  local_bind_address=('0.0.0.0', 10022)
                                                 )
             
+            print ("start tunnel")
             self.tunnel.start()
+            print ("self.client = paramiko.SSHClient()")
             self.client = paramiko.SSHClient()
+            
             self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+            print ('self.client.connect("127.0.0.1", 10022, username=self.username, password=self.password)')
             self.client.connect("127.0.0.1", 10022, username=self.username, password=self.password)
+            print ("done")
 
         elif self.password is None or self.password == "":
             raise IOError("Password not set for %s"%self.host)         
@@ -171,48 +182,80 @@ class SSHClient(object):
         
         
         assert self.client is not None
-        self.sftp = paramiko.SFTPClient.from_transport(self.client._transport)
+        #self.sftp = paramiko.SFTPClient.from_transport(self.client._transport)
         return self
 
+        
+    def __del__(self):
+        self.close()
+
+    @property
+    def sftp(self):
+        return paramiko.SFTPClient.from_transport(self.client._transport)
+    
+#     @sftp.setter
+#     def sftp(self, values):
+#         pass
+
     def __exit__(self, *args):
         self.disconnect -= 1
         if self.disconnect == 0:
             self.close()
 
 
-    def download(self, remotefilepath, localfile, verbose=False, retry=1):
+    def download(self, remotefilepath, localfile, verbose=False, retry=1, callback=None):
         if verbose:
             ret = None
             print ("Download %s > %s" % (remotefilepath, str(localfile)))
-        with self:
+        if callback is None:
+            callback = self.ui.progress_callback()
+        
         for i in range(retry):
             if i>0:
                 print ("Retry download %s, #%d"%(remotefilepath, i))
+
             try:
+                #print ("start download enter", threadnames.name())
+                SSHClient.__enter__(self)
+                #print ("start download", threadnames.name())
                 if isinstance(localfile, (str, bytes, int)):
-                        ret = self.sftp.get(remotefilepath, localfile)
+                    ret = self.sftp.get(remotefilepath, localfile, callback=callback)
                 elif hasattr(localfile, 'write'):
-                        ret = self.sftp.putfo(remotefilepath, localfile)
+                    ret = self.sftp.putfo(remotefilepath, localfile, callback=callback)
                 break
             except:
                 pass
+            finally:
+                #print ("End download", threadnames.name())
+                SSHClient.__exit__(self)
+                #print ("End download exit", threadnames.name())
+            
             print ("Download %s failed from %s"%(remotefilepath, self.host))
         if verbose:
             print (ret)
 
-    def upload(self, localfile, filepath, verbose=False):
+    def upload(self, localfile, filepath, verbose=False, callback=None):
         if verbose:
             print ("Upload %s > %s" % (localfile, filepath))
-        with self:
+        if callback is None:
+            callback = self.ui.progress_callback()
+        try:
+            SSHClient.__enter__(self)
             if isinstance(localfile, (str, bytes, int)):
-                ret = self.sftp.put(localfile, filepath)
+                ret = self.sftp.put(localfile, filepath, callback=callback)
             elif hasattr(localfile, 'read'):
-                ret = self.sftp.putfo(localfile, filepath)
+                size = len(localfile.read())
+                localfile.seek(0)
+                ret = self.sftp.putfo(localfile, filepath, file_size=size, callback=callback)
+        finally:
+            #print ("End upload", threadnames.name())
+            SSHClient.__exit__(self)
+            #print ("End upload exit", threadnames.name())
         if verbose:
             print (ret)
             
         
-    def upload_files(self, localpath, remotepath, file_lst=["."], compression_level=1):
+    def upload_files(self, localpath, remotepath, file_lst=["."], compression_level=1, callback=None):
         assert os.path.isdir(localpath)
         if not isinstance(file_lst, (tuple, list)):
             file_lst = [file_lst]
@@ -228,9 +271,10 @@ class SSHClient(object):
                 zipf.write(f, os.path.relpath(f, localpath))
             zipf.close()
             remote_zn = os.path.join(remotepath, zn).replace("\\","/")
+            with self:
                 self.execute("mkdir -p %s"%(remotepath))
                 
-            self.upload(zn, remote_zn)
+                self.upload(zn, remote_zn, callback=callback)
                 self.execute("unzip %s -d %s && rm %s"%(remote_zn, remotepath, remote_zn))
         except:
             raise
@@ -238,17 +282,19 @@ class SSHClient(object):
             os.remove(zn)
         
     
-    def download_files(self, remote_path, localpath, file_lst=["."], compression_level=1):
+    def download_files(self, remote_path, localpath, file_lst=["."], compression_level=1, callback=None):
         if not isinstance(file_lst, (tuple, list)):
             file_lst = [file_lst]
         file_lst = [f.replace("\\","/") for f in file_lst]
-        remote_zip = os.path.join(remote_path, "tmp.zip").replace("\\","/")
-        self.execute("cd %s && zip -r tmp.zip %s"%(remote_path, " ".join(file_lst)))
+        zn =  'tmp_%s_%s.zip'%(id(self),time.time())
+        
+        remote_zip = os.path.join(remote_path, zn).replace("\\","/")
+        self.execute("cd %s && zip -r %s %s"%(remote_path, zn, " ".join(file_lst)))
         
-        local_zip = os.path.join(localpath, "tmp.zip")
+        local_zip = os.path.join(localpath, zn)
         if not os.path.isdir(localpath):
             os.makedirs(localpath)
-        self.download(remote_zip, local_zip)
+        self.download(remote_zip, local_zip, callback=callback)
         self.execute("rm -f %s" % remote_zip)
         with zipfile.ZipFile(local_zip, "r") as z:
             z.extractall(localpath)
@@ -256,7 +302,7 @@ class SSHClient(object):
         
 
     def close(self):
-        for x in ["sftp", "client", 'tunnel' ]:
+        for x in ["client", 'tunnel' ]:
             try:
                 getattr(self, x).close()
                 setattr(self, x, None)
@@ -319,6 +365,9 @@ class SSHClient(object):
             ret = self.execute('wine regedit tmp.reg')
 
     def glob(self, filepattern, cwd="", recursive=False):
+        if isinstance(filepattern, list):
+            with self:
+                return [f for fp in filepattern for f in self.glob(fp, cwd, recursive)]
         cwd = os.path.join(cwd, os.path.split(filepattern)[0]).replace("\\", "/")
         filepattern = os.path.split(filepattern)[1]
         if recursive:
@@ -333,36 +382,48 @@ class SSHClient(object):
 class SharedSSHClient(SSHClient):
     def __init__(self, host, username, password=None, port=22, key=None, passphrase=None, interactive_auth_handler=None, gateway=None):
         SSHClient.__init__(self, host, username, password=password, port=port, key=key, passphrase=passphrase, interactive_auth_handler=interactive_auth_handler, gateway=gateway)
-        self.shared_ssh_lock = threading.RLock()
+        
         self.shared_ssh_queue = deque()
         self.next = None
         
         
+
+
+
     def execute(self, command, sudo=False, verbose=False):
         res = SSHClient.execute(self, command, sudo=sudo, verbose=verbose)
         return res
 
-    def __enter__(self):
-        with self.shared_ssh_lock:
-            if self.next == threading.currentThread():
-                return self.client
-            self.shared_ssh_queue.append(threading.current_thread())
-            if self.next is None:
-                self.next = self.shared_ssh_queue.popleft()
 
-        while self.next != threading.currentThread():
-            time.sleep(1)
+
+
+    def __enter__(self):
+        with self.ssh_lock:
             SSHClient.__enter__(self)
+            #print ("request SSH", threading.currentThread())
+#             if len(self.shared_ssh_queue)>0 and self.shared_ssh_queue[0] == threading.get_ident():
+#                 # SSH already allocated to this thread ( multiple use-statements in "with ssh:" block 
+#                 self.shared_ssh_queue.appendleft(threading.get_ident())
+#             else:
+#                 self.shared_ssh_queue.append(threading.get_ident())
+
+            if len(self.shared_ssh_queue)>0 and self.shared_ssh_queue[0] == threading.get_ident():
+                # SSH already allocated to this thread ( multiple use-statements in "with ssh:" block 
+                self.shared_ssh_queue.popleft()
+            
+            self.shared_ssh_queue.append(threading.get_ident())
+            
+        while self.shared_ssh_queue[0] != threading.get_ident():
+            #print ("Waiting for ssh", threadnames.name(), [threadnames.name(id) for id in self.shared_ssh_queue])
+            time.sleep(2)
+        #print ("Got SSH", threadnames.name())
+        
         return self.client
 
     def __exit__(self, *args):
-        with self.shared_ssh_lock:
-            if next != threading.current_thread():
-                with self.shared_ssh_lock:
-                    if len(self.shared_ssh_queue) > 0:
-                        self.next = self.shared_ssh_queue.popleft()
-                    else:
-                        self.next = None
+        with self.ssh_lock:
+            if len(self.shared_ssh_queue)>0 and self.shared_ssh_queue[0] == threading.get_ident():
+                self.shared_ssh_queue.popleft()
             
 
 

wetb/utils/tests/test_ssh_client.py

@@ -6,6 +6,8 @@ Created on 23. dec. 2016
 import unittest
 from wetb.utils.cluster_tools.ssh_client import SSHClient
 import os
+from wetb.utils.text_ui import TextUI
+
 
 try:
     import sys
@@ -54,10 +56,12 @@ class TestSSHClient(unittest.TestCase):
                 self.ssh.execute("rm -f tmp.txt")
                 io.StringIO()
                 
-                txt = "Hello world"
+                txt = "Hello world"*1000000
                 f = io.StringIO(txt)
                 f.seek(0)
-                self.ssh.upload(f, "tmp.txt")
+                print ("start upload")
+                self.ssh.upload(f, "tmp.txt", callback = TextUI().progress_callback("Uploading"))
+                print ("endupload")
                 _,out,_ = self.ssh.execute("cat tmp.txt")
                 self.assertEqual(out, txt)
                 fn = tfp + "tmp.txt"
@@ -114,7 +118,7 @@ class TestSSHClient(unittest.TestCase):
                 self.assertEqual(out.strip(), "g-000.risoe.dk")
                 
     def test_ssh_g047(self):
-        if 1 or all:
+        if 0 or all:
             if x:
                 gateway = SSHClient('gorm.risoe.dk', 'mmpe', x.mmpe)
                 ssh = SSHClient('g-047', "mmpe", x.mmpe, gateway=gateway)
@@ -128,10 +132,10 @@ class TestSSHClient(unittest.TestCase):
                 self.assertEqual(out.strip(), "ssh-03.risoe.dk")
 
     def test_ssh_risoe_gorm(self):
-        if 0 or all:
+        if 1 or all:
             if x:
             
-                gateway = SSHClient('ssh.risoe.dk', 'mmpe', interactive_auth_handler = sshrisoe_interactive_auth_handler(x.mmpe))
+                gateway = SSHClient('ssh.risoe.dk', 'mmpe', password="xxx", interactive_auth_handler = sshrisoe_interactive_auth_handler(x.mmpe))
                 ssh = SSHClient('10.40.23.49', 'mmpe', x.mmpe, gateway = gateway)
                 _,out,_ = ssh.execute("hostname")
                 self.assertEqual(out.strip(), "g-000.risoe.dk")