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

dpdispatcher/__init__.py

@@ -1,103 +1,21 @@
-import logging
-import os
-import sys
-import warnings
-
-ROOT_PATH = tuple(__path__)[0]
-dlog = logging.getLogger(__name__)
-dlog.propagate = False
-dlog.setLevel(logging.INFO)
-try:
-    dlogf = logging.FileHandler(
-        os.getcwd() + os.sep + "dpdispatcher" + ".log", delay=True
-    )
-except PermissionError:
-    warnings.warn(
-        "dpdispatcher.log meet permission error. redirect the log to ~/dpdispatcher.log"
-    )
-    dlogf = logging.FileHandler(
-        os.path.join(os.path.expanduser("~"), "dpdispatcher.log")
-    )
-
-# dlogf = logging.FileHandler('./'+os.sep+SHORT_CMD+'.log')
-# dlogf = logging.FileHandler(os.path.join(os.environ['HOME'], SHORT_CMD+'.log'))
-# dlogf = logging.FileHandler(os.path.join(os.path.expanduser('~'), SHORT_CMD+'.log'))
-# dlogf = logging.FileHandler(os.path.join("/tmp/", SHORT_CMD+'.log'))
-dlogf_formatter = logging.Formatter("%(asctime)s - %(levelname)s : %(message)s")
-# dlogf_formatter=logging.Formatter('%(asctime)s - %(name)s - [%(filename)s:%(funcName)s - %(lineno)d ] - %(levelname)s \n %(message)s')
-dlogf.setFormatter(dlogf_formatter)
-dlog.addHandler(dlogf)
-
-dlog_stdout = logging.StreamHandler(sys.stdout)
-dlog_stdout.setFormatter(dlogf_formatter)
-dlog.addHandler(dlog_stdout)
-
 __author__ = "DeepModeling Team"
-__copyright__ = "Copyright 2019"
-__status__ = "Development"
+__copyright__ = "Copyright 2019-2023, DeepModeling"
+__status__ = "Production"
 try:
     from ._version import version as __version__
 except ImportError:
-    __version__ = "unkown"
-
-from .distributed_shell import DistributedShell
-from .dp_cloud_server import DpCloudServer, Lebesgue
-from .dp_cloud_server_context import DpCloudServerContext, LebesgueContext
-from .fugaku import Fugaku
-from .hdfs_context import HDFSContext
-from .lazy_local_context import LazyLocalContext
-from .local_context import LocalContext
-from .lsf import LSF
-from .machine import Machine
-from .openapi import OpenAPI
-from .openapi_context import OpenAPIContext
-from .pbs import PBS, Torque
-from .shell import Shell
-from .slurm import Slurm
-from .ssh_context import SSHContext
-from .submission import Job, Resources, Submission, Task
-
-
-def info():
-    """Show basic information about dpdispatcher, its location and version."""
-    print("DeepModeling\n------------")
-    print("Version: " + __version__)
-    print("Path:    " + ROOT_PATH)
-    print("")
-    print("Dependency")
-    print("------------")
-    for modui in ["psutil", "paramiko", "dargs", "oss2"]:
-        try:
-            mm = __import__(modui)
-            print("%10s %10s   %s" % (modui, mm.__version__, mm.__path__[0]))
-        except ImportError:
-            print("%10s %10s Not Found" % (modui, ""))
-    print()
+    __version__ = "unknown"
 
+import dpdispatcher.contexts  # noqa: F401
+import dpdispatcher.machines  # noqa: F401
+from dpdispatcher.machine import Machine
+from dpdispatcher.submission import Job, Resources, Submission, Task
 
 __all__ = [
     "__version__",
-    "DistributedShell",
-    "DpCloudServer",
-    "OpenAPI",
-    "OpenAPIContext",
-    "DpCloudServerContext",
-    "HDFSContext",
-    "LazyLocalContext",
-    "LocalContext",
-    "LSF",
     "Machine",
-    "PBS",
-    "Shell",
-    "Slurm",
-    "Fugaku",
-    "SSHContext",
     "Submission",
     "Task",
-    "Torque",
-    "info",
-    "Lebesgue",
-    "LebesgueContext",
     "Job",
     "Resources",
 ]

dpdispatcher/__main__.py

@@ -0,0 +1,8 @@
+"""Package dp entry point."""
+
+from dpdispatcher.dpdisp import (
+    main,
+)
+
+if __name__ == "__main__":
+    main()

