dpdispatcher 0.5.6__py3-none-any.whl → 0.5.8__py3-none-any.whl

dpdispatcher/__init__.py

@@ -43,6 +43,7 @@ except ImportError:
 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
@@ -85,6 +86,7 @@ __all__ = [
     "PBS",
     "Shell",
     "Slurm",
+    "Fugaku",
     "SSHContext",
     "Submission",
     "Task",

dpdispatcher/_version.py

@@ -1,4 +1,4 @@
 # file generated by setuptools_scm
 # don't change, don't track in version control
-__version__ = version = '0.5.6'
-__version_tuple__ = version_tuple = (0, 5, 6)
+__version__ = version = '0.5.8'
+__version_tuple__ = version_tuple = (0, 5, 8)

dpdispatcher/base_context.py

@@ -70,9 +70,6 @@ class BaseContext(metaclass=ABCMeta):
     def read_file(self, fname):
         raise NotImplementedError("abstract method")
 
-    def kill(self, proc):
-        raise NotImplementedError("abstract method")
-
     def check_finish(self, proc):
         raise NotImplementedError("abstract method")
 

dpdispatcher/distributed_shell.py

@@ -136,17 +136,16 @@ class DistributedShell(Machine):
 
         resources = job.resources
         submit_command = (
-            "hadoop jar %s/hadoop-yarn-applications-distributedshell-*.jar "
+            "hadoop jar {}/hadoop-yarn-applications-distributedshell-*.jar "
             "org.apache.hadoop.yarn.applications.distributedshell.Client "
-            "-jar %s/hadoop-yarn-applications-distributedshell-*.jar "
-            '-queue %s -appname "distributedshell_dpgen_%s" '
+            "-jar {}/hadoop-yarn-applications-distributedshell-*.jar "
+            '-queue {} -appname "distributedshell_dpgen_{}" '
             "-shell_env YARN_CONTAINER_RUNTIME_TYPE=docker "
-            "-shell_env YARN_CONTAINER_RUNTIME_DOCKER_IMAGE=%s "
+            "-shell_env YARN_CONTAINER_RUNTIME_DOCKER_IMAGE={} "
             "-shell_env ENV_DOCKER_CONTAINER_SHM_SIZE='600m' "
             "-master_memory 1024 -master_vcores 2 -num_containers 1 "
-            "-container_resources memory-mb=%s,vcores=%s "
-            "-shell_script /tmp/%s"
-            % (
+            "-container_resources memory-mb={},vcores={} "
+            "-shell_script /tmp/{}".format(
                 resources.kwargs.get("yarn_path", ""),
                 resources.kwargs.get("yarn_path", ""),
                 resources.queue_name,

dpdispatcher/dp_cloud_server.py

@@ -106,7 +106,9 @@ class Bohrium(Machine):
 
         input_data = self.input_data.copy()
 
-        input_data["job_resources"] = job_resources
+        if not input_data.get("job_resources"):
+            input_data["job_resources"] = []
+        input_data["job_resources"].append(job_resources)
         input_data["command"] = f"bash {job.script_file_name}"
         if not input_data.get("backward_files"):
             input_data["backward_files"] = self._gen_backward_files_list(job)

dpdispatcher/dp_cloud_server_context.py

@@ -270,9 +270,6 @@ class BohriumContext(BaseContext):
     #         retcode = cmd_pipes['stdout'].channel.recv_exit_status()
     #         return retcode, cmd_pipes['stdout'], cmd_pipes['stderr']
 
-    def kill(self, cmd_pipes):
-        pass
-
     @classmethod
     def machine_subfields(cls) -> List[Argument]:
         """Generate the machine subfields.

dpdispatcher/dpcloudserver/client.py

@@ -198,7 +198,7 @@ class Client:
     ):
         post_data = {
             "job_type": job_type,
-            "oss_path": [oss_path],
+            "oss_path": oss_path,
         }
         if program_id is not None:
             post_data["project_id"] = program_id

dpdispatcher/fugaku.py

@@ -0,0 +1,94 @@
+import shlex
+
+from dpdispatcher import dlog
+from dpdispatcher.JobStatus import JobStatus
+from dpdispatcher.machine import Machine
+
+fugaku_script_header_template = """\
+{queue_name_line}
+{fugaku_node_number_line}
+{fugaku_ntasks_per_node_line}
+"""
+
+
+class Fugaku(Machine):
+    def gen_script(self, job):
+        fugaku_script = super().gen_script(job)
+        return fugaku_script
+
+    def gen_script_header(self, job):
+        resources = job.resources
+        fugaku_script_header_dict = {}
+        fugaku_script_header_dict[
+            "fugaku_node_number_line"
+        ] = f'#PJM -L "node={resources.number_node}" '
+        fugaku_script_header_dict[
+            "fugaku_ntasks_per_node_line"
+        ] = '#PJM --mpi "max-proc-per-node={cpu_per_node}"'.format(
+            cpu_per_node=resources.cpu_per_node
+        )
+        fugaku_script_header_dict[
+            "queue_name_line"
+        ] = f'#PJM -L "rscgrp={resources.queue_name}"'
+        fugaku_script_header = fugaku_script_header_template.format(
+            **fugaku_script_header_dict
+        )
+        return fugaku_script_header
+
+    def do_submit(self, job):
+        script_file_name = job.script_file_name
+        script_str = self.gen_script(job)
+        job_id_name = job.job_hash + "_job_id"
+        # script_str = self.sub_script(job_dirs, cmd, args=args, resources=resources, outlog=outlog, errlog=errlog)
+        self.context.write_file(fname=script_file_name, write_str=script_str)
+        # self.context.write_file(fname=os.path.join(self.context.submission.work_base, script_file_name), write_str=script_str)
+        # script_file_dir = os.path.join(self.context.submission.work_base)
+        script_file_dir = self.context.remote_root
+        # stdin, stdout, stderr = self.context.block_checkcall('cd %s && %s %s' % (self.context.remote_root, 'pjsub', script_file_name))
+
+        stdin, stdout, stderr = self.context.block_checkcall(
+            "cd {} && {} {}".format(
+                shlex.quote(script_file_dir), "pjsub", shlex.quote(script_file_name)
+            )
+        )
+        subret = stdout.readlines()
+        job_id = subret[0].split()[5]
+        self.context.write_file(job_id_name, job_id)
+        return job_id
+
+    def default_resources(self, resources):
+        pass
+
+    def check_status(self, job):
+        job_id = job.job_id
+        if job_id == "":
+            return JobStatus.unsubmitted
+        ret, stdin, stdout, stderr = self.context.block_call("pjstat " + job_id)
+        err_str = stderr.read().decode("utf-8")
+        try:
+            status_line = stdout.read().decode("utf-8").split("\n")[-2]
+        # pjstat only retrun 0 if the job is not waiting or running
+        except Exception:
+            ret, stdin, stdout, stderr = self.context.block_call("pjstat -H  " + job_id)
+            status_line = stdout.read().decode("utf-8").split("\n")[-2]
+            status_word = status_line.split()[3]
+            if status_word in ["EXT", "CCL", "ERR"]:
+                if self.check_finish_tag(job):
+                    dlog.info(f"job: {job.job_hash} {job.job_id} finished")
+                    return JobStatus.finished
+                else:
+                    return JobStatus.terminated
+            else:
+                return JobStatus.unknown
+        status_word = status_line.split()[3]
+        # dlog.info (status_word)
+        if status_word in ["QUE", "HLD", "RNA", "SPD"]:
+            return JobStatus.waiting
+        elif status_word in ["RUN", "RNE"]:
+            return JobStatus.running
+        else:
+            return JobStatus.unknown
+
+    def check_finish_tag(self, job):
+        job_tag_finished = job.job_hash + "_job_tag_finished"
+        return self.context.check_file_exists(job_tag_finished)

dpdispatcher/hdfs_context.py

@@ -247,6 +247,3 @@ class HDFSContext(BaseContext):
 
     def read_file(self, fname):
         return HDFS.read_hdfs_file(os.path.join(self.remote_root, fname))
-
-    def kill(self, job_id):
-        pass

dpdispatcher/lazy_local_context.py

@@ -1,5 +1,4 @@
 import os
-import signal
 import subprocess as sp
 
 from dpdispatcher.base_context import BaseContext
@@ -167,9 +166,6 @@ class LazyLocalContext(BaseContext):
         )
         return proc
 
-    def kill(self, job_id):
-        os.kill(job_id, signal.SIGTERM)
-
     def check_finish(self, proc):
         return proc.poll() is not None
 

dpdispatcher/local_context.py

@@ -1,7 +1,6 @@
 import hashlib
 import os
 import shutil
-import signal
 import subprocess as sp
 from glob import glob
 from subprocess import TimeoutExpired
@@ -291,9 +290,6 @@ class LocalContext(BaseContext):
         )
         return proc
 
-    def kill(self, job_id):
-        os.kill(job_id, signal.SIGTERM)
-
     def check_finish(self, proc):
         return proc.poll() is not None
 

dpdispatcher/lsf.py

@@ -83,8 +83,7 @@ class LSF(Machine):
 
         try:
             stdin, stdout, stderr = self.context.block_checkcall(
-                "cd %s && %s %s"
-                % (
+                "cd {} && {} {}".format(
                     shlex.quote(self.context.remote_root),
                     "bsub < ",
                     shlex.quote(script_file_name),
@@ -211,3 +210,14 @@ class LSF(Machine):
                 doc="Extra arguments.",
             )
         ]
+
+    def kill(self, job):
+        """Kill the job.
+
+        Parameters
+        ----------
+        job : Job
+            job
+        """
+        job_id = job.job_id
+        ret, stdin, stdout, stderr = self.context.block_call("bkill " + str(job_id))

dpdispatcher/machine.py

@@ -377,8 +377,12 @@ class Machine(metaclass=ABCMeta):
         machine_args = [
             Argument("batch_type", str, optional=False, doc=doc_batch_type),
             # TODO: add default to local_root and remote_root after refactor the code
-            Argument("local_root", [str, None], optional=False, doc=doc_local_root),
-            Argument("remote_root", [str, None], optional=True, doc=doc_remote_root),
+            Argument(
+                "local_root", [str, type(None)], optional=False, doc=doc_local_root
+            ),
+            Argument(
+                "remote_root", [str, type(None)], optional=True, doc=doc_remote_root
+            ),
             Argument(
                 "clean_asynchronously",
                 bool,
@@ -439,3 +443,15 @@ class Machine(metaclass=ABCMeta):
                 "kwargs", dict, optional=True, doc="This field is empty for this batch."
             )
         ]
+
+    def kill(self, job):
+        """Kill the job.
+
+        If not implemented, pass and let the user manually kill it.
+
+        Parameters
+        ----------
+        job : Job
+            job
+        """
+        dlog.warning("Job %s should be manually killed" % job.job_id)

dpdispatcher/pbs.py

@@ -46,8 +46,9 @@ class PBS(Machine):
         script_file_dir = self.context.remote_root
         # stdin, stdout, stderr = self.context.block_checkcall('cd %s && %s %s' % (self.context.remote_root, 'qsub', script_file_name))
         stdin, stdout, stderr = self.context.block_checkcall(
-            "cd %s && %s %s"
-            % (shlex.quote(script_file_dir), "qsub", shlex.quote(script_file_name))
+            "cd {} && {} {}".format(
+                shlex.quote(script_file_dir), "qsub", shlex.quote(script_file_name)
+            )
         )
         subret = stdout.readlines()
         job_id = subret[0].split()[0]
@@ -94,6 +95,17 @@ class PBS(Machine):
         job_tag_finished = job.job_hash + "_job_tag_finished"
         return self.context.check_file_exists(job_tag_finished)
 
+    def kill(self, job):
+        """Kill the job.
+
+        Parameters
+        ----------
+        job : Job
+            job
+        """
+        job_id = job.job_id
+        ret, stdin, stdout, stderr = self.context.block_call("qdel " + str(job_id))
+
 
 class Torque(PBS):
     def check_status(self, job):

dpdispatcher/shell.py

@@ -25,8 +25,7 @@ class Shell(Machine):
         output_name = job.job_hash + ".out"
         self.context.write_file(fname=script_file_name, write_str=script_str)
         ret, stdin, stdout, stderr = self.context.block_call(
-            "cd %s && { nohup bash %s 1>>%s 2>>%s & } && echo $!"
-            % (
+            "cd {} && {{ nohup bash {} 1>>{} 2>>{} & }} && echo $!".format(
                 shlex.quote(self.context.remote_root),
                 script_file_name,
                 output_name,
@@ -66,7 +65,7 @@ class Shell(Machine):
 
         # mark defunct process as terminated
         ret, stdin, stdout, stderr = self.context.block_call(
-            f"if ps -p {job_id} > /dev/null && ! (ps -p {job_id} | grep defunct >/dev/null) ; then echo 1; fi"
+            f"if ps -p {job_id} > /dev/null && ! (ps -o command -p {job_id} | grep defunct >/dev/null) ; then echo 1; fi"
         )
         if ret != 0:
             err_str = stderr.read().decode("utf-8")
@@ -101,3 +100,15 @@ class Shell(Machine):
         job_tag_finished = job.job_hash + "_job_tag_finished"
         # print('job finished: ',job.job_id, job_tag_finished)
         return self.context.check_file_exists(job_tag_finished)
+
+    def kill(self, job):
+        """Kill the job.
+
+        Parameters
+        ----------
+        job : Job
+            job
+        """
+        job_id = job.job_id
+        # 9 means exit, cannot be blocked
+        ret, stdin, stdout, stderr = self.context.block_call("kill -9 " + str(job_id))

dpdispatcher/slurm.py

@@ -1,3 +1,4 @@
+import math
 import pathlib
 import shlex
 from typing import List
@@ -45,9 +46,12 @@ class Slurm(Machine):
             )
         else:
             script_header_dict["slurm_number_gpu_line"] = custom_gpu_line
-        script_header_dict[
-            "slurm_partition_line"
-        ] = f"#SBATCH --partition {resources.queue_name}"
+        if resources.queue_name != "":
+            script_header_dict[
+                "slurm_partition_line"
+            ] = f"#SBATCH --partition {resources.queue_name}"
+        else:
+            script_header_dict["slurm_partition_line"] = ""
         slurm_script_header = slurm_script_header_template.format(**script_header_dict)
         return slurm_script_header
 
@@ -60,8 +64,7 @@ class Slurm(Machine):
         self.context.write_file(fname=script_file_name, write_str=script_str)
         # self.context.write_file(fname=os.path.join(self.context.submission.work_base, script_file_name), write_str=script_str)
         ret, stdin, stdout, stderr = self.context.block_call(
-            "cd %s && %s %s"
-            % (
+            "cd {} && {} {}".format(
                 shlex.quote(self.context.remote_root),
                 "sbatch",
                 shlex.quote(script_file_name),
@@ -78,7 +81,12 @@ class Slurm(Machine):
                     "Get error code %d in submitting through ssh with job: %s . message: %s"
                     % (ret, job.job_hash, err_str)
                 )
-            elif "Job violates accounting/QOS policy" in err_str:
+            elif (
+                "Job violates accounting/QOS policy" in err_str
+                # the number of jobs exceeds DEFAULT_MAX_JOB_COUNT (by default 10000)
+                or "Slurm temporarily unable to accept job, sleeping and retrying"
+                in err_str
+            ):
                 # job number exceeds, skip the submitting
                 return ""
             raise RuntimeError(
@@ -115,6 +123,7 @@ class Slurm(Machine):
             elif (
                 "Socket timed out on send/recv operation" in err_str
                 or "Unable to contact slurm controller" in err_str
+                or "Invalid user for SlurmUser" in err_str
             ):
                 # retry 3 times
                 raise RetrySignal(
@@ -194,30 +203,47 @@ class Slurm(Machine):
             )
         ]
 
+    def kill(self, job):
+        """Kill the job.
+
+        Parameters
+        ----------
+        job : Job
+            job
+        """
+        job_id = job.job_id
+        # -Q Do not report an error if the specified job is already completed.
+        ret, stdin, stdout, stderr = self.context.block_call(
+            "scancel -Q " + str(job_id)
+        )
+        # we do not need to stop here if scancel failed; just continue
+
 
 class SlurmJobArray(Slurm):
     """Slurm with job array enabled for multiple tasks in a job."""
 
     def gen_script_header(self, job):
+        slurm_job_size = job.resources.kwargs.get("slurm_job_size", 1)
         if job.fail_count > 0:
             # resubmit jobs, check if some of tasks have been finished
-            job_array = []
+            job_array = set()
             for ii, task in enumerate(job.job_task_list):
                 task_tag_finished = (
                     pathlib.PurePath(task.task_work_path)
                     / (task.task_hash + "_task_tag_finished")
                 ).as_posix()
                 if not self.context.check_file_exists(task_tag_finished):
-                    job_array.append(ii)
+                    job_array.add(ii // slurm_job_size)
             return super().gen_script_header(job) + "\n#SBATCH --array=%s" % (
                 ",".join(map(str, job_array))
             )
         return super().gen_script_header(job) + "\n#SBATCH --array=0-%d" % (
-            len(job.job_task_list) - 1
+            math.ceil(len(job.job_task_list) / slurm_job_size) - 1
         )
 
     def gen_script_command(self, job):
         resources = job.resources
+        slurm_job_size = resources.kwargs.get("slurm_job_size", 1)
         # SLURM_ARRAY_TASK_ID: 0 ~ n_jobs-1
         script_command = "case $SLURM_ARRAY_TASK_ID in\n"
         for ii, task in enumerate(job.job_task_list):
@@ -243,10 +269,16 @@ class SlurmJobArray(Slurm):
                 task_tag_finished=task_tag_finished,
                 log_err_part=log_err_part,
             )
-            script_command += f"{ii})\n"
+            if ii % slurm_job_size == 0:
+                script_command += f"{ii // slurm_job_size})\n"
             script_command += single_script_command
             script_command += self.gen_script_wait(resources=resources)
-            script_command += "\n;;\n"
+            script_command += "\n"
+            if (
+                ii % slurm_job_size == slurm_job_size - 1
+                or ii == len(job.job_task_list) - 1
+            ):
+                script_command += ";;\n"
         script_command += "*)\nexit 1\n;;\nesac\n"
         return script_command
 
@@ -337,9 +369,30 @@ class SlurmJobArray(Slurm):
     def check_finish_tag(self, job):
         results = []
         for task in job.job_task_list:
-            task_tag_finished = (
-                pathlib.PurePath(task.task_work_path)
-                / (task.task_hash + "_task_tag_finished")
-            ).as_posix()
-            results.append(self.context.check_file_exists(task_tag_finished))
+            task.get_task_state(self.context)
+            results.append(task.task_state == JobStatus.finished)
         return all(results)
+
+    @classmethod
+    def resources_subfields(cls) -> List[Argument]:
+        """Generate the resources subfields.
+
+        Returns
+        -------
+        list[Argument]
+            resources subfields
+        """
+        doc_slurm_job_size = "Number of tasks in a Slurm job"
+        arg = super().resources_subfields()[0]
+        arg.extend_subfields(
+            [
+                Argument(
+                    "slurm_job_size",
+                    int,
+                    optional=True,
+                    default=1,
+                    doc=doc_slurm_job_size,
+                ),
+            ]
+        )
+        return [arg]

dpdispatcher/ssh_context.py

@@ -116,7 +116,7 @@ class SSHSession:
     #     transport = self.ssh.get_transport()
     #     transport.set_keepalive(60)
 
-    @retry(max_retry=3, sleep=1)
+    @retry(max_retry=6, sleep=1)
     def _setup_ssh(self):
         # machine = self.machine
         self.ssh = paramiko.SSHClient()
@@ -199,7 +199,7 @@ class SSHSession:
                 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 3 times
+                # retry for up to 6 times
                 raise RetrySignal("Authentication failed")
         elif key_ok:
             pass
@@ -213,7 +213,12 @@ class SSHSession:
             raise RuntimeError("Please provide at least one form of authentication")
         assert ts.is_active()
         # Opening a session creates a channel along the socket to the server
-        ts.open_session(timeout=self.timeout)
+        try:
+            ts.open_session(timeout=self.timeout)
+        except paramiko.ssh_exception.SSHException:
+            # retry for up to 6 times
+            # ref: https://github.com/paramiko/paramiko/issues/1508
+            raise RetrySignal("Opening session failed")
         ts.set_keepalive(60)
         self.ssh._transport = ts  # type: ignore
         # reset sftp
@@ -323,14 +328,14 @@ class SSHSession:
             Argument("port", int, optional=True, default=22, doc=doc_port),
             Argument(
                 "key_filename",
-                [str, None],
+                [str, type(None)],
                 optional=True,
                 default=None,
                 doc=doc_key_filename,
             ),
             Argument(
                 "passphrase",
-                [str, None],
+                [str, type(None)],
                 optional=True,
                 default=None,
                 doc=doc_passphrase,
@@ -497,6 +502,14 @@ class SSHContext(BaseContext):
             self.block_checkcall(
                 f"mv {shlex.quote(old_remote_root)} {shlex.quote(self.remote_root)}"
             )
+        elif (
+            old_remote_root is not None
+            and old_remote_root != self.remote_root
+            and self.check_file_exists(old_remote_root)
+            and not len(self.ssh_session.sftp.listdir(old_remote_root))
+        ):
+            # if the new directory exists and the old directory does not contain files, then move the old directory
+            self._rmtree(old_remote_root)
 
         sftp = self.ssh_session.ssh.open_sftp()
         try:
@@ -762,12 +775,6 @@ class SSHContext(BaseContext):
             retcode = cmd_pipes["stdout"].channel.recv_exit_status()
             return retcode, cmd_pipes["stdout"], cmd_pipes["stderr"]
 
-    def kill(self, cmd_pipes):
-        raise RuntimeError(
-            "dose not work! we do not know how to kill proc through paramiko.SSHClient"
-        )
-        # self.block_checkcall('kill -15 %s' % cmd_pipes['pid'])
-
     def _rmtree(self, remotepath, verbose=False):
         """Remove the remote path."""
         # The original implementation method removes files one by one using sftp.
@@ -847,8 +854,7 @@ class SSHContext(BaseContext):
             self.ssh_session.put(from_f, to_f)
         except FileNotFoundError:
             raise FileNotFoundError(
-                "from %s to %s @ %s : %s Error!"
-                % (from_f, self.ssh_session.username, self.ssh_session.hostname, to_f)
+                f"from {from_f} to {self.ssh_session.username} @ {self.ssh_session.hostname} : {to_f} Error!"
             )
         # remote extract
         self.block_checkcall("tar xf %s" % of)
@@ -877,8 +883,7 @@ class SSHContext(BaseContext):
         ntar = len(files) // per_nfile + 1
         if ntar <= 1:
             self.block_checkcall(
-                "tar %s %s %s"
-                % (
+                "tar {} {} {}".format(
                     tar_command,
                     shlex.quote(of),
                     " ".join([shlex.quote(file) for file in files]),
@@ -890,8 +895,7 @@ class SSHContext(BaseContext):
             )
             self.write_file(file_list_file, "\n".join(files))
             self.block_checkcall(
-                "tar %s %s -T %s"
-                % (tar_command, shlex.quote(of), shlex.quote(file_list_file))
+                f"tar {tar_command} {shlex.quote(of)} -T {shlex.quote(file_list_file)}"
             )
         # trans
         from_f = pathlib.PurePath(os.path.join(self.remote_root, of)).as_posix()

dpdispatcher/submission.py

@@ -1,7 +1,10 @@
 # %%
+import asyncio
 import copy
+import functools
 import json
 import os
+import pathlib
 import random
 import time
 import uuid
@@ -198,7 +201,9 @@ class Submission:
             self.local_root = machine.context.temp_local_root
         return self
 
-    def run_submission(self, *, dry_run=False, exit_on_submit=False, clean=True):
+    def run_submission(
+        self, *, dry_run=False, exit_on_submit=False, clean=True, check_interval=30
+    ):
         """Main method to execute the submission.
         First, check whether old Submission exists on the remote machine, and try to recover from it.
         Second, upload the local files to the remote machine where the tasks to be executed.
@@ -235,11 +240,11 @@ class Submission:
                 dlog.info(f"at {self.machine.context.remote_root}")
                 return self.serialize()
             if ratio_unfinished > 0.0 and self.check_ratio_unfinished(ratio_unfinished):
-                self.remove_unfinished_jobs()
+                self.remove_unfinished_tasks()
                 break
 
             try:
-                time.sleep(30)
+                time.sleep(check_interval)
             except (Exception, KeyboardInterrupt, SystemExit) as e:
                 self.submission_to_json()
                 dlog.exception(e)
@@ -253,12 +258,72 @@ class Submission:
             finally:
                 pass
         self.handle_unexpected_submission_state()
-        self.download_jobs()
+        self.try_download_result()
         self.submission_to_json()
         if clean:
             self.clean_jobs()
         return self.serialize()
 
+    def try_download_result(self):
+        start_time = time.time()
+        retry_interval = 60  # 每1分钟重试一次
+        success = False
+        while not success:
+            try:
+                self.download_jobs()
+                success = True
+            except (EOFError, Exception) as e:
+                dlog.exception(e)
+                elapsed_time = time.time() - start_time
+                if elapsed_time < 3600:  # 1小时内
+                    dlog.info("Retrying in 1 minute...")
+                    time.sleep(retry_interval)
+                elif elapsed_time < 86400:  # 1小时后，但在24小时内
+                    retry_interval = 600  # 每10分钟重试一次
+                    dlog.info("Retrying in 10 minutes...")
+                    time.sleep(retry_interval)
+                else:  # 超过24小时
+                    dlog.info("Maximum retries time reached. Exiting.")
+                    break
+
+    async def async_run_submission(self, **kwargs):
+        """Async interface of run_submission.
+
+        Examples
+        --------
+        >>> import asyncio
+        >>> from dpdispacher import Machine, Resource, Submission
+        >>> async def run_jobs():
+        ...     backgroud_task = set()
+        ...     # task1
+        ...     task1 = Task(...)
+        ...     submission1 = Submission(..., task_list=[task1])
+        ...     background_task = asyncio.create_task(
+        ...         submission1.async_run_submission(check_interval=2, clean=False)
+        ...     )
+        ...     # task2
+        ...     task2 = Task(...)
+        ...     submission2 = Submission(..., task_list=[task1])
+        ...     background_task = asyncio.create_task(
+        ...         submission2.async_run_submission(check_interval=2, clean=False)
+        ...     )
+        ...     background_tasks.add(background_task)
+        ...     result = await asyncio.gather(*background_tasks)
+        ...     return result
+        >>> run_jobs()
+
+        May raise Error if pass `clean=True` explicitly when submit to pbs or slurm.
+        """
+        kwargs = {**{"clean": False}, **kwargs}
+        if kwargs["clean"]:
+            dlog.warning(
+                "Using async submission with `clean=True`, "
+                "job may fail in queue system"
+            )
+        loop = asyncio.get_event_loop()
+        wrapped_submission = functools.partial(self.run_submission, **kwargs)
+        return await loop.run_in_executor(None, wrapped_submission)
+
     def update_submission_state(self):
         """Check whether all the jobs in the submission.
 
@@ -306,41 +371,53 @@ class Submission:
 
     # def update_submi
 
-    def check_ratio_unfinished(self, ratio_unfinished):
-        status_list = [job.job_state for job in self.belonging_jobs]
-        finished_num = status_list.count(JobStatus.finished)
-        if finished_num / len(self.belonging_jobs) < (1 - ratio_unfinished):
-            return False
+    def check_ratio_unfinished(self, ratio_unfinished: float) -> bool:
+        """Calculate the ratio of unfinished tasks in the submission.
+
+        Parameters
+        ----------
+        ratio_unfinished : float
+            the ratio of unfinished tasks in the submission
+
+        Returns
+        -------
+        bool
+            whether the ratio of unfinished tasks in the submission is larger than ratio_unfinished
+        """
+        assert self.resources is not None
+        if self.resources.group_size == 1:
+            # if group size is 1, calculate job state is enough and faster
+            status_list = [job.job_state for job in self.belonging_jobs]
         else:
-            return True
+            # get task state is more accurate
+            status_list = []
+            for task in self.belonging_tasks:
+                task.get_task_state(self.machine.context)
+                status_list.append(task.task_state)
+        finished_num = status_list.count(JobStatus.finished)
+        return finished_num / len(self.belonging_tasks) >= (1 - ratio_unfinished)
 
-    def remove_unfinished_jobs(self):
-        removed_jobs = [
-            job
-            for job in self.belonging_jobs
-            if job.job_state not in [JobStatus.finished]
-        ]
-        self.belonging_jobs = [
-            job for job in self.belonging_jobs if job.job_state in [JobStatus.finished]
-        ]
-        for job in removed_jobs:
-            # kill unfinished jobs
-            try:
-                self.machine.context.kill(job.job_id)
-            except Exception as e:
-                dlog.info("Can not kill job %s" % job.job_id)
-
-            # remove unfinished tasks
-            import os
-            import shutil
-
-            for task in job.job_task_list:
-                shutil.rmtree(
-                    os.path.join(self.machine.context.local_root, task.task_work_path),
-                    ignore_errors=True,
-                )
-            self.belonging_tasks = [
-                task for task in self.belonging_tasks if task not in job.job_task_list
+    def remove_unfinished_tasks(self):
+        dlog.info("Remove unfinished tasks")
+        # kill all jobs and mark them as finished
+        for job in self.belonging_jobs:
+            if job.job_state != JobStatus.finished:
+                self.machine.kill(job)
+                job.job_state = JobStatus.finished
+        # remove all unfinished tasks
+        finished_tasks = []
+        for task in self.belonging_tasks:
+            if task.task_state == JobStatus.finished:
+                finished_tasks.append(task)
+            # there is no need to remove actual remote directory
+            # as it should be cleaned anyway
+        self.belonging_tasks = finished_tasks
+        # clean removed tasks in jobs - although this should not be necessary
+        for job in self.belonging_jobs:
+            job.job_task_list = [
+                task
+                for task in job.job_task_list
+                if task.task_state == JobStatus.finished
             ]
 
     def check_all_finished(self):
@@ -463,6 +540,9 @@ class Submission:
             submission.bind_machine(machine=self.machine)
             if self == submission:
                 self.belonging_jobs = submission.belonging_jobs
+                self.belonging_tasks = [
+                    task for job in self.belonging_jobs for task in job.job_task_list
+                ]
                 self.bind_machine(machine=self.machine)
                 dlog.info(
                     f"Find old submission; recover submission from json file;"
@@ -518,6 +598,7 @@ class Task:
         self.task_hash = self.get_hash()
         # self.task_need_resources="<to be completed in the future>"
         # self.uuid =
+        self.task_state = JobStatus.unsubmitted
 
     def __repr__(self):
         return str(self.serialize())
@@ -602,15 +683,44 @@ class Task:
                 default=[],
             ),
             Argument(
-                "outlog", [None, str], optional=False, doc=doc_outlog, default="log"
+                "outlog",
+                [type(None), str],
+                optional=False,
+                doc=doc_outlog,
+                default="log",
             ),
             Argument(
-                "errlog", [None, str], optional=False, doc=doc_errlog, default="err"
+                "errlog",
+                [type(None), str],
+                optional=False,
+                doc=doc_errlog,
+                default="err",
             ),
         ]
         task_format = Argument("task", dict, task_args)
         return task_format
 
+    def get_task_state(self, context):
+        """Get the task state by checking the tag file.
+
+        Parameters
+        ----------
+        context : Context
+            the context of the task
+        """
+        if self.task_state in (JobStatus.finished, JobStatus.unsubmitted):
+            # finished task should always be finished
+            # unsubmitted task do not need to check tag
+            return
+        # check tag
+        task_tag_finished = (
+            pathlib.PurePath(self.task_work_path)
+            / (self.task_hash + "_task_tag_finished")
+        ).as_posix()
+        result = context.check_file_exists(task_tag_finished)
+        if result:
+            self.task_state = JobStatus.finished
+
 
 class Job:
     """Job is generated by Submission automatically.
@@ -700,6 +810,8 @@ class Job:
         job.job_id = job_dict[job_hash]["job_id"]
         job.fail_count = job_dict[job_hash]["fail_count"]
         # job.job_uuid = job_dict[job_hash]['job_uuid']
+        for task in job.job_task_list:
+            task.task_state = job.job_state
         return job
 
     def get_job_state(self):
@@ -715,6 +827,11 @@ class Job:
         assert self.machine is not None
         job_state = self.machine.check_status(self)
         self.job_state = job_state
+        # update general task_state, which should be faster than checking tags
+        for task in self.job_task_list:
+            # only update if the task is not finished
+            if task.task_state != JobStatus.finished:
+                task.task_state = job_state
 
     def handle_unexpected_job_state(self):
         job_state = self.job_state
@@ -838,7 +955,7 @@ class Resources:
             If true, dpdispatcher will manually export environment variable CUDA_VISIBLE_DEVICES to different task.
             Usually, this option will be used with Task.task_need_resources variable simultaneously.
         ratio_unfinished : float
-            The ratio of `jobs` that can be unfinished.
+            The ratio of `task` that can be unfinished.
     para_deg : int
         Decide how many tasks will be run in parallel.
         Usually run with `strategy['if_cuda_multi_devices']`
@@ -1010,7 +1127,7 @@ class Resources:
             "If true, dpdispatcher will manually export environment variable CUDA_VISIBLE_DEVICES to different task."
             "Usually, this option will be used with Task.task_need_resources variable simultaneously."
         )
-        doc_ratio_unfinished = "The ratio of `jobs` that can be unfinished."
+        doc_ratio_unfinished = "The ratio of `tasks` that can be unfinished."
 
         strategy_args = [
             Argument(

{dpdispatcher-0.5.6.dist-info → dpdispatcher-0.5.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dpdispatcher
-Version: 0.5.6
+Version: 0.5.8
 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
@@ -204,15 +204,20 @@ Provides-Extra: test
 
 # DPDispatcher
 
-DPDispatcher is a python package used to generate HPC(High Performance Computing) scheduler systems (Slurm/PBS/LSF/dpcloudserver) jobs input scripts and submit these  scripts to HPC systems and poke until they finish.
+[![conda-forge](https://img.shields.io/conda/dn/conda-forge/dpdispatcher?color=red&label=conda-forge&logo=conda-forge)](https://anaconda.org/conda-forge/dpdispatcher)
+[![pip install](https://img.shields.io/pypi/dm/dpdispatcher?label=pip%20install&logo=pypi)](https://pypi.org/project/dpdispatcher)
+[![docker pull](https://img.shields.io/docker/pulls/dptechnology/dpdispatcher?logo=docker)](https://hub.docker.com/r/dptechnology/dpdispatcher)
+[![Documentation Status](https://readthedocs.org/projects/dpdispatcher/badge/)](https://dpdispatcher.readthedocs.io/)
+
+DPDispatcher is a Python package used to generate HPC (High-Performance Computing) scheduler systems (Slurm/PBS/LSF/Bohrium) jobs input scripts, submit them to HPC systems, and poke until they finish.
 ​
-DPDispatcher will monitor (poke) until these jobs finish and download the results files (if these jobs is running on remote systems connected by SSH).
+DPDispatcher will monitor (poke) until these jobs finish and download the results files (if these jobs are running on remote systems connected by SSH).
 
 For more information, check the [documentation](https://dpdispatcher.readthedocs.io/).
 
 ## Installation
 
-DPDispatcher can installed by `pip`:
+DPDispatcher can be installed by `pip`:
 
 ```bash
 pip install dpdispatcher
@@ -224,5 +229,9 @@ See [Getting Started](https://dpdispatcher.readthedocs.io/en/latest/getting-star
 
 ## Contributing
 
-DPDispatcher is maintained by Deep Modeling's developers and welcome other people.
+DPDispatcher is maintained by Deep Modeling's developers and welcomes other people.
 See [Contributing Guide](CONTRIBUTING.md) to become a contributor! 🤓
+
+## References
+
+DPDispatcher is derivated from the [DP-GEN](https://github.com/deepmodeling/dpgen) package. To mention DPDispatcher in a scholarly publication, please read Section 3.3 in the [DP-GEN paper](https://doi.org/10.1016/j.cpc.2020.107206).

dpdispatcher-0.5.8.dist-info/RECORD

@@ -0,0 +1,34 @@
+dpdispatcher/JobStatus.py,sha256=Eszs4TPLfszCuf6zLaFonf25feXDUguF28spYOjJpQE,233
+dpdispatcher/__init__.py,sha256=2GIz4niyzHTbxros1G7Mi4uBJbD3AMSnTPxXSJMJmUs,2907
+dpdispatcher/_version.py,sha256=iqWtoISytDDNpYe-atC8Kl-rZhTojPnDQKAEcFNtIhg,160
+dpdispatcher/arginfo.py,sha256=pNaxYIE6ahBidpR7OCKZdw8iGt003uTXGSlVzwiuvRg,188
+dpdispatcher/base_context.py,sha256=Hfri0x41XC4MRUjxc0-WMiZB_E4NvLp94ZYaHfYCWHM,3610
+dpdispatcher/distributed_shell.py,sha256=XMcXt8g1f2DY5HYhhyiN5ehV2ihKULY5ng-sB0B7YaI,6933
+dpdispatcher/dp_cloud_server.py,sha256=xVpDI0exBwHNSZECLJdfrQsvBzeUn5a0gx5Bzt9UAdU,9857
+dpdispatcher/dp_cloud_server_context.py,sha256=VfRRo4ruorWC8NVjW19EjmxQ0Rbz6XzxrHrJKl4cCZk,11255
+dpdispatcher/dpdisp.py,sha256=_dyH8xEgUR-s2xKkB20D9FIYhSHUCmzc2PxWgo9ildQ,94
+dpdispatcher/fugaku.py,sha256=wSjY0XB3TNNWAPKHgMpoPl5jyYJIlijBcEkYXp6nrZQ,3733
+dpdispatcher/hdfs_cli.py,sha256=9Vrf7Kz_kJgXP2xEdZqNVNxRGbui5RrtnLtEjxfcq9A,6047
+dpdispatcher/hdfs_context.py,sha256=1jT1nzx7VGJFJ42MHTXoFWhfEu4KBkMBJO84klRAnPI,8938
+dpdispatcher/lazy_local_context.py,sha256=ZdWNqK3QF8SsoqnCjpFt3ZDRCIagjzJNlKPUYutRUC8,5692
+dpdispatcher/local_context.py,sha256=anYJqQASOnkcAhfckUcFD8_DcjNUZ1KE0GuksxR5Mxw,11772
+dpdispatcher/lsf.py,sha256=zy-WEnC7f2Dy5hJGnRBl5jpjYZ_H3-KMcE0lxDG6ejo,7790
+dpdispatcher/machine.py,sha256=31xG5ksN8mBVwD8taLsk5KXLhjM0ZTjlHlbbPgiig1c,15296
+dpdispatcher/pbs.py,sha256=LiULEKNDuisrKmOpZyB1af6sGDQ35xrAhMh7VMwpFbY,6327
+dpdispatcher/shell.py,sha256=kEP7za-qN71y_21p0uBNkopZ5s63Adq54904hjUHv48,4141
+dpdispatcher/slurm.py,sha256=krlyjzxK8gIhSsqcKHFvNiUwVE7411wTUwuW9xGzS-E,14648
+dpdispatcher/ssh_context.py,sha256=7Xrm8biVA7tAEDJ6YJZzC3nbdQrVBr_5UOhQNQ7qJ2g,35032
+dpdispatcher/submission.py,sha256=r_F05nHTpN86b2os8RZAjZsCILNarDko2BjAEUYSntw,46643
+dpdispatcher/utils.py,sha256=RXUHJl3S2z26Em3SeltnxtdVM3kv7weXJKvBEjG6I34,5035
+dpdispatcher/dpcloudserver/__init__.py,sha256=FnX9HH-2dXADluNfucg98JPMfruMoBpN9ER9lZkVQvQ,49
+dpdispatcher/dpcloudserver/client.py,sha256=w1wQ8g-FMQlyh00LIAbJLE1xirGXocpp7zAnhbeM4V0,11152
+dpdispatcher/dpcloudserver/config.py,sha256=vBRtzExJXTGfXPeBObXrZNAhBNXoFFzMkzSuSrrjHEQ,635
+dpdispatcher/dpcloudserver/retcode.py,sha256=1qAF8gFZx55u2sO8KbtYSIIrjcO-IGufEUlwbkSfC1g,721
+dpdispatcher/dpcloudserver/temp_test.py,sha256=jklOSu7tZ_wW5gycGRiUsbBWMLZDqCBslSYOCb2hTHw,2932
+dpdispatcher/dpcloudserver/zip_file.py,sha256=f9WrlktwHW0YipaWg5Y0kxjMZlhD1cJYa6EUpvu4Cro,2611
+dpdispatcher-0.5.8.dist-info/LICENSE,sha256=46mU2C5kSwOnkqkw9XQAJlhBL2JAf1_uCD8lVcXyMRg,7652
+dpdispatcher-0.5.8.dist-info/METADATA,sha256=o2oD8_6Ohc04mRTkJWi51-KOPamYqH0kvUD-E0iW-c0,12280
+dpdispatcher-0.5.8.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
+dpdispatcher-0.5.8.dist-info/entry_points.txt,sha256=3bKn6IB6SYhKOUbbcOdBBevz4gsDmhmbogKMVn4ptOQ,52
+dpdispatcher-0.5.8.dist-info/top_level.txt,sha256=35jAQoXY-b-e9fJ1_mxhZUiaCoJNt1ZI7mpFRf07Qjs,13
+dpdispatcher-0.5.8.dist-info/RECORD,,

dpdispatcher-0.5.6.dist-info/RECORD

@@ -1,33 +0,0 @@
-dpdispatcher/JobStatus.py,sha256=Eszs4TPLfszCuf6zLaFonf25feXDUguF28spYOjJpQE,233
-dpdispatcher/__init__.py,sha256=U8OLDjSGHxILiz8XH-HYBxjIlhD429HEqqxQ-vVK1a4,2866
-dpdispatcher/_version.py,sha256=J0O-QTcfk70wRsnrg-XApMYr8T5heaCiqHkl3PJ9zfs,160
-dpdispatcher/arginfo.py,sha256=pNaxYIE6ahBidpR7OCKZdw8iGt003uTXGSlVzwiuvRg,188
-dpdispatcher/base_context.py,sha256=XTKN0T_ffhVipEsbVEYNPmbKds8qMuwWbCixAhV8mUc,3690
-dpdispatcher/distributed_shell.py,sha256=vbNT8VHaYwEu2zIVFyosQGji4C3_QSpWKMFZURkJC7c,6941
-dpdispatcher/dp_cloud_server.py,sha256=mwXt2rtQeW4uMYBP05JcYE5vm5dp2JeitxiS8KjpFQc,9758
-dpdispatcher/dp_cloud_server_context.py,sha256=t47Kfn3cyQ223cBU-HVVdKvTm8yQiqHsFcMnQcIXdgk,11300
-dpdispatcher/dpdisp.py,sha256=_dyH8xEgUR-s2xKkB20D9FIYhSHUCmzc2PxWgo9ildQ,94
-dpdispatcher/hdfs_cli.py,sha256=9Vrf7Kz_kJgXP2xEdZqNVNxRGbui5RrtnLtEjxfcq9A,6047
-dpdispatcher/hdfs_context.py,sha256=IGvXsw9wdR8aemQ9kOE5WaciwVLtZbr-t2mrCQjxywU,8980
-dpdispatcher/lazy_local_context.py,sha256=V0jVuAgOHKw_PkYPCKnr3OkMyWfLRTJ8B7As3VCzLX8,5775
-dpdispatcher/local_context.py,sha256=8tML77WRSydJoPA6DseYsIshV-id5xO-6kWsqDsJHQ0,11855
-dpdispatcher/lsf.py,sha256=PjsjNO8YZkWBzFe_277G1oVrLcAm1Qz1fN-3FZ4dsK4,7553
-dpdispatcher/machine.py,sha256=vafq9zTW6NYZ-ZOMEfVEiZkoZzbEjco09d4VwEb9Jk8,14949
-dpdispatcher/pbs.py,sha256=RScX8rX1lGlvilgCEFxzUAeTA-Em5AeAcM2yT2OKY3s,6057
-dpdispatcher/shell.py,sha256=SZoJynOmqldMrl-lIMsNOY1RELFNFWWffeSg7XJsi9g,3843
-dpdispatcher/slurm.py,sha256=bsSTaRe5t3-z5BdjMyerGGBEXjJ_BfzMcawNMaFULfs,12886
-dpdispatcher/ssh_context.py,sha256=aeGiTUmBzfwkZ4xgNsyXBSpHAUHoUEVltidkNFWLXUE,34670
-dpdispatcher/submission.py,sha256=5qt1Nw3qRQeFq-w5DvzoWb9K-qMLiOFUjy_JhtnxFfI,41960
-dpdispatcher/utils.py,sha256=RXUHJl3S2z26Em3SeltnxtdVM3kv7weXJKvBEjG6I34,5035
-dpdispatcher/dpcloudserver/__init__.py,sha256=FnX9HH-2dXADluNfucg98JPMfruMoBpN9ER9lZkVQvQ,49
-dpdispatcher/dpcloudserver/client.py,sha256=a1KzBbbBKz6ZMH9iWhQfrdhL5BwrWevLS1vlTK4WP8w,11154
-dpdispatcher/dpcloudserver/config.py,sha256=vBRtzExJXTGfXPeBObXrZNAhBNXoFFzMkzSuSrrjHEQ,635
-dpdispatcher/dpcloudserver/retcode.py,sha256=1qAF8gFZx55u2sO8KbtYSIIrjcO-IGufEUlwbkSfC1g,721
-dpdispatcher/dpcloudserver/temp_test.py,sha256=jklOSu7tZ_wW5gycGRiUsbBWMLZDqCBslSYOCb2hTHw,2932
-dpdispatcher/dpcloudserver/zip_file.py,sha256=f9WrlktwHW0YipaWg5Y0kxjMZlhD1cJYa6EUpvu4Cro,2611
-dpdispatcher-0.5.6.dist-info/LICENSE,sha256=46mU2C5kSwOnkqkw9XQAJlhBL2JAf1_uCD8lVcXyMRg,7652
-dpdispatcher-0.5.6.dist-info/METADATA,sha256=KtfW7Uy8R5YHqMkRH3oDE5jh3JHXk_RcZ1W5H_gv2EM,11481
-dpdispatcher-0.5.6.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
-dpdispatcher-0.5.6.dist-info/entry_points.txt,sha256=3bKn6IB6SYhKOUbbcOdBBevz4gsDmhmbogKMVn4ptOQ,52
-dpdispatcher-0.5.6.dist-info/top_level.txt,sha256=35jAQoXY-b-e9fJ1_mxhZUiaCoJNt1ZI7mpFRf07Qjs,13
-dpdispatcher-0.5.6.dist-info/RECORD,,