improve upload

Author: KZHIWEI

dpdispatcher/dp_cloud_server.py

@@ -1,5 +1,8 @@
+import shutil
+
 from dpdispatcher.JobStatus import JobStatus
 from dpdispatcher import dlog
+from dpdispatcher.dpcloudserver import zip_file
 from dpdispatcher.machine import Machine
 from dpdispatcher.dpcloudserver.api import API
 from dpdispatcher.dpcloudserver.config import ALI_OSS_BUCKET_URL
@@ -26,7 +29,7 @@ def __init__(self, context):
         username = context.remote_profile.get('username', None)
         password = context.remote_profile.get('password', None)
         if email is None and username is not None:
-            raise dlog.exception("username is no longer support in current version, "
+            dlog.exception("username is no longer support in current version, "
                                      "please consider use email instead of username.", DeprecationWarning)
         if email is None:
             raise ValueError("can not find email in remote_profile, please check your machine file.")
@@ -130,8 +133,27 @@ def check_status(self, job):
                 raise RuntimeError(f"cannot find job information in dpcloudserver's database for job {job.job_id} {check_return} {retry_return}")
 
         job_state = self.map_dp_job_state(dp_job_status)
+        if job_state == JobStatus.finished:
+            self._download_job(job)
         return job_state
 
+
+    def _download_job(self, job):
+        job_url = self.api.get_job_result_url(job.job_id)
+        if not job_url:
+            return
+        job_hash = job.job_hash
+        result_filename = job_hash + '_back.zip'
+        target_result_zip = os.path.join(self.context.local_root, result_filename)
+        self.api.download_from_url(job_url, target_result_zip)
+        zip_file.unzip_file(target_result_zip, out_dir=self.context.local_root)
+        try:
+            os.makedirs(os.path.join(self.context.local_root, 'backup'), exist_ok=True)
+            shutil.move(target_result_zip,
+                        os.path.join(self.context.local_root, 'backup', os.path.split(target_result_zip)[1]))
+        except (OSError, shutil.Error) as e:
+            dlog.exception("unable to backup file, " + str(e))
+
     def check_finish_tag(self, job):
         job_tag_finished = job.job_hash + '_job_tag_finished'
         dlog.info('check if job finished: ',job.job_id, job_tag_finished)

dpdispatcher/dp_cloud_server_context.py

@@ -124,7 +124,6 @@ def upload(self, submission):
             )
             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))
-        print()  # empty print because tqdm may not print a newline along with dlog
         return result
         # return oss_task_zip
         # api.upload(self.oss_task_dir, zip_task_file)
@@ -152,27 +151,37 @@ def download(self, submission):
                         job_hash = job_hashs[each['task_id']]
                     job_infos[job_hash] = each
         bar_format = "{l_bar}{bar}| {n:.02f}/{total:.02f} %  [{elapsed}<{remaining}, {rate_fmt}{postfix}]"
-        for job_hash, info in tqdm.tqdm(job_infos.items(), desc="Downloading from Lebesgue", bar_format=bar_format):
+        for job_hash, info in tqdm.tqdm(job_infos.items(), desc="Validating download file from Lebesgue", bar_format=bar_format):
             result_filename = job_hash + '_back.zip'
             target_result_zip = os.path.join(self.local_root, result_filename)
+            if self._check_if_job_has_already_downloaded(target_result_zip, self.local_root):
+               continue
             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))
-        print()  # empty print because tqdm may not print a newline along with dlog
+        self._clean_backup(self.local_root, keep_backup=self.remote_profile.get('keep_backup', True))
         return True
 
+    def _check_if_job_has_already_downloaded(self, target, local_root):
+        backup_file_location = os.path.join(local_root, 'backup', os.path.split(target)[1])
+        if os.path.exists(backup_file_location):
+            return False
+        else:
+            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)
+            # 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]))
         except (OSError, shutil.Error) as e:
             dlog.exception("unable to backup file, " + str(e))
 
+    def _clean_backup(self, local_root, keep_backup=True):
+        if not keep_backup:
+            os.removedirs(os.path.join(local_root, 'backup'))
+
     def write_file(self, fname, write_str):
         result = self.write_home_file(fname, write_str)
         return result

dpdispatcher/dpcloudserver/api.py

@@ -251,4 +251,20 @@ def check_job_has_uploaded(self, job_id):
             dlog.error(e)
             return False
 
+    def get_job_result_url(self, job_id):
+        try:
+            if not job_id:
+                return None
+            if 'job_group_id' in job_id:
+                ids = job_id.split(":job_group_id:")
+                job_id, _ = int(ids[0]), int(ids[1])
+            ret = self.get(f'data/job/{job_id}', {})
+            if 'result_url' in ret and len(ret['result_url']) != 0:
+                return ret.get('result_url')
+            else:
+                return None
+        except ValueError as e:
+            dlog.error(e)
+            return None
+
 # %%