dpdispatcher/_version.py

@@ -1,4 +1,16 @@
 # file generated by setuptools_scm
 # don't change, don't track in version control
-__version__ = version = '0.5.11'
-__version_tuple__ = version_tuple = (0, 5, 11)
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple, Union
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.6.1'
+__version_tuple__ = version_tuple = (0, 6, 1)

dpdispatcher/base_context.py

@@ -3,7 +3,7 @@ from typing import List, Tuple
 
 from dargs import Argument
 
-from dpdispatcher import dlog
+from dpdispatcher.dlog import dlog
 
 
 class BaseContext(metaclass=ABCMeta):

dpdispatcher/contexts/__init__.py

@@ -0,0 +1,11 @@
+"""Contexts."""
+import importlib
+from pathlib import Path
+
+PACKAGE_BASE = "dpdispatcher.contexts"
+NOT_LOADABLE = ("__init__.py",)
+
+for module_file in Path(__file__).parent.glob("*.py"):
+    if module_file.name not in NOT_LOADABLE:
+        module_name = f".{module_file.stem}"
+        importlib.import_module(module_name, PACKAGE_BASE)

dpdispatcher/{dp_cloud_server_context.py → contexts/dp_cloud_server_context.py}

@@ -8,13 +8,16 @@ from typing import List
 import tqdm
 from dargs.dargs import Argument
 
-from dpdispatcher import dlog
 from dpdispatcher.base_context import BaseContext
-from dpdispatcher.dpcloudserver.config import ALI_STS_BUCKET_NAME, ALI_STS_ENDPOINT
+from dpdispatcher.dlog import dlog
 
 # from dpdispatcher.submission import Machine
 # from . import dlog
-from .dpcloudserver import Client, zip_file
+from dpdispatcher.utils.dpcloudserver import Client, zip_file
+from dpdispatcher.utils.dpcloudserver.config import (
+    ALI_STS_BUCKET_NAME,
+    ALI_STS_ENDPOINT,
+)
 
 # from zip_file import zip_files
 
@@ -113,6 +116,7 @@ class BohriumContext(BaseContext):
 
         upload_file_list = [
             job.script_file_name,
+            f"{job.script_file_name}.run",
         ]
         upload_file_list.extend(common_files)
 

dpdispatcher/{hdfs_context.py → contexts/hdfs_context.py}

@@ -3,9 +3,9 @@ import shutil
 import tarfile
 from glob import glob
 
-from dpdispatcher import dlog
 from dpdispatcher.base_context import BaseContext
-from dpdispatcher.hdfs_cli import HDFS
+from dpdispatcher.dlog import dlog
+from dpdispatcher.utils.hdfs_cli import HDFS
 
 
 class HDFSContext(BaseContext):

dpdispatcher/{local_context.py → contexts/local_context.py}

@@ -1,12 +1,11 @@
-import hashlib
 import os
 import shutil
 import subprocess as sp
 from glob import glob
 from subprocess import TimeoutExpired
 
-from dpdispatcher import dlog
 from dpdispatcher.base_context import BaseContext
+from dpdispatcher.dlog import dlog
 
 
 class SPRetObj:
@@ -30,14 +29,6 @@ def _check_file_path(fname):
         os.makedirs(dirname, exist_ok=True)
 
 
