lazy_local_context.py

import os,shutil,uuid
import subprocess as sp
from glob import glob

class SPRetObj(object) :
    def __init__ (self,
                  ret) :
        self.data = ret

    def read(self) :
        return self.data

    def readlines(self) :
        lines = self.data.decode('utf-8').splitlines()
        ret = []
        for aa in lines:
            ret.append(aa+'\n')
        return ret

class LazyLocalContext(object) :
    def __init__ (self,
                  local_root,
                  work_profile = None,
                  job_uuid = None) :
        """
        work_profile:
        local_root:
        """
        assert(type(local_root) == str)
        self.local_root = os.path.abspath(local_root)
        self.remote_root = self.local_root
        if job_uuid:
           self.job_uuid=job_uuid
        else:
           self.job_uuid = str(uuid.uuid4())
        
    def get_job_root(self) :
        return self.local_root

    def upload(self,
               job_dirs,
               local_up_files,
               dereference = True) :
        pass

    def download(self, 
                 job_dirs,
                 remote_down_files,
                 check_exists = False,
                 mark_failure = True,
                 back_error=False) :
        for ii in job_dirs :
            for jj in remote_down_files :
                fname = os.path.join(self.local_root, ii, jj)
                exists = os.path.exists(fname)
                if not exists:
                    if check_exists:
                        if mark_failure:
                            with open(os.path.join(self.local_root, ii, 'tag_failure_download_%s' % jj), 'w') as fp: pass
                        else:
                            pass
                    else:
                        raise RuntimeError('do not find download file ' + fname)

    def block_checkcall(self,
                        cmd) :
        cwd = os.getcwd()
        os.chdir(self.local_root)
        proc = sp.Popen(cmd, shell=True, stdout = sp.PIPE, stderr = sp.PIPE)
        o, e = proc.communicate()
        stdout = SPRetObj(o)
        stderr = SPRetObj(e)
        code = proc.returncode
        if code != 0:
            os.chdir(cwd)        
            raise RuntimeError("Get error code %d in locally calling %s with job: %s " % (code, cmd, self.job_uuid))
        os.chdir(cwd)        
        return None, stdout, stderr
        
    def block_call(self, cmd) :
        cwd = os.getcwd()
        os.chdir(self.local_root)
        proc = sp.Popen(cmd, shell=True, stdout = sp.PIPE, stderr = sp.PIPE)
        o, e = proc.communicate()
        stdout = SPRetObj(o)
        stderr = SPRetObj(e)
        code = proc.returncode
        os.chdir(cwd)        
        return code, None, stdout, stderr

    def clean(self):
        tmp_files = [f'{self.job_uuid}.sub', 
                     f'{self.job_uuid}_job_id',
                     f'{self.job_uuid}_tag_finished']
        for fn in tmp_files:
            if self.check_file_exists(fn):
                os.remove(os.path.join(self.local_root, fn))

    def write_file(self, fname, write_str):
        with open(os.path.join(self.local_root, fname), 'w') as fp :
            fp.write(write_str)

    def read_file(self, fname):
        with open(os.path.join(self.local_root, fname), 'r') as fp:
            ret = fp.read()
        return ret

    def check_file_exists(self, fname):
        return os.path.isfile(os.path.join(self.local_root, fname))
        
    def call(self, cmd) :
        cwd = os.getcwd()
        os.chdir(self.local_root)
        proc = sp.Popen(cmd, shell=True, stdout = sp.PIPE, stderr = sp.PIPE)
        os.chdir(cwd)        
        return proc

    def kill(self, proc):
        proc.kill()

    def check_finish(self, proc):
        return (proc.poll() != None)

    def get_return(self, proc):
        ret = proc.poll()
        if ret is None:
            return None, None, None
        else :
            try:
                o, e = proc.communicate()
                stdout = SPRetObj(o)
                stderr = SPRetObj(e)
            except:
                stdout = None
                stderr = None
        return ret, stdout, stderr