dpdispatcher 0.6.1__py3-none-any.whl → 1.0.0__py3-none-any.whl

dpdispatcher/_version.py

@@ -1,16 +1,34 @@
-# file generated by setuptools_scm
+# file generated by setuptools-scm
 # don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
 TYPE_CHECKING = False
 if TYPE_CHECKING:
-    from typing import Tuple, Union
+    from typing import Tuple
+    from typing import Union
+
     VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
 else:
     VERSION_TUPLE = object
+    COMMIT_ID = object
 
 version: str
 __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = '1.0.0'
+__version_tuple__ = version_tuple = (1, 0, 0)
 
-__version__ = version = '0.6.1'
-__version_tuple__ = version_tuple = (0, 6, 1)
+__commit_id__ = commit_id = None

dpdispatcher/base_context.py

@@ -1,5 +1,5 @@
 from abc import ABCMeta, abstractmethod
-from typing import List, Tuple
+from typing import Any, List, Tuple
 
 from dargs import Argument
 
@@ -73,6 +73,65 @@ class BaseContext(metaclass=ABCMeta):
     def check_finish(self, proc):
         raise NotImplementedError("abstract method")
 
+    def block_checkcall(self, cmd, asynchronously=False) -> Tuple[Any, Any, Any]:
+        """Run command with arguments. Wait for command to complete.
+
+        Parameters
+        ----------
+        cmd : str
+            The command to run.
+        asynchronously : bool, optional, default=False
+            Run command asynchronously. If True, `nohup` will be used to run the command.
+
+        Returns
+        -------
+        stdin
+            standard inout
+        stdout
+            standard output
+        stderr
+            standard error
+
+        Raises
+        ------
+        RuntimeError
+            when the return code is not zero
+        """
+        if asynchronously:
+            cmd = f"nohup {cmd} >/dev/null &"
+        exit_status, stdin, stdout, stderr = self.block_call(cmd)
+        if exit_status != 0:
+            raise RuntimeError(
+                "Get error code {} in calling {} with job: {} . message: {}".format(
+                    exit_status,
+                    cmd,
+                    self.submission.submission_hash,
+                    stderr.read().decode("utf-8"),
+                )
+            )
+        return stdin, stdout, stderr
+
+    @abstractmethod
+    def block_call(self, cmd) -> Tuple[int, Any, Any, Any]:
+        """Run command with arguments. Wait for command to complete.
+
+        Parameters
+        ----------
+        cmd : str
+            The command to run.
+
+        Returns
+        -------
+        exit_status
+            exit code
+        stdin
+            standard inout
+        stdout
+            standard output
+        stderr
+            standard error
+        """
+
     @classmethod
     def machine_arginfo(cls) -> Argument:
         """Generate the machine arginfo.

dpdispatcher/contexts/__init__.py

@@ -1,4 +1,5 @@
 """Contexts."""
+
 import importlib
 from pathlib import Path
 

dpdispatcher/contexts/dp_cloud_server_context.py

@@ -161,7 +161,9 @@ class BohriumContext(BaseContext):
         # return oss_task_zip
         # api.upload(self.oss_task_dir, zip_task_file)
 
-    def download(self, submission):
+    def download(
+        self, submission, check_exists=False, mark_failure=True, back_error=False
+    ):
         jobs = submission.belonging_jobs
         job_hashs = {}
         job_infos = {}
@@ -335,6 +337,11 @@ class BohriumContext(BaseContext):
             )
         ]
 
+    def block_call(self, cmd):
+        raise RuntimeError(
+            "Unsupported method. You may use an unsupported combination of the machine and the context."
+        )
+
 
 DpCloudServerContext = BohriumContext
 LebesgueContext = BohriumContext

dpdispatcher/contexts/hdfs_context.py

@@ -88,7 +88,7 @@ class HDFSContext(BaseContext):
             for ff in task.forward_files:
                 abs_file_list = glob(os.path.join(local_job, ff))
                 if not abs_file_list:
-                    raise RuntimeError(
+                    raise FileNotFoundError(
                         "cannot find upload file " + os.path.join(local_job, ff)
                     )
                 rel_file_list = [
@@ -100,7 +100,7 @@ class HDFSContext(BaseContext):
         for fc in submission.forward_common_files:
             abs_file_list = glob(os.path.join(local_job, fc))
             if not abs_file_list:
-                raise RuntimeError(
+                raise FileNotFoundError(
                     "cannot find upload file " + os.path.join(local_job, fc)
                 )
             rel_file_list = [
@@ -138,7 +138,7 @@ class HDFSContext(BaseContext):
             shutil.rmtree(gz_dir, ignore_errors=True)
         os.mkdir(os.path.join(self.local_root, "tmp"))
         rfile_tgz = f"{self.remote_root}/{submission.submission_hash}_*_download.tar.gz"
-        lfile_tgz = "%s/tmp/" % (self.local_root)
+        lfile_tgz = f"{self.local_root}/tmp/"
         HDFS.copy_to_local(rfile_tgz, lfile_tgz)
 
         tgz_file_list = glob(os.path.join(self.local_root, "tmp/*_download.tar.gz"))
@@ -164,15 +164,17 @@ class HDFSContext(BaseContext):
                                 os.path.join(
                                     self.local_root,
                                     task.task_work_path,
-                                    "tag_failure_download_%s" % jj,
+                                    f"tag_failure_download_{jj}",
                                 ),
                                 "w",
                             ) as fp:
                                 pass
                         else:
-                            raise RuntimeError("do not find download file " + rfile)
+                            raise FileNotFoundError(
+                                "do not find download file " + rfile
+                            )
                     else:
-                        raise RuntimeError("do not find download file " + rfile)
+                        raise FileNotFoundError("do not find download file " + rfile)
                 else:
                     if os.path.exists(lfile):
                         dlog.info(f"find existing {lfile}, replacing by {rfile}")
@@ -196,16 +198,14 @@ class HDFSContext(BaseContext):
                 if check_exists:
                     if mark_failure:
                         with open(
-                            os.path.join(
-                                self.local_root, "tag_failure_download_%s" % jj
-                            ),
+                            os.path.join(self.local_root, f"tag_failure_download_{jj}"),
                             "w",
                         ) as fp:
                             pass
                     else:
-                        raise RuntimeError("do not find download file " + rfile)
+                        raise FileNotFoundError("do not find download file " + rfile)
                 else:
-                    raise RuntimeError("do not find download file " + rfile)
+                    raise FileNotFoundError("do not find download file " + rfile)
             else:
                 if os.path.exists(lfile):
                     dlog.info(f"find existing {lfile}, replacing by {rfile}")
@@ -244,3 +244,8 @@ class HDFSContext(BaseContext):
 
     def read_file(self, fname):
         return HDFS.read_hdfs_file(os.path.join(self.remote_root, fname))
+
+    def block_call(self, cmd):
+        raise RuntimeError(
+            "Unsupported method. You may use an unsupported combination of the machine and the context."
+        )

dpdispatcher/contexts/lazy_local_context.py

@@ -83,7 +83,7 @@ class LazyLocalContext(BaseContext):
 
     def upload(
         self,
-        jobs,
+        submission,
         # local_up_files,
         dereference=True,
     ):
@@ -91,7 +91,7 @@ class LazyLocalContext(BaseContext):
 
     def download(
         self,
-        jobs,
+        submission,
         # remote_down_files,
         check_exists=False,
         mark_failure=True,
@@ -112,23 +112,6 @@ class LazyLocalContext(BaseContext):
     #                else:
     #                    raise RuntimeError('do not find download file ' + fname)
 
-    def block_checkcall(self, cmd):
-        # script_dir = os.path.join(self.local_root, self.submission.work_base)
-        # os.chdir(script_dir)
-        proc = sp.Popen(
-            cmd, cwd=self.local_root, shell=True, stdout=sp.PIPE, stderr=sp.PIPE
-        )
-        o, e = proc.communicate()
-        stdout = SPRetObj(o)
-        stderr = SPRetObj(e)
-        code = proc.returncode
-        if code != 0:
-            raise RuntimeError(
-                "Get error code %d in locally calling %s with job: %s ",
-                (code, cmd, self.submission.submission_hash),
-            )
-        return None, stdout, stderr
-
     def block_call(self, cmd):
         proc = sp.Popen(
             cmd, cwd=self.local_root, shell=True, stdout=sp.PIPE, stderr=sp.PIPE

dpdispatcher/contexts/local_context.py

@@ -3,6 +3,9 @@ import shutil
 import subprocess as sp
 from glob import glob
 from subprocess import TimeoutExpired
+from typing import List
+
+from dargs import Argument
 
 from dpdispatcher.base_context import BaseContext
 from dpdispatcher.dlog import dlog
@@ -60,6 +63,7 @@ class LocalContext(BaseContext):
         self.temp_local_root = os.path.abspath(local_root)
         self.temp_remote_root = os.path.abspath(remote_root)
         self.remote_profile = remote_profile
+        self.symlink = remote_profile.get("symlink", True)
 
     @classmethod
     def load_from_dict(cls, context_dict):
@@ -83,6 +87,25 @@ class LocalContext(BaseContext):
             self.temp_remote_root, submission.submission_hash
         )
 
+    def _copy_from_local_to_remote(self, local_path, remote_path):
+        if not os.path.exists(local_path):
+            raise FileNotFoundError(
+                f"cannot find uploaded file {os.path.join(local_path)}"
+            )
+        if os.path.exists(remote_path):
+            os.remove(remote_path)
+        _check_file_path(remote_path)
+
+        if self.symlink:
+            # ensure the file exist
+            os.symlink(local_path, remote_path)
+        elif os.path.isfile(local_path):
+            shutil.copyfile(local_path, remote_path)
+        elif os.path.isdir(local_path):
+            shutil.copytree(local_path, remote_path)
+        else:
+            raise ValueError(f"Unknown file type: {local_path}")
+
     def upload(self, submission):
         os.makedirs(self.remote_root, exist_ok=True)
         for ii in submission.belonging_tasks:
@@ -94,7 +117,7 @@ class LocalContext(BaseContext):
             for kk in ii.forward_files:
                 abs_file_list = glob(os.path.join(local_job, kk))
                 if not abs_file_list:
-                    raise RuntimeError(
+                    raise FileNotFoundError(
                         "cannot find upload file " + os.path.join(local_job, kk)
                     )
                 rel_file_list = [
@@ -103,14 +126,9 @@ class LocalContext(BaseContext):
                 file_list.extend(rel_file_list)
 
             for jj in file_list:
-                if not os.path.exists(os.path.join(local_job, jj)):
-                    raise RuntimeError(
-                        "cannot find upload file " + os.path.join(local_job, jj)
-                    )
-                if os.path.exists(os.path.join(remote_job, jj)):
-                    os.remove(os.path.join(remote_job, jj))
-                _check_file_path(os.path.join(remote_job, jj))
-                os.symlink(os.path.join(local_job, jj), os.path.join(remote_job, jj))
+                self._copy_from_local_to_remote(
+                    os.path.join(local_job, jj), os.path.join(remote_job, jj)
+                )
 
         local_job = self.local_root
         remote_job = self.remote_root
@@ -119,7 +137,7 @@ class LocalContext(BaseContext):
         for kk in submission.forward_common_files:
             abs_file_list = glob(os.path.join(local_job, kk))
             if not abs_file_list:
-                raise RuntimeError(
+                raise FileNotFoundError(
                     "cannot find upload file " + os.path.join(local_job, kk)
                 )
             rel_file_list = [
@@ -128,14 +146,9 @@ class LocalContext(BaseContext):
             file_list.extend(rel_file_list)
 
         for jj in file_list:
-            if not os.path.exists(os.path.join(local_job, jj)):
-                raise RuntimeError(
-                    "cannot find upload file " + os.path.join(local_job, jj)
-                )
-            if os.path.exists(os.path.join(remote_job, jj)):
-                os.remove(os.path.join(remote_job, jj))
-            _check_file_path(os.path.join(remote_job, jj))
-            os.symlink(os.path.join(local_job, jj), os.path.join(remote_job, jj))
+            self._copy_from_local_to_remote(
+                os.path.join(local_job, jj), os.path.join(remote_job, jj)
+            )
 
     def download(
         self, submission, check_exists=False, mark_failure=True, back_error=False
@@ -153,14 +166,14 @@ class LocalContext(BaseContext):
                             tag_file_path = os.path.join(
                                 self.local_root,
                                 ii.task_work_path,
-                                "tag_failure_download_%s" % kk,
+                                f"tag_failure_download_{kk}",
                             )
                             with open(tag_file_path, "w") as fp:
                                 pass
                         else:
                             pass
                     else:
-                        raise RuntimeError(
+                        raise FileNotFoundError(
                             "cannot find download file " + os.path.join(remote_job, kk)
                         )
                 rel_flist = [
@@ -181,14 +194,16 @@ class LocalContext(BaseContext):
                                 tag_file_path = os.path.join(
                                     self.local_root,
                                     ii.task_work_path,
-                                    "tag_failure_download_%s" % jj,
+                                    f"tag_failure_download_{jj}",
                                 )
                                 with open(tag_file_path, "w") as fp:
                                     pass
                             else:
                                 pass
                         else:
-                            raise RuntimeError("do not find download file " + rfile)
+                            raise FileNotFoundError(
+                                "do not find download file " + rfile
+                            )
                     elif (not os.path.exists(rfile)) and (os.path.exists(lfile)):
                         # already downloaded
                         pass
@@ -206,7 +221,10 @@ class LocalContext(BaseContext):
                             shutil.rmtree(lfile, ignore_errors=True)
                         elif os.path.isfile(lfile) or os.path.islink(lfile):
                             os.remove(lfile)
-                        shutil.copyfile(rfile, lfile)
+                        if not os.path.islink(rfile):
+                            shutil.move(rfile, lfile)
+                        else:
+                            shutil.copyfile(rfile, lfile)
                     else:
                         raise RuntimeError("should not reach here!")
                 else:
@@ -222,14 +240,14 @@ class LocalContext(BaseContext):
                 if check_exists:
                     if mark_failure:
                         tag_file_path = os.path.join(
-                            self.local_root, "tag_failure_download_%s" % kk
+                            self.local_root, f"tag_failure_download_{kk}"
                         )
                         with open(tag_file_path, "w") as fp:
                             pass
                     else:
                         pass
                 else:
-                    raise RuntimeError(
+                    raise FileNotFoundError(
                         "cannot find download file " + os.path.join(remote_job, kk)
                     )
             rel_flist = [os.path.relpath(ii, start=remote_job) for ii in abs_flist_r]
@@ -247,7 +265,7 @@ class LocalContext(BaseContext):
                         if mark_failure:
                             with open(
                                 os.path.join(
-                                    self.local_root, "tag_failure_download_%s" % jj
+                                    self.local_root, f"tag_failure_download_{jj}"
                                 ),
                                 "w",
                             ) as fp:
@@ -255,7 +273,7 @@ class LocalContext(BaseContext):
                         else:
                             pass
                     else:
-                        raise RuntimeError("do not find download file " + rfile)
+                        raise FileNotFoundError("do not find download file " + rfile)
                 elif (not os.path.exists(rfile)) and (os.path.exists(lfile)):
                     # already downloaded
                     pass
@@ -273,28 +291,16 @@ class LocalContext(BaseContext):
                         shutil.rmtree(lfile, ignore_errors=True)
                     elif os.path.isfile(lfile) or os.path.islink(lfile):
                         os.remove(lfile)
-                    shutil.copyfile(rfile, lfile)
+                    if not os.path.islink(rfile):
+                        shutil.move(rfile, lfile)
+                    else:
+                        shutil.copyfile(rfile, lfile)
                 else:
                     raise RuntimeError("should not reach here!")
             else:
                 # no nothing in the case of linked files
                 pass
 
-    def block_checkcall(self, cmd):
-        proc = sp.Popen(
-            cmd, cwd=self.remote_root, shell=True, stdout=sp.PIPE, stderr=sp.PIPE
-        )
-        o, e = proc.communicate()
-        stdout = SPRetObj(o)
-        stderr = SPRetObj(e)
-        code = proc.returncode
-        if code != 0:
-            raise RuntimeError(
-                "Get error code %d in locally calling %s with job: %s ",
-                (code, cmd, self.submission.submission_hash),
-            )
-        return None, stdout, stderr
-
     def block_call(self, cmd):
         proc = sp.Popen(
             cmd, cwd=self.remote_root, shell=True, stdout=sp.PIPE, stderr=sp.PIPE
@@ -343,3 +349,31 @@ class LocalContext(BaseContext):
                 stdout = None
                 stderr = None
         return ret, stdout, stderr
+
+    @classmethod
+    def machine_subfields(cls) -> List[Argument]:
+        """Generate the machine subfields.
+
+        Returns
+        -------
+        list[Argument]
+            machine subfields
+        """
+        doc_remote_profile = "The information used to maintain the local machine."
+        return [
+            Argument(
+                "remote_profile",
+                dict,
+                optional=True,
+                doc=doc_remote_profile,
+                sub_fields=[
+                    Argument(
+                        "symlink",
+                        bool,
+                        optional=True,
+                        default=True,
+                        doc="Whether to use symbolic links to replace copy. This option should be turned off if the local directory is not accessible on the Batch system.",
+                    ),
+                ],
+            )
+        ]

dpdispatcher/contexts/openapi_context.py

@@ -1,18 +1,20 @@
+import glob
 import os
 import shutil
 import uuid
+from zipfile import ZipFile
 
 import tqdm
 
 try:
-    from bohriumsdk.client import Client
-    from bohriumsdk.job import Job
-    from bohriumsdk.storage import Storage
-    from bohriumsdk.util import Util
-except ModuleNotFoundError:
+    from bohrium import Bohrium
+    from bohrium.resources import Job, Tiefblue
+except ModuleNotFoundError as e:
     found_bohriumsdk = False
+    import_bohrium_error = e
 else:
     found_bohriumsdk = True
+    import_bohrium_error = None
 
 from dpdispatcher.base_context import BaseContext
 from dpdispatcher.dlog import dlog
@@ -23,6 +25,36 @@ DP_CLOUD_SERVER_HOME_DIR = os.path.join(
 )
 
 
+def unzip_file(zip_file, out_dir="./"):
+    obj = ZipFile(zip_file, "r")
+    for item in obj.namelist():
+        obj.extract(item, out_dir)
+
+
+def zip_file_list(root_path, zip_filename, file_list=[]):
+    out_zip_file = os.path.join(root_path, zip_filename)
+    # print('debug: file_list', file_list)
+    zip_obj = ZipFile(out_zip_file, "w")
+    for f in file_list:
+        matched_files = os.path.join(root_path, f)
+        for ii in glob.glob(matched_files):
+            # print('debug: matched_files:ii', ii)
+            if os.path.isdir(ii):
+                arcname = os.path.relpath(ii, start=root_path)
+                zip_obj.write(ii, arcname)
+                for root, dirs, files in os.walk(ii):
+                    for file in files:
+                        filename = os.path.join(root, file)
+                        arcname = os.path.relpath(filename, start=root_path)
+                        # print('debug: filename:arcname:root_path', filename, arcname, root_path)
+                        zip_obj.write(filename, arcname)
+            else:
+                arcname = os.path.relpath(ii, start=root_path)
+                zip_obj.write(ii, arcname)
+    zip_obj.close()
+    return out_zip_file
+
+
 class OpenAPIContext(BaseContext):
     def __init__(
         self,
@@ -35,16 +67,41 @@ class OpenAPIContext(BaseContext):
         if not found_bohriumsdk:
             raise ModuleNotFoundError(
                 "bohriumsdk not installed. Install dpdispatcher with `pip install dpdispatcher[bohrium]`"
-            )
+            ) from import_bohrium_error
         self.init_local_root = local_root
         self.init_remote_root = remote_root
         self.temp_local_root = os.path.abspath(local_root)
         self.remote_profile = remote_profile
-        self.client = Client()
-        self.storage = Storage(client=self.client)
-        self.job = Job(client=self.client)
-        self.util = Util()
+        access_key = (
+            remote_profile.get("access_key", None)
+            or os.getenv("BOHRIUM_ACCESS_KEY", None)
+            or os.getenv("ACCESS_KEY", None)
+        )
+        project_id = (
+            remote_profile.get("project_id", None)
+            or os.getenv("BOHRIUM_PROJECT_ID", None)
+            or os.getenv("PROJECT_ID", None)
+        )
+        app_key = (
+            remote_profile.get("app_key", None)
+            or os.getenv("BOHRIUM_APP_KEY", None)
+            or os.getenv("APP_KEY", None)
+        )
+        if access_key is None:
+            raise ValueError(
+                "remote_profile must contain 'access_key' or set environment variable 'BOHRIUM_ACCESS_KEY'"
+            )
+        if project_id is None:
+            raise ValueError(
+                "remote_profile must contain 'project_id' or set environment variable 'BOHRIUM_PROJECT_ID'"
+            )
+        self.client = Bohrium(  # type: ignore[reportPossiblyUnboundVariable]
+            access_key=access_key, project_id=project_id, app_key=app_key
+        )
+        self.storage = Tiefblue()  # type: ignore[reportPossiblyUnboundVariable]
+        self.job = Job(client=self.client)  # type: ignore[reportPossiblyUnboundVariable]
         self.jgid = None
+        os.makedirs(DP_CLOUD_SERVER_HOME_DIR, exist_ok=True)
 
     @classmethod
     def load_from_dict(cls, context_dict):
@@ -97,7 +154,7 @@ class OpenAPIContext(BaseContext):
             for file in task.forward_files:
                 upload_file_list.append(os.path.join(task.task_work_path, file))
 
-        upload_zip = Util.zip_file_list(
+        upload_zip = zip_file_list(
             self.local_root, zip_task_file, file_list=upload_file_list
         )
         project_id = self.remote_profile.get("project_id", 0)
@@ -113,7 +170,7 @@ class OpenAPIContext(BaseContext):
         object_key = os.path.join(data["storePath"], zip_filename)  # type: ignore
         job.upload_path = object_key
         job.job_id = data["jobId"]  # type: ignore
-        job.jgid = data["jobGroupId"]  # type: ignore
+        job.jgid = data.get("jobGroupId", "")  # type: ignore
         self.storage.upload_From_file_multi_part(
             object_key=object_key, file_path=upload_zip, token=token
         )
@@ -149,7 +206,9 @@ class OpenAPIContext(BaseContext):
         # return oss_task_zip
         # api.upload(self.oss_task_dir, zip_task_file)
 
-    def download(self, submission):
+    def download(
+        self, submission, check_exists=False, mark_failure=True, back_error=False
+    ):
         jobs = submission.belonging_jobs
         job_hashs = {}
         job_infos = {}
@@ -189,7 +248,7 @@ class OpenAPIContext(BaseContext):
             ):
                 continue
             self.storage.download_from_url(info["resultUrl"], target_result_zip)
-            Util.unzip_file(target_result_zip, out_dir=self.local_root)
+            unzip_file(target_result_zip, out_dir=self.local_root)
             self._backup(self.local_root, target_result_zip)
         self._clean_backup(
             self.local_root, keep_backup=self.remote_profile.get("keep_backup", True)
@@ -258,3 +317,8 @@ class OpenAPIContext(BaseContext):
             dir_to_be_removed = os.path.join(local_root, "backup")
             if os.path.exists(dir_to_be_removed):
                 shutil.rmtree(dir_to_be_removed)
+
+    def block_call(self, cmd):
+        raise RuntimeError(
+            "Unsupported method. You may use an unsupported combination of the machine and the context."
+        )