dpdispatcher 0.5.7__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.7'
-__version_tuple__ = version_tuple = (0, 5, 7)
+__version__ = version = '0.5.8'
+__version_tuple__ = version_tuple = (0, 5, 8)

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/submission.py

@@ -1,5 +1,7 @@
 # %%
+import asyncio
 import copy
+import functools
 import json
 import os
 import pathlib
@@ -199,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.
@@ -240,7 +244,7 @@ class Submission:
                 break
 
             try:
-                time.sleep(30)
+                time.sleep(check_interval)
             except (Exception, KeyboardInterrupt, SystemExit) as e:
                 self.submission_to_json()
                 dlog.exception(e)
@@ -254,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.
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dpdispatcher
-Version: 0.5.7
+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

{dpdispatcher-0.5.7.dist-info → dpdispatcher-0.5.8.dist-info}/RECORD

@@ -1,12 +1,13 @@
 dpdispatcher/JobStatus.py,sha256=Eszs4TPLfszCuf6zLaFonf25feXDUguF28spYOjJpQE,233
-dpdispatcher/__init__.py,sha256=U8OLDjSGHxILiz8XH-HYBxjIlhD429HEqqxQ-vVK1a4,2866
-dpdispatcher/_version.py,sha256=JZGzq-Js5LHjsQ55qR1dohZK98vvZwStS3mcoNYbvzE,160
+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
@@ -17,7 +18,7 @@ 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=V8vdPSYgFJULbFs1H-K5Xj9_ynWC3vSmMekgngGPs44,44077
+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
@@ -25,9 +26,9 @@ dpdispatcher/dpcloudserver/config.py,sha256=vBRtzExJXTGfXPeBObXrZNAhBNXoFFzMkzSu
 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.7.dist-info/LICENSE,sha256=46mU2C5kSwOnkqkw9XQAJlhBL2JAf1_uCD8lVcXyMRg,7652
-dpdispatcher-0.5.7.dist-info/METADATA,sha256=f-kzhWa-DzQnSOHvFxoT6A9rIZy7iOfFrZQpNvQCpPY,12280
-dpdispatcher-0.5.7.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
-dpdispatcher-0.5.7.dist-info/entry_points.txt,sha256=3bKn6IB6SYhKOUbbcOdBBevz4gsDmhmbogKMVn4ptOQ,52
-dpdispatcher-0.5.7.dist-info/top_level.txt,sha256=35jAQoXY-b-e9fJ1_mxhZUiaCoJNt1ZI7mpFRf07Qjs,13
-dpdispatcher-0.5.7.dist-info/RECORD,,
+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,,