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

dpdispatcher/dpdisp.py

@@ -1,8 +1,120 @@
 #!/usr/bin/env python
+import argparse
+from typing import List, Optional
+
+from dpdispatcher.entrypoints.gui import start_dpgui
+from dpdispatcher.entrypoints.submission import handle_submission
+
+
+def main_parser() -> argparse.ArgumentParser:
+    """Dpdispatcher commandline options argument parser.
+
+    Notes
+    -----
+    This function is used by documentation.
+
+    Returns
+    -------
+    argparse.ArgumentParser
+        the argument parser
+    """
+    parser = argparse.ArgumentParser(
+        description="dpdispatcher: Generate HPC scheduler systems jobs input scripts, submit these scripts to HPC systems, and poke until they finish",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    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",
+        help="Serve DP-GUI.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser_gui.add_argument(
+        "-p",
+        "--port",
+        type=int,
+        default=6042,
+        help="The port to serve DP-GUI on.",
+    )
+    parser_gui.add_argument(
+        "--bind_all",
+        action="store_true",
+        help=(
+            "Serve on all public interfaces. This will expose your DP-GUI instance "
+            "to the network on both IPv4 and IPv6 (where available)."
+        ),
+    )
+    return parser
+
+
+def parse_args(args: Optional[List[str]] = None):
+    """Dpdispatcher commandline options argument parsing.
+
+    Parameters
+    ----------
+    args : List[str]
+        list of command line arguments, main purpose is testing default option None
+        takes arguments from sys.argv
+    """
+    parser = main_parser()
+
+    parsed_args = parser.parse_args(args=args)
+    if parsed_args.command is None:
+        parser.print_help()
+
+    return parsed_args
 
 
 def main():
-    print("test")
+    args = parse_args()
+    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,
+        )
+    elif args.command is None:
+        pass
+    else:
+        raise RuntimeError(f"unknown command {args.command}")
 
 
 if __name__ == "__main__":

dpdispatcher/entrypoints/__init__.py

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

dpdispatcher/entrypoints/gui.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""DP-GUI entrypoint."""
+
+
+def start_dpgui(*, port: int, bind_all: bool, **kwargs):
+    """Host DP-GUI server.
+
+    Parameters
+    ----------
+    port : int
+        The port to serve DP-GUI on.
+    bind_all : bool
+        Serve on all public interfaces. This will expose your DP-GUI instance
+        to the network on both IPv4 and IPv6 (where available).
+    **kwargs
+        additional arguments
+
+    Raises
+    ------
+    ModuleNotFoundError
+        The dpgui package is not installed
+    """
+    try:
+        from dpgui import (
+            start_dpgui,
+        )
+    except ModuleNotFoundError as e:
+        raise ModuleNotFoundError(
+            "To use DP-GUI, please install the dpgui package:\npip install dpgui"
+        ) from e
+    start_dpgui(port=port, bind_all=bind_all)

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/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}
@@ -35,7 +36,7 @@ cd {task_work_path}
 test $? -ne 0 && exit 1
 if [ ! -f {task_tag_finished} ] ;then
   {command_env} ( {command} ) {log_err_part}
-  if test $? -eq 0; then touch {task_tag_finished}; else echo 1 > $REMOTE_ROOT/{flag_if_job_task_fail};fi
+  if test $? -eq 0; then touch {task_tag_finished}; else echo 1 > $REMOTE_ROOT/{flag_if_job_task_fail};tail -v -c 1000 $REMOTE_ROOT/{task_work_path}/{err_file} > $REMOTE_ROOT/{last_err_file};fi
 fi &
 """
 
@@ -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"]
@@ -191,17 +199,20 @@ class Machine(metaclass=ABCMeta):
             "abstract method do_submit should be implemented by derived class"
         )
 
+    def gen_script_run_command(self, job):
+        return f"source $REMOTE_ROOT/{job.script_file_name}.run"
+
     def gen_script(self, job):
         script_header = self.gen_script_header(job)
         script_custom_flags = self.gen_script_custom_flags_lines(job)
         script_env = self.gen_script_env(job)
-        script_command = self.gen_script_command(job)
+        script_run_command = self.gen_script_run_command(job)
         script_end = self.gen_script_end(job)
         script = script_template.format(
             script_header=script_header,
             script_custom_flags=script_custom_flags,
             script_env=script_env,
-            script_command=script_command,
+            script_command=script_run_command,
             script_end=script_end,
         )
         return script
@@ -295,6 +306,7 @@ class Machine(metaclass=ABCMeta):
                 log_err_part += f"2>>{shlex.quote(task.errlog)} "
 
             flag_if_job_task_fail = job.job_hash + "_flag_if_job_task_fail"
