dpdispatcher 0.6.0__tar.gz → 0.6.1__tar.gz

{dpdispatcher-0.6.0 → dpdispatcher-0.6.1}/.github/workflows/test.yml

@@ -27,7 +27,11 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - run: pip install .[test] coverage
     - name: Test
-      run: coverage run --source=./dpdispatcher -m unittest -v && coverage report
+      run: |
+        coverage run -p --source=./dpdispatcher -m unittest -v
+        coverage run -p --source=./dpdispatcher -m dpdispatcher -h
+        coverage combine
+        coverage report
     - uses: codecov/codecov-action@v3
   pass:
     needs: [test]

{dpdispatcher-0.6.0 → dpdispatcher-0.6.1}/.pre-commit-config.yaml

@@ -16,16 +16,13 @@ repos:
     -   id: check-symlinks
     -   id: check-toml
 # Python
--   repo: https://github.com/psf/black
-    rev: 23.9.1
-    hooks:
-    -   id: black-jupyter
 -   repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.0.292
+    rev: v0.1.5
     hooks:
     - id: ruff
       args: ["--fix"]
+    - id: ruff-format
 # numpydoc
 -   repo: https://github.com/Carreau/velin
     rev: 0.0.12

{dpdispatcher-0.6.0 → dpdispatcher-0.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dpdispatcher
-Version: 0.6.0
+Version: 0.6.1
 Summary: Generate HPC scheduler systems jobs input scripts, submit these scripts to HPC systems, and poke until they finish
 Author: DeepModeling
 License:                    GNU LESSER GENERAL PUBLIC LICENSE
@@ -186,10 +186,11 @@ Requires-Python: >=3.7
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: paramiko
-Requires-Dist: dargs>=0.2.9
+Requires-Dist: dargs>=0.4.1
 Requires-Dist: requests
 Requires-Dist: tqdm>=4.9.0
 Requires-Dist: typing_extensions; python_version < "3.7"
+Requires-Dist: pyyaml
 Provides-Extra: docs
 Requires-Dist: sphinx; extra == "docs"
 Requires-Dist: myst-parser; extra == "docs"

dpdispatcher-0.6.1/doc/examples/template.md

@@ -0,0 +1,16 @@
+# Customizing the submission script header
+
+When submitting jobs to some clusters, such as the [Tiger Cluster](https://researchcomputing.princeton.edu/systems/tiger) at Princeton University, the Slurm header is quite different from the standard one. In this case, DPDispatcher allows users to customize the templates by setting {dargs:argument}`strategy/customized_script_header_template_file <resources/strategy/customized_script_header_template_file>` to a template file:
+
+```{literalinclude} ../../examples/resources/tiger.json
+:language: json
+:linenos:
+```
+
+`template.slurm` is the template file, where {meth}`str.format` is used to format the template with [Resources Parameters](resources):
+
+```{literalinclude} ../../examples/resources/template.slurm
+:linenos:
+```
+
+See [Python Format String Syntax](https://docs.python.org/3/library/string.html#formatstrings) for how to insert parameters inside the template.

{dpdispatcher-0.6.0 → dpdispatcher-0.6.1}/doc/index.rst

@@ -29,7 +29,10 @@ DPDispatcher will monitor (poke) until these jobs finish and download the result
    :caption: Examples
    :glob:
 
-   examples/*
+   examples/expanse
+   examples/g16
+   examples/shell
+   examples/template
 
 .. toctree::
    :caption: Project details

dpdispatcher-0.6.1/dpdispatcher/__init__.py

@@ -0,0 +1,21 @@
+__author__ = "DeepModeling Team"
+__copyright__ = "Copyright 2019-2023, DeepModeling"
+__status__ = "Production"
+try:
+    from ._version import version as __version__
+except ImportError:
+    __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__",
+    "Machine",
+    "Submission",
+    "Task",
+    "Job",
+    "Resources",
+]

dpdispatcher-0.6.1/dpdispatcher/__main__.py

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

{dpdispatcher-0.6.0 → dpdispatcher-0.6.1}/dpdispatcher/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.6.0'
-__version_tuple__ = version_tuple = (0, 6, 0)
+__version__ = version = '0.6.1'
+__version_tuple__ = version_tuple = (0, 6, 1)

{dpdispatcher-0.6.0 → dpdispatcher-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-0.6.1/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-0.6.0/dpdispatcher → dpdispatcher-0.6.1/dpdispatcher/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
 

{dpdispatcher-0.6.0/dpdispatcher → dpdispatcher-0.6.1/dpdispatcher/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-0.6.0/dpdispatcher → dpdispatcher-0.6.1/dpdispatcher/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.
 

{dpdispatcher-0.6.0/dpdispatcher → dpdispatcher-0.6.1/dpdispatcher/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/"

{dpdispatcher-0.6.0/dpdispatcher → dpdispatcher-0.6.1/dpdispatcher/contexts}/ssh_context.py

@@ -18,11 +18,17 @@ 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:
@@ -140,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
@@ -939,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-0.6.1/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",
+]

{dpdispatcher-0.6.0 → dpdispatcher-0.6.1}/dpdispatcher/dpdisp.py

@@ -2,7 +2,8 @@
 import argparse
 from typing import List, Optional
 
-from dpdispatcher.gui import start_dpgui
+from dpdispatcher.entrypoints.gui import start_dpgui
+from dpdispatcher.entrypoints.submission import handle_submission
 
 
 def main_parser() -> argparse.ArgumentParser:
@@ -23,6 +24,37 @@ def main_parser() -> argparse.ArgumentParser:
     )
     subparsers = parser.add_subparsers(title="Valid subcommands", dest="command")
     ##########################################
+    # backward
+    parser_submission = subparsers.add_parser(
+        "submission",
+        help="Handle terminated submission.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser_submission.add_argument(
+        "SUBMISSION_HASH",
+        type=str,
+        help="Submission hash to download.",
+    )
+    parser_submission_action = parser_submission.add_argument_group(
+        "Actions",
+        description="One or more actions to take on submission.",
+    )
+    parser_submission_action.add_argument(
+        "--download-terminated-log",
+        action="store_true",
+        help="Download log files of terminated tasks.",
+    )
+    parser_submission_action.add_argument(
+        "--download-finished-task",
+        action="store_true",
+        help="Download finished tasks.",
+    )
+    parser_submission_action.add_argument(
+        "--clean",
+        action="store_true",
+        help="Clean submission.",
+    )
+    ##########################################
     # gui
     parser_gui = subparsers.add_parser(
         "gui",
@@ -67,7 +99,14 @@ def parse_args(args: Optional[List[str]] = None):
 
 def main():
     args = parse_args()
-    if args.command == "gui":
+    if args.command == "submission":
+        handle_submission(
+            submission_hash=args.SUBMISSION_HASH,
+            download_terminated_log=args.download_terminated_log,
+            download_finished_task=args.download_finished_task,
+            clean=args.clean,
+        )
+    elif args.command == "gui":
         start_dpgui(
             port=args.port,
             bind_all=args.bind_all,

dpdispatcher-0.6.1/dpdispatcher/entrypoints/__init__.py

@@ -0,0 +1 @@
+"""Entry points."""

dpdispatcher-0.6.1/dpdispatcher/entrypoints/submission.py

@@ -0,0 +1,83 @@
+from pathlib import Path
+
+from dpdispatcher.dlog import dlog
+from dpdispatcher.submission import Submission
+from dpdispatcher.utils.job_status import JobStatus
+from dpdispatcher.utils.record import record
+
+
+def handle_submission(
+    *,
+    submission_hash: str,
+    download_terminated_log: bool = False,
+    download_finished_task: bool = False,
+    clean: bool = False,
+):
+    """Handle terminated submission.
+
+    Parameters
+    ----------
+    submission_hash : str
+        Submission hash to download.
+    download_terminated_log : bool, optional
+        Download log files of terminated tasks.
+    download_finished_task : bool, optional
+        Download finished tasks.
+    clean : bool, optional
+        Clean submission.
+
+    Raises
+    ------
+    ValueError
+        At least one action should be specified.
+    """
+    if int(download_terminated_log) + int(download_finished_task) + int(clean) == 0:
+        raise ValueError("At least one action should be specified.")
+
+    submission_file = record.get_submission(submission_hash)
+    submission = Submission.submission_from_json(str(submission_file))
+    submission.belonging_tasks = [
+        task for job in submission.belonging_jobs for task in job.job_task_list
+    ]
+    # TODO: for unclear reason, the submission_hash may be changed
+    submission.submission_hash = submission_hash
+    submission.machine.context.bind_submission(submission)
+    submission.update_submission_state()
+
+    terminated_tasks = []
+    finished_tasks = []
+    for task in submission.belonging_tasks:
+        task.get_task_state(submission.machine.context)
+        if task.task_state == JobStatus.terminated:
+            terminated_tasks.append(task)
+        elif task.task_state == JobStatus.finished:
+            finished_tasks.append(task)
+    submission.belonging_tasks = []
+
+    if download_terminated_log:
+        for task in terminated_tasks:
+            task.backward_files = [task.outlog, task.errlog]
+        submission.belonging_tasks += terminated_tasks
+    if download_finished_task:
+        submission.belonging_tasks += finished_tasks
+
+    submission.download_jobs()
+
+    if download_terminated_log:
+        terminated_log_files = []
+        for task in terminated_tasks:
+            assert submission.local_root is not None
+            terminated_log_files.append(
+                Path(submission.local_root) / task.task_work_path / task.outlog
+            )
+            terminated_log_files.append(
+                Path(submission.local_root) / task.task_work_path / task.errlog
+            )
+
+        dlog.info(
+            "Terminated logs are downloaded into:\n  "
+            + "\n  ".join([str(f) for f in terminated_log_files])
+        )
+
+    if clean:
+        submission.clean_jobs()

{dpdispatcher-0.6.0 → dpdispatcher-0.6.1}/dpdispatcher/machine.py

@@ -4,10 +4,11 @@ import shlex
 from abc import ABCMeta, abstractmethod
 from typing import List, Tuple
 
+import yaml
 from dargs import Argument, Variant
 
-from dpdispatcher import dlog
 from dpdispatcher.base_context import BaseContext
+from dpdispatcher.dlog import dlog
 
 script_template = """\
 {script_header}
@@ -124,6 +125,13 @@ class Machine(metaclass=ABCMeta):
         machine = cls.load_from_dict(machine_dict=machine_dict)
         return machine
 
+    @classmethod
+    def load_from_yaml(cls, yaml_path):
+        with open(yaml_path) as f:
+            machine_dict = yaml.safe_load(f)
+        machine = cls.load_from_dict(machine_dict=machine_dict)
+        return machine
+
     @classmethod
     def load_from_dict(cls, machine_dict):
         batch_type = machine_dict["batch_type"]

dpdispatcher-0.6.1/dpdispatcher/machines/__init__.py

@@ -0,0 +1,11 @@
+"""Machines."""
+import importlib
+from pathlib import Path
+
+PACKAGE_BASE = "dpdispatcher.machines"
+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-0.6.0/dpdispatcher → dpdispatcher-0.6.1/dpdispatcher/machines}/distributed_shell.py

@@ -1,7 +1,7 @@
-from dpdispatcher import dlog
-from dpdispatcher.JobStatus import JobStatus
+from dpdispatcher.dlog import dlog
 from dpdispatcher.machine import Machine
-from dpdispatcher.utils import (
+from dpdispatcher.utils.job_status import JobStatus
+from dpdispatcher.utils.utils import (
     customized_script_header_template,
     run_cmd_with_all_output,
 )

{dpdispatcher-0.6.0/dpdispatcher → dpdispatcher-0.6.1/dpdispatcher/machines}/dp_cloud_server.py

@@ -4,12 +4,12 @@ import time
 import uuid
 import warnings
 
-from dpdispatcher import dlog
-from dpdispatcher.dpcloudserver import Client, zip_file
-from dpdispatcher.dpcloudserver.config import ALI_OSS_BUCKET_URL
-from dpdispatcher.JobStatus import JobStatus
+from dpdispatcher.dlog import dlog
 from dpdispatcher.machine import Machine
-from dpdispatcher.utils import customized_script_header_template
+from dpdispatcher.utils.dpcloudserver import Client, zip_file
+from dpdispatcher.utils.dpcloudserver.config import ALI_OSS_BUCKET_URL
+from dpdispatcher.utils.job_status import JobStatus
+from dpdispatcher.utils.utils import customized_script_header_template
 
 shell_script_header_template = """
 #!/bin/bash -l

{dpdispatcher-0.6.0/dpdispatcher → dpdispatcher-0.6.1/dpdispatcher/machines}/fugaku.py

@@ -1,9 +1,9 @@
 import shlex
 
-from dpdispatcher import dlog
-from dpdispatcher.JobStatus import JobStatus
+from dpdispatcher.dlog import dlog
 from dpdispatcher.machine import Machine
-from dpdispatcher.utils import customized_script_header_template
+from dpdispatcher.utils.job_status import JobStatus
+from dpdispatcher.utils.utils import customized_script_header_template
 
 fugaku_script_header_template = """\
 {queue_name_line}

{dpdispatcher-0.6.0/dpdispatcher → dpdispatcher-0.6.1/dpdispatcher/machines}/lsf.py

@@ -3,10 +3,14 @@ from typing import List
 
 from dargs import Argument
 
-from dpdispatcher import dlog
-from dpdispatcher.JobStatus import JobStatus
+from dpdispatcher.dlog import dlog
 from dpdispatcher.machine import Machine
-from dpdispatcher.utils import RetrySignal, customized_script_header_template, retry
+from dpdispatcher.utils.job_status import JobStatus
+from dpdispatcher.utils.utils import (
+    RetrySignal,
+    customized_script_header_template,
+    retry,
+)
 
 lsf_script_header_template = """\
 #!/bin/bash -l

{dpdispatcher-0.6.0/dpdispatcher → dpdispatcher-0.6.1/dpdispatcher/machines}/openapi.py

@@ -2,7 +2,7 @@ import os
 import shutil
 import time
 
-from dpdispatcher.utils import customized_script_header_template
+from dpdispatcher.utils.utils import customized_script_header_template
 
 try:
     from bohriumsdk.client import Client
@@ -14,9 +14,9 @@ except ModuleNotFoundError:
 else:
     found_bohriumsdk = True
 
-from dpdispatcher import dlog
-from dpdispatcher.JobStatus import JobStatus
+from dpdispatcher.dlog import dlog
 from dpdispatcher.machine import Machine
+from dpdispatcher.utils.job_status import JobStatus
 
 shell_script_header_template = """
 #!/bin/bash -l
@@ -147,7 +147,9 @@ class OpenAPI(Machine):
                 )
 
         job_state = self.map_dp_job_state(
-            dp_job_status, check_return.get("exitCode", 0), self.ignore_exit_code  # type: ignore
+            dp_job_status,
+            check_return.get("exitCode", 0),  # type: ignore
+            self.ignore_exit_code,
         )
         if job_state == JobStatus.finished:
             job_log = self.job.log(job_id)

{dpdispatcher-0.6.0/dpdispatcher → dpdispatcher-0.6.1/dpdispatcher/machines}/pbs.py

@@ -1,9 +1,9 @@
 import shlex
 
-from dpdispatcher import dlog
-from dpdispatcher.JobStatus import JobStatus
+from dpdispatcher.dlog import dlog
 from dpdispatcher.machine import Machine
-from dpdispatcher.utils import customized_script_header_template
+from dpdispatcher.utils.job_status import JobStatus
+from dpdispatcher.utils.utils import customized_script_header_template
 
 pbs_script_header_template = """
 #!/bin/bash -l

{dpdispatcher-0.6.0/dpdispatcher → dpdispatcher-0.6.1/dpdispatcher/machines}/shell.py

@@ -1,9 +1,9 @@
 import shlex
 
-from dpdispatcher import dlog
-from dpdispatcher.JobStatus import JobStatus
+from dpdispatcher.dlog import dlog
 from dpdispatcher.machine import Machine
-from dpdispatcher.utils import customized_script_header_template
+from dpdispatcher.utils.job_status import JobStatus
+from dpdispatcher.utils.utils import customized_script_header_template
 
 shell_script_header_template = """
 #!/bin/bash -l