add module purge, improve envs and upload oss path with random address. (#155)

* add envs , module purge and lebesgue context

* fix format

* fix backup

* add module_purge to test

* fix test case

* refract backup

* fix unexpected lebesgue job status

* change logger method to exception

* fix dlog

* fix type

Author: KZHIWEI

dpdispatcher/__init__.py

@@ -46,14 +46,14 @@
 from .pbs import Torque
 from .shell import Shell
 from .lsf import LSF
-from .dp_cloud_server import DpCloudServer
+from .dp_cloud_server import DpCloudServer, Lebesgue
 from .distributed_shell import DistributedShell
 from .machine import Machine
 
 from .lazy_local_context import LazyLocalContext
 from .local_context import LocalContext
 from .ssh_context import SSHContext
-from .dp_cloud_server_context import DpCloudServerContext
+from .dp_cloud_server_context import DpCloudServerContext, LebesgueContext
 from .hdfs_context import HDFSContext
 
 def info():

dpdispatcher/distributed_shell.py

@@ -46,6 +46,9 @@ def gen_script_env(self, job):
         source_files_part = ""
 
         module_unload_part = ""
+        module_purge = job.resources.module_purge
+        if module_purge:
+            module_unload_part += "module purge\n"
         module_unload_list = job.resources.module_unload_list
         for ii in module_unload_list:
             module_unload_part += f"module unload {ii}\n"
@@ -63,7 +66,11 @@ def gen_script_env(self, job):
         export_envs_part = ""
         envs = job.resources.envs
         for k, v in envs.items():
-            export_envs_part += f"export {k}={v}\n"
+            if isinstance(v, list):
+                for each_value in v:
+                    export_envs_part += f"export {k}={each_value}\n"
+            else:
+                export_envs_part += f"export {k}={v}\n"
 
         flag_if_job_task_fail = job.job_hash + '_flag_if_job_task_fail'
 

dpdispatcher/dp_cloud_server.py

@@ -6,6 +6,7 @@
 import time
 import warnings
 import os
+import uuid
 
 shell_script_header_template="""
 #!/bin/bash -l
@@ -60,10 +61,24 @@ def _gen_backward_files_list(self, job):
             result_file_list.extend([ os.path.join(task.task_work_path,b_f) for b_f in task.backward_files])
         return result_file_list
 
+    def _gen_oss_path(self, job, zip_filename):
+        if hasattr(job, 'upload_path') and job.upload_path:
+            return job.upload_path
+        else:
+            program_id = self.context.remote_profile.get('program_id')
+            if program_id is None:
+                dlog.info("can not find program id in remote profile, upload to default program id.")
+                program_id = 0
+            uid = uuid.uuid4()
+            path = os.path.join("program", str(program_id), str(uid), zip_filename)
+            setattr(job, 'upload_path', path)
+            return path
+
     def do_submit(self, job):
         self.gen_local_script(job)
         zip_filename = job.job_hash + '.zip'
-        oss_task_zip = 'indicate/' + job.job_hash + '/' + zip_filename
+        # oss_task_zip = 'indicate/' + job.job_hash + '/' + zip_filename
+        oss_task_zip = self._gen_oss_path(job, zip_filename)
         job_resources = ALI_OSS_BUCKET_URL + oss_task_zip
 
         input_data = self.input_data.copy()
@@ -167,3 +182,6 @@ def map_dp_job_state(status):
     #     job_tag_finished = job.job_hash + '_job_tag_finished'
     #     return self.context.check_file_exists(job_tag_finished)
 
+
+class Lebesgue(DpCloudServer):
+    pass

dpdispatcher/dp_cloud_server_context.py

@@ -1,14 +1,18 @@
 #!/usr/bin/env python
 # coding: utf-8
 # %%
+import uuid
+
 from dargs.dargs import Argument
 from dpdispatcher.base_context import BaseContext
 from typing import List
 import os
-# from dpdispatcher import dlog
+from dpdispatcher import dlog
 # from dpdispatcher.submission import Machine
+# from . import dlog
 from .dpcloudserver.api import API
 from .dpcloudserver import zip_file
+import shutil
 # from zip_file import zip_files
 DP_CLOUD_SERVER_HOME_DIR = os.path.join(
     os.path.expanduser('~'),
@@ -66,6 +70,18 @@ def bind_submission(self, submission):
         #     file_uuid = uuid.uuid1().hex
         # oss_task_dir = os.path.join()
 
+    def _gen_oss_path(self, job, zip_filename):
+        if hasattr(job, 'upload_path') and job.upload_path:
+            return job.upload_path
+        else:
+            program_id = self.remote_profile.get('program_id')
+            if program_id is None:
+                program_id = 0
+            uid = uuid.uuid4()
+            path = os.path.join("program", str(program_id), str(uid), zip_filename)
+            setattr(job, 'upload_path', path)
+            return path
+
     def upload(self, submission):
         # oss_task_dir = os.path.join('%s/%s/%s.zip' % ('indicate', file_uuid, file_uuid))
         # zip_filename = submission.submission_hash + '.zip'
@@ -77,7 +93,7 @@ def upload(self, submission):
         for job in submission.belonging_jobs:
             self.machine.gen_local_script(job)
             zip_filename = job.job_hash + '.zip'
-            oss_task_zip = 'indicate/' + job.job_hash + '/' + zip_filename
+            oss_task_zip = self._gen_oss_path(job, zip_filename)
             zip_task_file = os.path.join(self.local_root, zip_filename)
 
             upload_file_list = [job.script_file_name, ]
@@ -97,6 +113,7 @@ def upload(self, submission):
                 file_list=upload_file_list
             )
             result = self.api.upload(oss_task_zip, upload_zip, ENDPOINT, BUCKET_NAME)
+            self._backup(self.local_root, upload_zip, keep_backup=self.remote_profile.get('keep_backup', True))
         return result
         # return oss_task_zip
         # api.upload(self.oss_task_dir, zip_task_file)
@@ -110,23 +127,42 @@ def download(self, submission):
             if isinstance(job.job_id, str) and ':job_group_id:' in job.job_id:
                 ids = job.job_id.split(":job_group_id:")
                 jid, gid = int(ids[0]), int(ids[1])
-                job_hashs[jid] = job.job_hash 
+                job_hashs[jid] = job.job_hash
                 group_id = gid
             else:
-                job_infos[job.job_hash] = self.get_tasks(job.job_id)[0]
+                job_infos[job.job_hash] = self.api.get_tasks(job.job_id)[0]
         if group_id is not None:
             job_result = self.api.get_tasks_v2_list(group_id)
             for each in job_result:
                 if 'result_url' in each and each['result_url'] != '' and each['status'] == 2:
-                    job_hash = job_hashs[each['task_id']]
+                    job_hash = ''
+                    if each['task_id'] not in job_hashs:
+                        dlog.info(f"find unexpect job_hash, but task {each['task_id']} still been download.")
+                        dlog.debug(str(job_hashs))
+                        job_hash = str(each['task_id'])
+                    else:
+                        job_hash = job_hashs[each['task_id']]
                     job_infos[job_hash] = each
-        for hash, info in job_infos.items():
-            result_filename = hash + '_back.zip'
+        for job_hash, info in job_infos.items():
+            result_filename = job_hash + '_back.zip'
             target_result_zip = os.path.join(self.local_root, result_filename)
             self.api.download_from_url(info['result_url'], target_result_zip)
             zip_file.unzip_file(target_result_zip, out_dir=self.local_root)
+            self._backup(self.local_root, target_result_zip, keep_backup=self.remote_profile.get('keep_backup', True))
         return True
 
+    def _backup(self, local_root, target, keep_backup=True):
+        try:
+            if keep_backup:
+                # move to backup directory
+                os.makedirs(os.path.join(local_root, 'backup'), exist_ok=True)
+                shutil.move(target,
+                            os.path.join(local_root, 'backup', os.path.split(target)[1]))
+            else:
+                os.remove(target)
+        except (OSError, shutil.Error) as e:
+            dlog.exception("unable to backup file, " + str(e))
+
     def write_file(self, fname, write_str):
         result = self.write_home_file(fname, write_str)
         return result
@@ -214,4 +250,9 @@ def machine_subfields(cls) -> List[Argument]:
                                                                          'None or empty.')
             ], optional=False, doc="Configuration of job"),
         ], doc=doc_remote_profile)]
+
+
+class LebesgueContext(DpCloudServerContext):
+    pass
+
 #%%

dpdispatcher/machine.py

@@ -201,6 +201,9 @@ def gen_script_env(self, job):
         source_files_part = ""
 
         module_unload_part = ""
+        module_purge = job.resources.module_purge
+        if module_purge:
+            module_unload_part += "module purge\n"
         module_unload_list = job.resources.module_unload_list
         for ii in module_unload_list:
             module_unload_part += f"module unload {ii}\n"
@@ -218,7 +221,11 @@ def gen_script_env(self, job):
         export_envs_part = ""
         envs = job.resources.envs
         for k,v in envs.items():
-            export_envs_part += f"export {k}={v}\n"
+            if isinstance(v, list):
+                for each_value in v:
+                    export_envs_part += f"export {k}={each_value}\n"
+            else:
+                export_envs_part += f"export {k}={v}\n"
 
         flag_if_job_task_fail = job.job_hash + '_flag_if_job_task_fail'
 

dpdispatcher/submission.py

@@ -659,6 +659,7 @@ def __init__(self,
                 strategy=default_strategy,
                 para_deg=1,
                 module_unload_list=[],
+                module_purge=False,
                 module_list=[],
                 source_list=[],
                 envs={},
@@ -673,6 +674,7 @@ def __init__(self,
         self.custom_flags = custom_flags
         self.strategy = strategy
         self.para_deg = para_deg
+        self.module_purge = module_purge
         self.module_unload_list = module_unload_list
         self.module_list = module_list
         self.source_list = source_list
@@ -707,6 +709,7 @@ def serialize(self):
         resources_dict['custom_flags'] = self.custom_flags
         resources_dict['strategy'] = self.strategy
         resources_dict['para_deg'] = self.para_deg
+        resources_dict['module_purge'] = self.module_purge
         resources_dict['module_unload_list'] = self.module_unload_list
         resources_dict['module_list'] = self.module_list
         resources_dict['source_list'] = self.source_list
@@ -725,6 +728,7 @@ def deserialize(cls, resources_dict):
                         custom_flags=resources_dict.get('custom_flags', []),
                         strategy=resources_dict.get('strategy', default_strategy),
                         para_deg=resources_dict.get('para_deg', 1),
+                        module_purge=resources_dict.get('module_purge', False),
                         module_unload_list=resources_dict.get('module_unload_list', []),
                         module_list=resources_dict.get('module_list', []),
                         source_list=resources_dict.get('source_list', []),
@@ -757,6 +761,7 @@ def arginfo():
         doc_custom_flags = 'The extra lines pass to job submitting script header'
         doc_para_deg = 'Decide how many tasks will be run in parallel.'
         doc_source_list = 'The env file to be sourced before the command execution.'
+        doc_module_purge = 'Remove all modules on HPC system before module load (module_list)'
         doc_module_unload_list = 'The modules to be unloaded on HPC system before submitting jobs'
         doc_module_list = 'The modules to be loaded on HPC system before submitting jobs'
         doc_envs = 'The environment variables to be exported on before submitting jobs'
@@ -779,6 +784,7 @@ def arginfo():
             strategy_format,
             Argument("para_deg", int, optional=True, doc=doc_para_deg, default=1),
             Argument("source_list", list, optional=True, doc=doc_source_list, default=[]),
+            Argument("module_purge", bool, optional=True, doc=doc_module_purge, default=False),
             Argument("module_unload_list", list, optional=True, doc=doc_module_unload_list, default=[]),
             Argument("module_list", list, optional=True, doc=doc_module_list, default=[]),
             Argument("envs", dict, optional=True, doc=doc_envs, default={}),

tests/sample_class.py

@@ -42,6 +42,7 @@ def get_sample_resources_dict(cls):
             'custom_flags':[],
             'strategy':{'if_cuda_multi_devices': False}, 
             'para_deg':1,
+            'module_purge':False,
             'module_unload_list':[],
             'module_list':[],
             'source_list':[],