+            last_err_file = job.job_hash + "_last_err_file"
             single_script_command = script_command_template.format(
                 flag_if_job_task_fail=flag_if_job_task_fail,
                 command_env=command_env,
@@ -304,6 +316,8 @@ class Machine(metaclass=ABCMeta):
                 command=task.command,
                 task_tag_finished=task_tag_finished,
                 log_err_part=log_err_part,
+                err_file=shlex.quote(task.errlog),
+                last_err_file=shlex.quote(last_err_file),
             )
             script_command += single_script_command
 

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/{distributed_shell.py → machines/distributed_shell.py}

@@ -1,7 +1,10 @@
-from dpdispatcher import dlog
-from dpdispatcher.JobStatus import JobStatus
+from dpdispatcher.dlog import dlog
 from dpdispatcher.machine import Machine
-from dpdispatcher.utils import run_cmd_with_all_output
+from dpdispatcher.utils.job_status import JobStatus
+from dpdispatcher.utils.utils import (
+    customized_script_header_template,
+    run_cmd_with_all_output,
+)
 
 shell_script_header_template = """
 #!/bin/bash -l
@@ -112,7 +115,17 @@ class DistributedShell(Machine):
         return script_end
 
     def gen_script_header(self, job):
-        shell_script_header = shell_script_header_template
+        resources = job.resources
+        if (
+            resources["strategy"].get("customized_script_header_template_file")
+            is not None
+        ):
+            shell_script_header = customized_script_header_template(
+                resources["strategy"]["customized_script_header_template_file"],
+                resources,
+            )
+        else:
+            shell_script_header = shell_script_header_template
         return shell_script_header
 
     def do_submit(self, job):
@@ -133,6 +146,9 @@ class DistributedShell(Machine):
         job_id_name = job.job_hash + "_job_id"
         output_name = job.job_hash + ".out"
         self.context.write_file(fname=script_file_name, write_str=script_str)
+        script_run_str = self.gen_script_command(job)
+        script_run_file_name = f"{job.script_file_name}.run"
+        self.context.write_file(fname=script_run_file_name, write_str=script_run_str)
 
         resources = job.resources
         submit_command = (

dpdispatcher/{dp_cloud_server.py → machines/dp_cloud_server.py}

@@ -4,11 +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.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
@@ -71,13 +72,28 @@ class Bohrium(Machine):
         return shell_script
 
     def gen_script_header(self, job):
-        shell_script_header = shell_script_header_template
+        resources = job.resources
+        if (
+            resources["strategy"].get("customized_script_header_template_file")
+            is not None
+        ):
+            shell_script_header = customized_script_header_template(
+                resources["strategy"]["customized_script_header_template_file"],
+                resources,
+            )
+        else:
+            shell_script_header = shell_script_header_template
         return shell_script_header
 
     def gen_local_script(self, job):
         script_str = self.gen_script(job)
         script_file_name = job.script_file_name
         self.context.write_local_file(fname=script_file_name, write_str=script_str)
+        script_run_str = self.gen_script_command(job)
+        script_run_file_name = f"{job.script_file_name}.run"
+        self.context.write_local_file(
+            fname=script_run_file_name, write_str=script_run_str
+        )
         return script_file_name
 
     def _gen_backward_files_list(self, job):

dpdispatcher/{fugaku.py → machines/fugaku.py}

@@ -1,8 +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.job_status import JobStatus
+from dpdispatcher.utils.utils import customized_script_header_template
 
 fugaku_script_header_template = """\
 {queue_name_line}
@@ -28,9 +29,18 @@ class Fugaku(Machine):
         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
-        )
+        if (
+            resources["strategy"].get("customized_script_header_template_file")
+            is not None
+        ):
+            fugaku_script_header = customized_script_header_template(
+                resources["strategy"]["customized_script_header_template_file"],
+                resources,
+            )
+        else:
+            fugaku_script_header = fugaku_script_header_template.format(
+                **fugaku_script_header_dict
+            )
         return fugaku_script_header
 
     def do_submit(self, job):
@@ -39,6 +49,9 @@ class Fugaku(Machine):
         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)