-def _identical_files(fname0, fname1):
-    with open(fname0) as fp:
-        code0 = hashlib.sha1(fp.read().encode("utf-8")).hexdigest()
-    with open(fname1) as fp:
-        code1 = hashlib.sha1(fp.read().encode("utf-8")).hexdigest()
-    return code0 == code1
-
-
 class LocalContext(BaseContext):
     """Run jobs in the local server and remote directory.
 
@@ -152,9 +143,34 @@ class LocalContext(BaseContext):
         for ii in submission.belonging_tasks:
             local_job = os.path.join(self.local_root, ii.task_work_path)
             remote_job = os.path.join(self.remote_root, ii.task_work_path)
-            flist = ii.backward_files
+            flist = []
+            for kk in ii.backward_files:
+                abs_flist_r = glob(os.path.join(remote_job, kk))
+                abs_flist_l = glob(os.path.join(local_job, kk))
+                if not abs_flist_r and not abs_flist_l:
+                    if check_exists:
+                        if mark_failure:
+                            tag_file_path = os.path.join(
+                                self.local_root,
+                                ii.task_work_path,
+                                "tag_failure_download_%s" % kk,
+                            )
+                            with open(tag_file_path, "w") as fp:
+                                pass
+                        else:
+                            pass
+                    else:
+                        raise RuntimeError(
+                            "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
+                ]
+                flist.extend(rel_flist)
             if back_error:
-                flist += glob(os.path.join(remote_job, "error*"))
+                abs_flist = glob(os.path.join(remote_job, "error*"))
+                rel_flist = [os.path.relpath(ii, start=remote_job) for ii in abs_flist]
+                flist.extend(rel_flist)
             for jj in flist:
                 rfile = os.path.join(remote_job, jj)
                 lfile = os.path.join(local_job, jj)
@@ -198,9 +214,30 @@ class LocalContext(BaseContext):
                     pass
         local_job = self.local_root
         remote_job = self.remote_root
-        flist = submission.backward_common_files
+        flist = []
+        for kk in submission.backward_common_files:
+            abs_flist_r = glob(os.path.join(remote_job, kk))
+            abs_flist_l = glob(os.path.join(local_job, kk))
+            if not abs_flist_r and not abs_flist_l:
+                if check_exists:
+                    if mark_failure:
+                        tag_file_path = os.path.join(
+                            self.local_root, "tag_failure_download_%s" % kk
+                        )
+                        with open(tag_file_path, "w") as fp:
+                            pass
+                    else:
+                        pass
+                else:
+                    raise RuntimeError(
+                        "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]
+            flist.extend(rel_flist)
         if back_error:
-            flist += glob(os.path.join(remote_job, "error*"))
+            abs_flist = glob(os.path.join(remote_job, "error*"))
+            rel_flist = [os.path.relpath(ii, start=remote_job) for ii in abs_flist]
+            flist.extend(rel_flist)
         for jj in flist:
             rfile = os.path.join(remote_job, jj)
             lfile = os.path.join(local_job, jj)

dpdispatcher/{openapi_context.py → contexts/openapi_context.py}

@@ -14,9 +14,9 @@ except ModuleNotFoundError:
 else:
     found_bohriumsdk = True
 
-from dpdispatcher import dlog
 from dpdispatcher.base_context import BaseContext
-from dpdispatcher.JobStatus import JobStatus
+from dpdispatcher.dlog import dlog
+from dpdispatcher.utils.job_status import JobStatus
 
 DP_CLOUD_SERVER_HOME_DIR = os.path.join(
     os.path.expanduser("~"), ".dpdispatcher/", "dp_cloud_server/"
@@ -88,6 +88,7 @@ class OpenAPIContext(BaseContext):
 
         upload_file_list = [
             job.script_file_name,
+            f"{job.script_file_name}.run",
         ]
 
         upload_file_list.extend(common_files)

dpdispatcher/{ssh_context.py → contexts/ssh_context.py}

@@ -1,5 +1,6 @@
 #!/usr/bin/env python
 
+import fnmatch
 import os
 import pathlib
 import shlex
@@ -10,17 +11,24 @@ import time
 import uuid
 from functools import lru_cache
 from glob import glob
+from stat import S_ISDIR, S_ISREG
 from typing import List
 
 import paramiko
 import paramiko.ssh_exception
 from dargs.dargs import Argument
 
-from dpdispatcher import dlog
 from dpdispatcher.base_context import BaseContext
+from dpdispatcher.dlog import dlog
 
 # from dpdispatcher.submission import Machine
-from dpdispatcher.utils import RetrySignal, generate_totp, get_sha256, retry, rsync
+from dpdispatcher.utils.utils import (
+    RetrySignal,
+    generate_totp,
+    get_sha256,
+    retry,
+    rsync,
+)
 
 
 class SSHSession:
@@ -48,6 +56,7 @@ class SSHSession:
         self.ssh = None
         self.tar_compress = tar_compress
         self.look_for_keys = look_for_keys
+        self._keyboard_interactive_auth = False
         self._setup_ssh()
 
     # @classmethod
@@ -137,6 +146,7 @@ class SSHSession:
         # Make a Paramiko Transport object using the socket
         ts = paramiko.Transport(sock)
         ts.banner_timeout = 60
+        ts.auth_timeout = self.timeout + 20
         ts.use_compression(compress=True)
 
         # Tell Paramiko that the Transport is going to be used as a client
@@ -165,6 +175,8 @@ class SSHSession:
                         pass
                     if key is not None:
                         break
+            else:
+                raise OSError(f"{key_path} not found!")
         elif self.look_for_keys:
             for keytype, name in [
                 (paramiko.RSAKey, "rsa"),
@@ -187,20 +199,22 @@ class SSHSession:
                 if key is not None:
                     break
 
+        allowed_types = set()
         if key is not None:
             try:
-                ts.auth_publickey(self.username, key)
+                allowed_types = set(ts.auth_publickey(self.username, key))
             except paramiko.ssh_exception.AuthenticationException as e:
                 key_error = e
             else:
                 key_ok = True
-        if self.totp_secret is not None:
+        if self.totp_secret is not None or "keyboard-interactive" in allowed_types:
             try:
                 ts.auth_interactive(self.username, self.inter_handler)
             except paramiko.ssh_exception.AuthenticationException:
                 # since the asynchrony of interactive authentication, one addtional try is added
                 # retry for up to 6 times
                 raise RetrySignal("Authentication failed")
+            self._keyboard_interactive_auth = True
         elif key_ok:
             pass
         elif self.password is not None:
@@ -390,6 +404,7 @@ class SSHSession:
             and self.exec_command("rsync --version")[1].channel.recv_exit_status() == 0
             and self.totp_secret is None
             and self.passphrase is None
+            and not self._keyboard_interactive_auth
         )
 
     @property
@@ -542,6 +557,8 @@ class SSHContext(BaseContext):
                         directory_list.append(root)
                     for name in files:
                         file_list.append(os.path.join(root, name))
+            elif os.path.islink(file_name) and not os.path.exists(file_name):
+                raise OSError(f"{file_name} is broken symbolic link")
             elif glob(file_name):
                 # If the file name contains a wildcard, os.path functions will fail to identify it. Use glob to get the complete list of filenames which match the wildcard.
                 abs_file_list = glob(file_name)
@@ -632,6 +649,18 @@ class SSHContext(BaseContext):
             tar_compress=self.remote_profile.get("tar_compress", None),
         )
 
+    def list_remote_dir(self, sftp, remote_dir, ref_remote_root, result_list):
+        for entry in sftp.listdir_attr(remote_dir):
+            remote_name = pathlib.PurePath(
+                os.path.join(remote_dir, entry.filename)
+            ).as_posix()
+            st_mode = entry.st_mode
+            if S_ISDIR(st_mode):
+                self.list_remote_dir(sftp, remote_name, ref_remote_root, result_list)
+            elif S_ISREG(st_mode):
+                rel_remote_name = os.path.relpath(remote_name, start=ref_remote_root)
+                result_list.append(rel_remote_name)
+
     def download(
         self,
         submission,
@@ -641,34 +670,70 @@ class SSHContext(BaseContext):
         mark_failure=True,
         back_error=False,
     ):
+        assert self.remote_root is not None
         self.ssh_session.ensure_alive()
         file_list = []
         # for ii in job_dirs :
-        for task in submission.belonging_tasks:
-            for jj in task.backward_files:
-                file_name = pathlib.PurePath(
-                    os.path.join(task.task_work_path, jj)
-                ).as_posix()
+        for ii in submission.belonging_tasks:
+            remote_file_list = None
+            for jj in ii.backward_files:
+                if "*" in jj or "?" in jj:
+                    if remote_file_list is not None:
+                        abs_file_list = fnmatch.filter(remote_file_list, jj)
+                    else:
+                        remote_file_list = []
+                        remote_job = pathlib.PurePath(
+                            os.path.join(self.remote_root, ii.task_work_path)
+                        ).as_posix()
+                        self.list_remote_dir(
+                            self.sftp, remote_job, remote_job, remote_file_list
+                        )
+
+                        abs_file_list = fnmatch.filter(remote_file_list, jj)
+                    rel_file_list = [
+                        pathlib.PurePath(os.path.join(ii.task_work_path, kk)).as_posix()
+                        for kk in abs_file_list
+                    ]
+
+                else:
+                    rel_file_list = [
+                        pathlib.PurePath(os.path.join(ii.task_work_path, jj)).as_posix()
+                    ]
                 if check_exists:
-                    if self.check_file_exists(file_name):
-                        file_list.append(file_name)
-                    elif mark_failure:
-                        with open(
-                            os.path.join(
-                                self.local_root,
-                                task.task_work_path,
-                                "tag_failure_download_%s" % jj,
-                            ),
-                            "w",
-                        ) as fp:
+                    for file_name in rel_file_list:
+                        if self.check_file_exists(file_name):
+                            file_list.append(file_name)
+                        elif mark_failure:
+                            with open(
+                                os.path.join(
+                                    self.local_root,
+                                    ii.task_work_path,
+                                    "tag_failure_download_%s" % jj,
+                                ),
+                                "w",
+                            ) as fp:
+                                pass
+                        else:
                             pass
-                    else:
-                        pass
                 else:
-                    file_list.append(file_name)
+                    file_list.extend(rel_file_list)
             if back_error:
-                errors = glob(os.path.join(task.task_work_path, "error*"))
-                file_list.extend(errors)
+                if remote_file_list is not None:
+                    abs_errors = fnmatch.filter(remote_file_list, "error*")
+                else:
+                    remote_file_list = []
+                    remote_job = pathlib.PurePath(
+                        os.path.join(self.remote_root, ii.task_work_path)
+                    ).as_posix()
+                    self.list_remote_dir(
+                        self.sftp, remote_job, remote_job, remote_file_list
+                    )
+                    abs_errors = fnmatch.filter(remote_file_list, "error*")
+                rel_errors = [
+                    pathlib.PurePath(os.path.join(ii.task_work_path, kk)).as_posix()
+                    for kk in abs_errors
+                ]
+                file_list.extend(rel_errors)
         file_list.extend(submission.backward_common_files)
         if len(file_list) > 0:
             self._get_files(
@@ -881,21 +946,35 @@ class SSHContext(BaseContext):
         per_nfile = 100
         ntar = len(files) // per_nfile + 1
         if ntar <= 1:
-            self.block_checkcall(
-                "tar {} {} {}".format(
-                    tar_command,
-                    shlex.quote(of),
-                    " ".join([shlex.quote(file) for file in files]),
+            try:
+                self.block_checkcall(
+                    "tar {} {} {}".format(
+                        tar_command,
+                        shlex.quote(of),
+                        " ".join([shlex.quote(file) for file in files]),
+                    )
                 )
-            )
+            except RuntimeError as e:
+                if "No such file or directory" in str(e):
+                    raise FileNotFoundError(
+                        "Any of the backward files does not exist in the remote directory."
+                    ) from e
+                raise e
         else:
             file_list_file = os.path.join(
                 self.remote_root, ".tmp.tar." + str(uuid.uuid4())
             )
             self.write_file(file_list_file, "\n".join(files))
-            self.block_checkcall(
-                f"tar {tar_command} {shlex.quote(of)} -T {shlex.quote(file_list_file)}"
-            )
+            try:
+                self.block_checkcall(
+                    f"tar {tar_command} {shlex.quote(of)} -T {shlex.quote(file_list_file)}"
+                )
+            except RuntimeError as e:
+                if "No such file or directory" in str(e):
+                    raise FileNotFoundError(
+                        "Any of the backward files does not exist in the remote directory."
+                    ) from e
+                raise e
         # trans
         from_f = pathlib.PurePath(os.path.join(self.remote_root, of)).as_posix()
         to_f = pathlib.PurePath(os.path.join(self.local_root, of)).as_posix()

dpdispatcher/dlog.py

@@ -0,0 +1,31 @@
+import logging
+import os
+import sys
+import warnings
+
+dlog = logging.getLogger("dpdispatcher")
+dlog.propagate = False
+dlog.setLevel(logging.INFO)
+try:
+    dlogf = logging.FileHandler(
+        os.getcwd() + os.sep + "dpdispatcher" + ".log", delay=True
+    )
+except PermissionError:
+    warnings.warn(
+        "dpdispatcher.log meet permission error. redirect the log to ~/dpdispatcher.log"
+    )
+    dlogf = logging.FileHandler(
+        os.path.join(os.path.expanduser("~"), "dpdispatcher.log"), delay=True
+    )
+
+dlogf_formatter = logging.Formatter("%(asctime)s - %(levelname)s : %(message)s")
+dlogf.setFormatter(dlogf_formatter)
+dlog.addHandler(dlogf)
+
+dlog_stdout = logging.StreamHandler(sys.stdout)
+dlog_stdout.setFormatter(dlogf_formatter)
+dlog.addHandler(dlog_stdout)
+
+__all__ = [
+    "dlog",
+]