+        script_run_str = self.gen_script_command(job)
+        script_run_file_name = f"{job.script_file_name}.run"
+        self.context.write_file(fname=script_run_file_name, write_str=script_run_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

dpdispatcher/{lsf.py → 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, 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
@@ -60,7 +64,16 @@ class LSF(Machine):
                 script_header_dict["lsf_number_gpu_line"] = ""
         else:
             script_header_dict["lsf_number_gpu_line"] = custom_gpu_line
-        lsf_script_header = lsf_script_header_template.format(**script_header_dict)
+        if (
+            resources["strategy"].get("customized_script_header_template_file")
+            is not None
+        ):
+            lsf_script_header = customized_script_header_template(
+                resources["strategy"]["customized_script_header_template_file"],
+                resources,
+            )
+        else:
+            lsf_script_header = lsf_script_header_template.format(**script_header_dict)
 
         return lsf_script_header
 
@@ -70,6 +83,9 @@ class LSF(Machine):
         script_str = self.gen_script(job)
         job_id_name = job.job_hash + "_job_id"
         self.context.write_file(fname=script_file_name, write_str=script_str)
+        script_run_str = self.gen_script_command(job)
+        script_run_file_name = f"{job.script_file_name}.run"
+        self.context.write_file(fname=script_run_file_name, write_str=script_run_str)
 
         try:
             stdin, stdout, stderr = self.context.block_checkcall(

dpdispatcher/{openapi.py → machines/openapi.py}

@@ -2,6 +2,8 @@ import os
 import shutil
 import time
 
+from dpdispatcher.utils.utils import customized_script_header_template
+
 try:
     from bohriumsdk.client import Client
     from bohriumsdk.job import Job
@@ -12,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
@@ -43,13 +45,28 @@ class OpenAPI(Machine):
         return shell_script
 
     def gen_script_header(self, job):
-        shell_script_header = shell_script_header_template
+        resources = job.resources
+        if (
+            resources["strategy"].get("customized_script_header_template_file")
+            is not None
+        ):
+            shell_script_header = customized_script_header_template(
+                resources["strategy"]["customized_script_header_template_file"],
+                resources,
+            )
+        else:
+            shell_script_header = shell_script_header_template
         return shell_script_header
 
     def gen_local_script(self, job):
         script_str = self.gen_script(job)
         script_file_name = job.script_file_name
         self.context.write_local_file(fname=script_file_name, write_str=script_str)
+        script_run_str = self.gen_script_command(job)
+        script_run_file_name = f"{job.script_file_name}.run"
+        self.context.write_local_file(
+            fname=script_run_file_name, write_str=script_run_str
+        )
         return script_file_name
 
     def _gen_backward_files_list(self, job):
@@ -130,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/{pbs.py → machines/pbs.py}

@@ -1,8 +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.job_status import JobStatus
+from dpdispatcher.utils.utils import customized_script_header_template
 
 pbs_script_header_template = """
 #!/bin/bash -l
@@ -28,7 +29,18 @@ class PBS(Machine):
                 "select_node_line"
             ] += f":ngpus={resources.gpu_per_node}"
         pbs_script_header_dict["queue_name_line"] = f"#PBS -q {resources.queue_name}"
-        pbs_script_header = pbs_script_header_template.format(**pbs_script_header_dict)
+        if (
+            resources["strategy"].get("customized_script_header_template_file")
+            is not None
+        ):
+            pbs_script_header = customized_script_header_template(
+                resources["strategy"]["customized_script_header_template_file"],
+                resources,
+            )
+        else:
+            pbs_script_header = pbs_script_header_template.format(
+                **pbs_script_header_dict
+            )
         return pbs_script_header
 
     def do_submit(self, job):
@@ -37,6 +49,9 @@ class PBS(Machine):
         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)
+        script_run_str = self.gen_script_command(job)
+        script_run_file_name = f"{job.script_file_name}.run"
+        self.context.write_file(fname=script_run_file_name, write_str=script_run_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
@@ -149,5 +164,16 @@ class Torque(PBS):
                 gpu_per_node=resources.gpu_per_node
             )
         pbs_script_header_dict["queue_name_line"] = f"#PBS -q {resources.queue_name}"
-        pbs_script_header = pbs_script_header_template.format(**pbs_script_header_dict)
+        if (
+            resources["strategy"].get("customized_script_header_template_file")
+            is not None
+        ):
+            pbs_script_header = customized_script_header_template(
+                resources["strategy"]["customized_script_header_template_file"],
+                resources,
+            )
+        else:
+            pbs_script_header = pbs_script_header_template.format(
+                **pbs_script_header_dict
+            )
         return pbs_script_header

dpdispatcher/{shell.py → machines/shell.py}

@@ -1,8 +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.job_status import JobStatus
+from dpdispatcher.utils.utils import customized_script_header_template
 
 shell_script_header_template = """
 #!/bin/bash -l
@@ -15,7 +16,17 @@ class Shell(Machine):
         return shell_script
 
     def gen_script_header(self, job):
-        shell_script_header = shell_script_header_template
+        resources = job.resources
+        if (
+            resources["strategy"].get("customized_script_header_template_file")
+            is not None
+        ):
+            shell_script_header = customized_script_header_template(
+                resources["strategy"]["customized_script_header_template_file"],
+                resources,
+            )
+        else:
+            shell_script_header = shell_script_header_template
         return shell_script_header
 
     def do_submit(self, job):
@@ -24,6 +35,9 @@ class Shell(Machine):
         job_id_name = job.job_hash + "_job_id"
         output_name = job.job_hash + ".out"
         self.context.write_file(fname=script_file_name, write_str=script_str)
+        script_run_str = self.gen_script_command(job)
+        script_run_file_name = f"{job.script_file_name}.run"
+        self.context.write_file(fname=script_run_file_name, write_str=script_run_str)
         ret, stdin, stdout, stderr = self.context.block_call(
             "cd {} && {{ nohup bash {} 1>>{} 2>>{} & }} && echo $!".format(
                 shlex.quote(self.context.remote_root),