dpdispatcher 0.6.5__py3-none-any.whl → 0.6.7__py3-none-any.whl

dpdispatcher/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.6.5'
-__version_tuple__ = version_tuple = (0, 6, 5)
+__version__ = version = '0.6.7'
+__version_tuple__ = version_tuple = (0, 6, 7)

dpdispatcher/base_context.py

@@ -1,5 +1,5 @@
 from abc import ABCMeta, abstractmethod
-from typing import List, Tuple
+from typing import Any, List, Tuple
 
 from dargs import Argument
 
@@ -73,6 +73,66 @@ class BaseContext(metaclass=ABCMeta):
     def check_finish(self, proc):
         raise NotImplementedError("abstract method")
 
+    def block_checkcall(self, cmd, asynchronously=False) -> Tuple[Any, Any, Any]:
+        """Run command with arguments. Wait for command to complete.
+
+        Parameters
+        ----------
+        cmd : str
+            The command to run.
+        asynchronously : bool, optional, default=False
+            Run command asynchronously. If True, `nohup` will be used to run the command.
+
+        Returns
+        -------
+        stdin
+            standard inout
+        stdout
+            standard output
+        stderr
+            standard error
+
+        Raises
+        ------
+        RuntimeError
+            when the return code is not zero
+        """
+        if asynchronously:
+            cmd = f"nohup {cmd} >/dev/null &"
+        exit_status, stdin, stdout, stderr = self.block_call(cmd)
+        if exit_status != 0:
+            raise RuntimeError(
+                "Get error code %d in calling %s with job: %s . message: %s"
+                % (
+                    exit_status,
+                    cmd,
+                    self.submission.submission_hash,
+                    stderr.read().decode("utf-8"),
+                )
+            )
+        return stdin, stdout, stderr
+
+    @abstractmethod
+    def block_call(self, cmd) -> Tuple[int, Any, Any, Any]:
+        """Run command with arguments. Wait for command to complete.
+
+        Parameters
+        ----------
+        cmd : str
+            The command to run.
+
+        Returns
+        -------
+        exit_status
+            exit code
+        stdin
+            standard inout
+        stdout
+            standard output
+        stderr
+            standard error
+        """
+
     @classmethod
     def machine_arginfo(cls) -> Argument:
         """Generate the machine arginfo.

dpdispatcher/contexts/dp_cloud_server_context.py

@@ -335,6 +335,11 @@ class BohriumContext(BaseContext):
             )
         ]
 
+    def block_call(self, cmd):
+        raise RuntimeError(
+            "Unsupported method. You may use an unsupported combination of the machine and the context."
+        )
+
 
 DpCloudServerContext = BohriumContext
 LebesgueContext = BohriumContext

dpdispatcher/contexts/hdfs_context.py

@@ -244,3 +244,8 @@ class HDFSContext(BaseContext):
 
     def read_file(self, fname):
         return HDFS.read_hdfs_file(os.path.join(self.remote_root, fname))
+
+    def block_call(self, cmd):
+        raise RuntimeError(
+            "Unsupported method. You may use an unsupported combination of the machine and the context."
+        )

dpdispatcher/contexts/lazy_local_context.py

@@ -112,23 +112,6 @@ class LazyLocalContext(BaseContext):
     #                else:
     #                    raise RuntimeError('do not find download file ' + fname)
 
-    def block_checkcall(self, cmd):
-        # script_dir = os.path.join(self.local_root, self.submission.work_base)
-        # os.chdir(script_dir)
-        proc = sp.Popen(
-            cmd, cwd=self.local_root, shell=True, stdout=sp.PIPE, stderr=sp.PIPE
-        )
-        o, e = proc.communicate()
-        stdout = SPRetObj(o)
-        stderr = SPRetObj(e)
-        code = proc.returncode
-        if code != 0:
-            raise RuntimeError(
-                "Get error code %d in locally calling %s with job: %s ",
-                (code, cmd, self.submission.submission_hash),
-            )
-        return None, stdout, stderr
-
     def block_call(self, cmd):
         proc = sp.Popen(
             cmd, cwd=self.local_root, shell=True, stdout=sp.PIPE, stderr=sp.PIPE

dpdispatcher/contexts/local_context.py

@@ -3,6 +3,9 @@ import shutil
 import subprocess as sp
 from glob import glob
 from subprocess import TimeoutExpired
+from typing import List
+
+from dargs import Argument
 
 from dpdispatcher.base_context import BaseContext
 from dpdispatcher.dlog import dlog
@@ -60,6 +63,7 @@ class LocalContext(BaseContext):
         self.temp_local_root = os.path.abspath(local_root)
         self.temp_remote_root = os.path.abspath(remote_root)
         self.remote_profile = remote_profile
+        self.symlink = remote_profile.get("symlink", True)
 
     @classmethod
     def load_from_dict(cls, context_dict):
@@ -83,6 +87,25 @@ class LocalContext(BaseContext):
             self.temp_remote_root, submission.submission_hash
         )
 
+    def _copy_from_local_to_remote(self, local_path, remote_path):
+        if not os.path.exists(local_path):
+            raise FileNotFoundError(
+                f"cannot find uploaded file {os.path.join(local_path)}"
+            )
+        if os.path.exists(remote_path):
+            os.remove(remote_path)
+        _check_file_path(remote_path)
+
+        if self.symlink:
+            # ensure the file exist
+            os.symlink(local_path, remote_path)
+        elif os.path.isfile(local_path):
+            shutil.copyfile(local_path, remote_path)
+        elif os.path.isdir(local_path):
+            shutil.copytree(local_path, remote_path)
+        else:
+            raise ValueError(f"Unknown file type: {local_path}")
+
     def upload(self, submission):
         os.makedirs(self.remote_root, exist_ok=True)
         for ii in submission.belonging_tasks:
@@ -103,14 +126,9 @@ class LocalContext(BaseContext):
                 file_list.extend(rel_file_list)
 
             for jj in file_list:
-                if not os.path.exists(os.path.join(local_job, jj)):
-                    raise FileNotFoundError(
-                        "cannot find upload file " + os.path.join(local_job, jj)
-                    )
-                if os.path.exists(os.path.join(remote_job, jj)):
-                    os.remove(os.path.join(remote_job, jj))
-                _check_file_path(os.path.join(remote_job, jj))
-                os.symlink(os.path.join(local_job, jj), os.path.join(remote_job, jj))
+                self._copy_from_local_to_remote(
+                    os.path.join(local_job, jj), os.path.join(remote_job, jj)
+                )
 
         local_job = self.local_root
         remote_job = self.remote_root
@@ -128,14 +146,9 @@ class LocalContext(BaseContext):
             file_list.extend(rel_file_list)
 
         for jj in file_list:
-            if not os.path.exists(os.path.join(local_job, jj)):
-                raise FileNotFoundError(
-                    "cannot find upload file " + os.path.join(local_job, jj)
-                )
-            if os.path.exists(os.path.join(remote_job, jj)):
-                os.remove(os.path.join(remote_job, jj))
-            _check_file_path(os.path.join(remote_job, jj))
-            os.symlink(os.path.join(local_job, jj), os.path.join(remote_job, jj))
+            self._copy_from_local_to_remote(
+                os.path.join(local_job, jj), os.path.join(remote_job, jj)
+            )
 
     def download(
         self, submission, check_exists=False, mark_failure=True, back_error=False
@@ -288,21 +301,6 @@ class LocalContext(BaseContext):
                 # no nothing in the case of linked files
                 pass
 
-    def block_checkcall(self, cmd):
-        proc = sp.Popen(
-            cmd, cwd=self.remote_root, shell=True, stdout=sp.PIPE, stderr=sp.PIPE
-        )
-        o, e = proc.communicate()
-        stdout = SPRetObj(o)
-        stderr = SPRetObj(e)
-        code = proc.returncode
-        if code != 0:
-            raise RuntimeError(
-                f"Get error code {code} in locally calling {cmd} with job: {self.submission.submission_hash}"
-                f"\nStandard error: {stderr}"
-            )
-        return None, stdout, stderr
-
     def block_call(self, cmd):
         proc = sp.Popen(
             cmd, cwd=self.remote_root, shell=True, stdout=sp.PIPE, stderr=sp.PIPE
@@ -351,3 +349,31 @@ class LocalContext(BaseContext):
                 stdout = None
                 stderr = None
         return ret, stdout, stderr
+
+    @classmethod
+    def machine_subfields(cls) -> List[Argument]:
+        """Generate the machine subfields.
+
+        Returns
+        -------
+        list[Argument]
+            machine subfields
+        """
+        doc_remote_profile = "The information used to maintain the local machine."
+        return [
+            Argument(
+                "remote_profile",
+                dict,
+                optional=True,
+                doc=doc_remote_profile,
+                sub_fields=[
+                    Argument(
+                        "symlink",
+                        bool,
+                        optional=True,
+                        default=True,
+                        doc="Whether to use symbolic links to replace copy. This option should be turned off if the local directory is not accessible on the Batch system.",
+                    ),
+                ],
+            )
+        ]

dpdispatcher/contexts/openapi_context.py

@@ -258,3 +258,8 @@ class OpenAPIContext(BaseContext):
             dir_to_be_removed = os.path.join(local_root, "backup")
             if os.path.exists(dir_to_be_removed):
                 shutil.rmtree(dir_to_be_removed)
+
+    def block_call(self, cmd):
+        raise RuntimeError(
+            "Unsupported method. You may use an unsupported combination of the machine and the context."
+        )

dpdispatcher/contexts/ssh_context.py

@@ -44,6 +44,7 @@ class SSHSession:
         totp_secret=None,
         tar_compress=True,
         look_for_keys=True,
+        execute_command=None,
     ):
         self.hostname = hostname
         self.username = username
@@ -56,6 +57,7 @@ class SSHSession:
         self.ssh = None
         self.tar_compress = tar_compress
         self.look_for_keys = look_for_keys
+        self.execute_command = execute_command
         self._keyboard_interactive_auth = False
         self._setup_ssh()
 
@@ -237,6 +239,8 @@ class SSHSession:
         self.ssh._transport = ts  # type: ignore
         # reset sftp
         self._sftp = None
+        if self.execute_command is not None:
+            self.exec_command(self.execute_command)
 
     def inter_handler(self, title, instructions, prompt_list):
         """inter_handler: the callback for paramiko.transport.auth_interactive.
@@ -295,7 +299,11 @@ class SSHSession:
         assert self.ssh is not None
         try:
             return self.ssh.exec_command(cmd)
-        except (paramiko.ssh_exception.SSHException, socket.timeout, EOFError) as e:
+        except (
+            paramiko.ssh_exception.SSHException,
+            socket.timeout,
+            EOFError,
+        ) as e:
             # SSH session not active
             # retry for up to 3 times
             # ensure alive
@@ -334,6 +342,7 @@ class SSHSession:
         doc_look_for_keys = (
             "enable searching for discoverable private key files in ~/.ssh/"
         )
+        doc_execute_command = "execute command after ssh connection is established."
         ssh_remote_profile_args = [
             Argument("hostname", str, optional=False, doc=doc_hostname),
             Argument("username", str, optional=False, doc=doc_username),
@@ -355,10 +364,18 @@ class SSHSession:
             ),
             Argument("timeout", int, optional=True, default=10, doc=doc_timeout),
             Argument(
-                "totp_secret", str, optional=True, default=None, doc=doc_totp_secret
+                "totp_secret",
+                str,
+                optional=True,
+                default=None,
+                doc=doc_totp_secret,
             ),
             Argument(
-                "tar_compress", bool, optional=True, default=True, doc=doc_tar_compress
+                "tar_compress",
+                bool,
+                optional=True,
+                default=True,
+                doc=doc_tar_compress,
             ),
             Argument(
                 "look_for_keys",
@@ -367,6 +384,13 @@ class SSHSession:
                 default=True,
                 doc=doc_look_for_keys,
             ),
+            Argument(
+                "execute_command",
+                str,
+                optional=True,
+                default=None,
+                doc=doc_execute_command,
+            ),
         ]
         ssh_remote_profile_format = Argument(
             "ssh_session", dict, ssh_remote_profile_args
@@ -603,7 +627,10 @@ class SSHContext(BaseContext):
                 directory_list,
             )
         self._walk_directory(
-            submission.forward_common_files, self.local_root, file_list, directory_list
+            submission.forward_common_files,
+            self.local_root,
+            file_list,
+            directory_list,
         )
 
         # convert to relative path to local_root
@@ -621,9 +648,9 @@ class SSHContext(BaseContext):
                 ).as_posix()
                 sha256_list.append(f"{sha256}  {jj_rel}")
             # write to remote
-            sha256_file = os.path.join(
-                self.remote_root, ".tmp.sha256." + str(uuid.uuid4())
-            )
+            sha256_file = pathlib.PurePath(
+                os.path.join(self.remote_root, ".tmp.sha256." + str(uuid.uuid4()))
+            ).as_posix()
             self.write_file(sha256_file, "\n".join(sha256_list))
             # check sha256
             # `:` means pass: https://stackoverflow.com/a/2421592/9567349
@@ -736,43 +763,9 @@ class SSHContext(BaseContext):
         file_list.extend(submission.backward_common_files)
         if len(file_list) > 0:
             self._get_files(
-                file_list, tar_compress=self.remote_profile.get("tar_compress", None)
-            )
-
-    def block_checkcall(self, cmd, asynchronously=False, stderr_whitelist=None):
-        """Run command with arguments. Wait for command to complete. If the return code
-        was zero then return, otherwise raise RuntimeError.
-
-        Parameters
-        ----------
-        cmd : str
-            The command to run.
-        asynchronously : bool, optional, default=False
-            Run command asynchronously. If True, `nohup` will be used to run the command.
-        stderr_whitelist : list of str, optional, default=None
-            If not None, the stderr will be checked against the whitelist. If the stderr
-            contains any of the strings in the whitelist, the command will be considered
-            successful.
-        """
-        assert self.remote_root is not None
-        self.ssh_session.ensure_alive()
-        if asynchronously:
-            cmd = f"nohup {cmd} >/dev/null &"
-        stdin, stdout, stderr = self.ssh_session.exec_command(
-            (f"cd {shlex.quote(self.remote_root)} ;") + cmd
-        )
-        exit_status = stdout.channel.recv_exit_status()
-        if exit_status != 0:
-            raise RuntimeError(
-                "Get error code %d in calling %s through ssh with job: %s . message: %s"
-                % (
-                    exit_status,
-                    cmd,
-                    self.submission.submission_hash,
-                    stderr.read().decode("utf-8"),
-                )
+                file_list,
+                tar_compress=self.remote_profile.get("tar_compress", None),
             )
-        return stdin, stdout, stderr
 
     def block_call(self, cmd):
         assert self.remote_root is not None
@@ -793,18 +786,23 @@ class SSHContext(BaseContext):
         fname = pathlib.PurePath(os.path.join(self.remote_root, fname)).as_posix()
         # to prevent old file from being overwritten but cancelled, create a temporary file first
         # when it is fully written, rename it to the original file name
-        with self.sftp.open(fname + "~", "w") as fp:
-            fp.write(write_str)
+        temp_fname = fname + "_tmp"
+        try:
+            with self.sftp.open(temp_fname, "w") as fp:
+                fp.write(write_str)
+            # Rename the temporary file
+            self.block_checkcall(f"mv {shlex.quote(temp_fname)} {shlex.quote(fname)}")
         # sftp.rename may throw OSError
-        self.block_checkcall(
-            "mv {} {}".format(shlex.quote(fname + "~"), shlex.quote(fname))
-        )
+        except OSError as e:
+            dlog.exception(f"Error writing to file {fname}")
+            raise e
 
     def read_file(self, fname):
         assert self.remote_root is not None
         self.ssh_session.ensure_alive()
         with self.sftp.open(
-            pathlib.PurePath(os.path.join(self.remote_root, fname)).as_posix(), "r"
+            pathlib.PurePath(os.path.join(self.remote_root, fname)).as_posix(),
+            "r",
         ) as fp:
             ret = fp.read().decode("utf-8")
         return ret
@@ -945,36 +943,28 @@ class SSHContext(BaseContext):
         per_nfile = 100
         ntar = len(files) // per_nfile + 1
         if ntar <= 1:
-            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
+            file_list = " ".join([shlex.quote(file) for file in files])
+            tar_cmd = f"tar {tar_command} {shlex.quote(of)} {file_list}"
         else:
-            file_list_file = os.path.join(
-                self.remote_root, ".tmp.tar." + str(uuid.uuid4())
-            )
+            file_list_file = pathlib.PurePath(
+                os.path.join(self.remote_root, f".tmp_tar_{uuid.uuid4()}")
+            ).as_posix()
             self.write_file(file_list_file, "\n".join(files))
-            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
+            tar_cmd = (
+                f"tar {tar_command} {shlex.quote(of)} -T {shlex.quote(file_list_file)}"
+            )
+
+        # Execute the tar command remotely
+        try:
+            self.block_checkcall(tar_cmd)
+        except RuntimeError as e:
+            if "No such file or directory" in str(e):
+                raise FileNotFoundError(
+                    "Backward files do not exist in the remote directory."
+                ) from e
+            raise e
+
+        # Transfer the archive from remote to local
         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()
         if os.path.isfile(to_f):

dpdispatcher/machine.py

@@ -161,6 +161,9 @@ class Machine(metaclass=ABCMeta):
             machine_dict["remote_profile"] = self.context.remote_profile
         else:
             machine_dict["remote_profile"] = {}
+        # normalize the dict
+        base = self.arginfo()
+        machine_dict = base.normalize_value(machine_dict, trim_pattern="_*")
         return machine_dict
 
     def __eq__(self, other):
@@ -261,11 +264,19 @@ class Machine(metaclass=ABCMeta):
 
         source_list = job.resources.source_list
         for ii in source_list:
-            line = f"{{ source {ii}; }} \n"
-            source_files_part += line
+            source_files_part += f"source {ii}\n"
 
         export_envs_part = ""
         envs = job.resources.envs
+        envs = {
+            # export resources information to the environment variables
+            "DPDISPATCHER_NUMBER_NODE": job.resources.number_node,
+            "DPDISPATCHER_CPU_PER_NODE": job.resources.cpu_per_node,
+            "DPDISPATCHER_GPU_PER_NODE": job.resources.gpu_per_node,
+            "DPDISPATCHER_QUEUE_NAME": job.resources.queue_name,
+            "DPDISPATCHER_GROUP_SIZE": job.resources.group_size,
+            **envs,
+        }
         for k, v in envs.items():
             if isinstance(v, list):
                 for each_value in v:

dpdispatcher/machines/JH_UniScheduler.py

@@ -105,7 +105,7 @@ class JH_UniScheduler(Machine):
         elif ret != 0:
             # just retry when any unknown error raised.
             raise RetrySignal(
-                "Get error code %d in checking status through ssh with job: %s . message: %s"
+                "Get error code %d in checking status with job: %s . message: %s"
                 % (ret, job.job_hash, err_str)
             )
         status_out = stdout.read().decode("utf-8").split("\n")

dpdispatcher/machines/distributed_shell.py

@@ -181,8 +181,8 @@ class DistributedShell(Machine):
         if ret != 0:
             err_str = stderr.decode("utf-8")
             raise RuntimeError(
-                "Command squeue fails to execute, error message:%s\nreturn code %d\n"
-                % (err_str, ret)
+                "Command %s fails to execute, error message:%s\nreturn code %d\n"
+                % (cmd, err_str, ret)
             )
         job_id = int(stdout.decode("utf-8").strip())
 

dpdispatcher/machines/lsf.py

@@ -129,7 +129,7 @@ class LSF(Machine):
         elif ret != 0:
             # just retry when any unknown error raised.
             raise RetrySignal(
-                "Get error code %d in checking status through ssh with job: %s . message: %s"
+                "Get error code %d in checking status with job: %s . message: %s"
                 % (ret, job.job_hash, err_str)
             )
         status_out = stdout.read().decode("utf-8").split("\n")

dpdispatcher/machines/pbs.py

@@ -1,4 +1,7 @@
 import shlex
+from typing import List
+
+from dargs import Argument
 
 from dpdispatcher.dlog import dlog
 from dpdispatcher.machine import Machine
@@ -73,7 +76,8 @@ class PBS(Machine):
         job_id = job.job_id
         if job_id == "":
             return JobStatus.unsubmitted
-        ret, stdin, stdout, stderr = self.context.block_call("qstat -x " + job_id)
+        command = "qstat -x " + job_id
+        ret, stdin, stdout, stderr = self.context.block_call(command)
         err_str = stderr.read().decode("utf-8")
         if ret != 0:
             if "qstat: Unknown Job Id" in err_str or "Job has finished" in err_str:
@@ -83,8 +87,8 @@ class PBS(Machine):
                     return JobStatus.terminated
             else:
                 raise RuntimeError(
-                    "status command qstat fails to execute. erro info: %s return code %d"
-                    % (err_str, ret)
+                    "status command %s fails to execute. erro info: %s return code %d"
+                    % (command, err_str, ret)
                 )
         status_line = stdout.read().decode("utf-8").split("\n")[-2]
         status_word = status_line.split()[-2]
@@ -123,7 +127,8 @@ class Torque(PBS):
         job_id = job.job_id
         if job_id == "":
             return JobStatus.unsubmitted
-        ret, stdin, stdout, stderr = self.context.block_call("qstat -l " + job_id)
+        command = "qstat -l " + job_id
+        ret, stdin, stdout, stderr = self.context.block_call(command)
         err_str = stderr.read().decode("utf-8")
         if ret != 0:
             if "qstat: Unknown Job Id" in err_str or "Job has finished" in err_str:
@@ -133,8 +138,8 @@ class Torque(PBS):
                     return JobStatus.terminated
             else:
                 raise RuntimeError(
-                    "status command qstat fails to execute. erro info: %s return code %d"
-                    % (err_str, ret)
+                    "status command %s fails to execute. erro info: %s return code %d"
+                    % (command, err_str, ret)
                 )
         status_line = stdout.read().decode("utf-8").split("\n")[-2]
         status_word = status_line.split()[-2]
@@ -181,10 +186,9 @@ class Torque(PBS):
 
 sge_script_header_template = """
 #!/bin/bash
-#$ -N dpdispatcher_submit
-{select_node_line}
+#$ -S /bin/bash
 #$ -cwd
-
+{select_node_line}
 """
 
 
@@ -209,14 +213,31 @@ class SGE(PBS):
         )
 
     def gen_script_header(self, job):
+        ### Ref:https://softpanorama.org/HPC/PBS_and_derivatives/Reference/pbs_command_vs_sge_commands.shtml
+        # resources.number_node is not used in SGE
         resources = job.resources
+        job_name = resources.kwargs.get("job_name", "wDPjob")
+        pe_name = resources.kwargs.get("pe_name", "mpi")
         sge_script_header_dict = {}
-        # resources.number_node is not used
-        sge_script_header_dict["select_node_line"] = (
-            f"#$ -pe mpi {resources.cpu_per_node} "
+        sge_script_header_dict["select_node_line"] = f"#$ -N {job_name}\n"
+        sge_script_header_dict["select_node_line"] += (
+            f"#$ -pe {pe_name} {resources.cpu_per_node}\n"
         )
-        # resources.queue_name is not necessary
-        sge_script_header = sge_script_header_template.format(**sge_script_header_dict)
+
+        if resources.queue_name != "":
+            sge_script_header_dict["select_node_line"] += (
+                f"#$ -q {resources.queue_name}"
+            )
+        if (
+            resources["strategy"].get("customized_script_header_template_file")
+            is not None
+        ):
+            file_name = resources["strategy"]["customized_script_header_template_file"]
+            sge_script_header = customized_script_header_template(file_name, resources)
+        else:
+            sge_script_header = sge_script_header_template.format(
+                **sge_script_header_dict
+            )
         return sge_script_header
 
     def do_submit(self, job):
@@ -224,6 +245,9 @@ class SGE(PBS):
         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)
         script_file_dir = self.context.remote_root
         stdin, stdout, stderr = self.context.block_checkcall(
             "cd {} && {} {}".format(script_file_dir, "qsub", script_file_name)
@@ -241,12 +265,12 @@ class SGE(PBS):
         status_line = None
         if job_id == "":
             return JobStatus.unsubmitted
-        ret, stdin, stdout, stderr = self.context.block_call("qstat")
+        command = "qstat"
+        ret, stdin, stdout, stderr = self.context.block_call(command)
         err_str = stderr.read().decode("utf-8")
         if ret != 0:
             raise RuntimeError(
-                "status command qstat fails to execute. erro info: %s return code %d"
-                % (err_str, ret)
+                f"status command {command} fails to execute. erro info: {err_str} return code {ret}"
             )
         status_text_list = stdout.read().decode("utf-8").split("\n")
         for txt in status_text_list:
@@ -259,8 +283,7 @@ class SGE(PBS):
                 if self.check_finish_tag(job=job):
                     return JobStatus.finished
                 dlog.info(
-                    "not tag_finished detected, execute sync command and wait. count "
-                    + str(count)
+                    f"not tag_finished detected, execute sync command and wait. count {count}"
                 )
                 self.context.block_call("sync")
                 import time
@@ -281,3 +304,44 @@ class SGE(PBS):
     def check_finish_tag(self, job):
         job_tag_finished = job.job_hash + "_job_tag_finished"
         return self.context.check_file_exists(job_tag_finished)
+
+    @classmethod
+    def resources_subfields(cls) -> List[Argument]:
+        """Generate the resources subfields.
+
+            pe_name : str
+        The parallel environment name of SGE.
+
+        Returns
+        -------
+        list[Argument]
+            resources subfields
+        """
+        doc_pe_name = "The parallel environment name of SGE system."
+        doc_job_name = "The name of SGE's job."
+
+        return [
+            Argument(
+                "kwargs",
+                dict,
+                [
+                    Argument(
+                        "pe_name",
+                        str,
+                        optional=True,
+                        default="mpi",
+                        doc=doc_pe_name,
+                        alias=["sge_pe_name"],
+                    ),
+                    Argument(
+                        "job_name",
+                        str,
+                        optional=True,
+                        default="wDPjob",
+                        doc=doc_job_name,
+                    ),
+                ],
+                optional=False,
+                doc="Extra arguments.",
+            )
+        ]

dpdispatcher/machines/shell.py

@@ -38,14 +38,13 @@ class Shell(Machine):
         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(
-            f"cd {shlex.quote(self.context.remote_root)} && {{ nohup bash {script_file_name} 1>>{output_name} 2>>{output_name} & }} && echo $!"
-        )
+        cmd = f"cd {shlex.quote(self.context.remote_root)} && {{ nohup bash {script_file_name} 1>>{output_name} 2>>{output_name} & }} && echo $!"
+        ret, stdin, stdout, stderr = self.context.block_call(cmd)
         if ret != 0:
             err_str = stderr.read().decode("utf-8")
             raise RuntimeError(
-                "status command squeue fails to execute\nerror message:%s\nreturn code %d\n"
-                % (err_str, ret)
+                "status command %s fails to execute\nerror message:%s\nreturn code %d\n"
+                % (cmd, err_str, ret)
             )
         job_id = int(stdout.read().decode("utf-8").strip())
         self.context.write_file(job_id_name, str(job_id))
@@ -73,14 +72,16 @@ class Shell(Machine):
             return JobStatus.unsubmitted
 
         # mark defunct process as terminated
-        ret, stdin, stdout, stderr = self.context.block_call(
+        cmd = (
+            r"""command -v ps >/dev/null 2>&1 || { echo >&2 "I require ps but it's not installed. Aborting."; exit 1; };"""
             f"if ps -p {job_id} > /dev/null && ! (ps -o command -p {job_id} | grep defunct >/dev/null) ; then echo 1; fi"
         )
+        ret, stdin, stdout, stderr = self.context.block_call(cmd)
         if ret != 0:
             err_str = stderr.read().decode("utf-8")
             raise RuntimeError(
-                "status command squeue fails to execute\nerror message:%s\nreturn code %d\n"
-                % (err_str, ret)
+                "status command %s fails to execute\nerror message:%s\nreturn code %d\n"
+                % (cmd, err_str, ret)
             )
 
         if_job_exists = bool(stdout.read().decode("utf-8").strip())

dpdispatcher/machines/slurm.py

@@ -83,13 +83,12 @@ class Slurm(Machine):
         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)
-        ret, stdin, stdout, stderr = self.context.block_call(
-            "cd {} && {} {}".format(
-                shlex.quote(self.context.remote_root),
-                "sbatch",
-                shlex.quote(script_file_name),
-            )
+        command = "cd {} && {} {}".format(
+            shlex.quote(self.context.remote_root),
+            "sbatch",
+            shlex.quote(script_file_name),
         )
+        ret, stdin, stdout, stderr = self.context.block_call(command)
         if ret != 0:
             err_str = stderr.read().decode("utf-8")
             if (
@@ -98,7 +97,7 @@ class Slurm(Machine):
             ):
                 # server network error, retry 3 times
                 raise RetrySignal(
-                    "Get error code %d in submitting through ssh with job: %s . message: %s"
+                    "Get error code %d in submitting with job: %s . message: %s"
                     % (ret, job.job_hash, err_str)
                 )
             elif (
@@ -110,8 +109,8 @@ class Slurm(Machine):
                 # job number exceeds, skip the submitting
                 return ""
             raise RuntimeError(
-                "status command squeue fails to execute\nerror message:%s\nreturn code %d\n"
-                % (err_str, ret)
+                "command %s fails to execute\nerror message:%s\nreturn code %d\n"
+                % (command, err_str, ret)
             )
         subret = stdout.readlines()
         # --parsable
@@ -129,9 +128,8 @@ class Slurm(Machine):
         job_id = job.job_id
         if job_id == "":
             return JobStatus.unsubmitted
-        ret, stdin, stdout, stderr = self.context.block_call(
-            'squeue -o "%.18i %.2t" -j ' + job_id
-        )
+        command = 'squeue -o "%.18i %.2t" -j ' + job_id
+        ret, stdin, stdout, stderr = self.context.block_call(command)
         if ret != 0:
             err_str = stderr.read().decode("utf-8")
             if "Invalid job id specified" in err_str:
@@ -147,13 +145,13 @@ class Slurm(Machine):
             ):
                 # retry 3 times
                 raise RetrySignal(
-                    "Get error code %d in checking status through ssh with job: %s . message: %s"
+                    "Get error code %d in checking status with job: %s . message: %s"
                     % (ret, job.job_hash, err_str)
                 )
             raise RuntimeError(
-                "status command squeue fails to execute."
+                "status command %s fails to execute."
                 "job_id:%s \n error message:%s\n return code %d\n"
-                % (job_id, err_str, ret)
+                % (command, job_id, err_str, ret)
             )
         status_line = stdout.read().decode("utf-8").split("\n")[-2]
         status_word = status_line.split()[-1]
@@ -319,9 +317,8 @@ class SlurmJobArray(Slurm):
         job_id = job.job_id
         if job_id == "":
             return JobStatus.unsubmitted
-        ret, stdin, stdout, stderr = self.context.block_call(
-            'squeue -h -o "%.18i %.2t" -j ' + job_id
-        )
+        command = 'squeue -h -o "%.18i %.2t" -j ' + job_id
+        ret, stdin, stdout, stderr = self.context.block_call(command)
         if ret != 0:
             err_str = stderr.read().decode("utf-8")
             if "Invalid job id specified" in err_str:
@@ -336,13 +333,13 @@ class SlurmJobArray(Slurm):
             ):
                 # retry 3 times
                 raise RetrySignal(
-                    "Get error code %d in checking status through ssh with job: %s . message: %s"
+                    "Get error code %d in checking status with job: %s . message: %s"
                     % (ret, job.job_hash, err_str)
                 )
             raise RuntimeError(
-                "status command squeue fails to execute."
+                "status command %s fails to execute."
                 "job_id:%s \n error message:%s\n return code %d\n"
-                % (job_id, err_str, ret)
+                % (command, job_id, err_str, ret)
             )
         status_lines = stdout.read().decode("utf-8").split("\n")[:-1]
         status = []

dpdispatcher/submission.py

@@ -1122,9 +1122,9 @@ class Resources:
 
     @staticmethod
     def arginfo(detail_kwargs=True):
-        doc_number_node = "The number of node need for each `job`"
-        doc_cpu_per_node = "cpu numbers of each node assigned to each job."
-        doc_gpu_per_node = "gpu numbers of each node assigned to each job."
+        doc_number_node = "The number of nodes required for each `job`."
+        doc_cpu_per_node = "CPU numbers of each node assigned to each job."
+        doc_gpu_per_node = "GPU numbers of each node assigned to each job."
         doc_queue_name = "The queue name of batch job scheduler system."
         doc_group_size = "The number of `tasks` in a `job`. 0 means infinity."
         doc_custom_flags = "The extra lines pass to job submitting script header"

{dpdispatcher-0.6.5.dist-info → dpdispatcher-0.6.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dpdispatcher
-Version: 0.6.5
+Version: 0.6.7
 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
@@ -207,7 +207,7 @@ Requires-Dist: sphinx-rtd-theme >=1.0.0rc1 ; extra == 'docs'
 Requires-Dist: numpydoc ; extra == 'docs'
 Requires-Dist: deepmodeling-sphinx >=0.1.1 ; extra == 'docs'
 Requires-Dist: dargs >=0.3.1 ; extra == 'docs'
-Requires-Dist: sphinx-argparse ; extra == 'docs'
+Requires-Dist: sphinx-argparse <0.5.0 ; extra == 'docs'
 Provides-Extra: gui
 Requires-Dist: dpgui ; extra == 'gui'
 Provides-Extra: test

{dpdispatcher-0.6.5.dist-info → dpdispatcher-0.6.7.dist-info}/RECORD

@@ -1,36 +1,36 @@
 dpdispatcher/__init__.py,sha256=CLZP_N5CTp14ujWCykEHuJjoIfKR6CwrclXhjWUgNoE,517
 dpdispatcher/__main__.py,sha256=BFhG-mSBzVZUEezQJqXWZnt2WsnhAHT_zpT8Y6gpOz0,116
-dpdispatcher/_version.py,sha256=PuC6q1U5hHaOMp2tDNeTKt6ExeuO2V9ihjqjMYIsVUo,411
+dpdispatcher/_version.py,sha256=iLXz9haw4jSV4Xm2-5_V8999GBAYoJkXg9-YOwMJpLY,411
 dpdispatcher/arginfo.py,sha256=pNaxYIE6ahBidpR7OCKZdw8iGt003uTXGSlVzwiuvRg,188
-dpdispatcher/base_context.py,sha256=NvaC_RHyspxq412z-eCq4Zn8-szZxvn8K6OkXvx7l4Y,3615
+dpdispatcher/base_context.py,sha256=FDie50yHuLxTwV-k3B_sgAYqR5rLLLVOwk24wSiu4SQ,5254
 dpdispatcher/dlog.py,sha256=QJKAwB6gV3Zb6zQUL9dZ_uIoTIEy9Z7ecmVQ-8WNmD8,1081
 dpdispatcher/dpdisp.py,sha256=jhuTmwPY7KBF4WukaQomEwZcfYoISaMbKwuxdDGSluc,4206
-dpdispatcher/machine.py,sha256=z5D0eLAPfdo5SZdO6NLvWBUUePE0VHRMWurRMzEV0U0,16138
+dpdispatcher/machine.py,sha256=k53ycs_v7xrl4D93URc5ht0shoO9NPrVl0rYr4v5OiU,16696
 dpdispatcher/run.py,sha256=tFHbJAioXXpgHTE5bhRRAuc8w7cX1ET9SBbiAg3Rw-I,5382
-dpdispatcher/submission.py,sha256=0_PCpRyiUwCHwYAzdXs-3rzq8YzZs0VZBU6tS7SixG0,48361
+dpdispatcher/submission.py,sha256=NaljgA88NLv0rvxoSMZvUMq0sQEggkgKlcT8gXUnqFs,48367
 dpdispatcher/contexts/__init__.py,sha256=jlvcIppmUnS39yBlkZEDvIQFV-j_BR75ZTbZALF_RB0,336
-dpdispatcher/contexts/dp_cloud_server_context.py,sha256=6XK0B2sLGEDeZmV2SZzQdVrMcWAWYZVLLK-IaShEXIY,12245
-dpdispatcher/contexts/hdfs_context.py,sha256=B6pjGUD8Xaa0G_Zrnoci2DZnEXxojE9fAcexMMvAZCM,8930
-dpdispatcher/contexts/lazy_local_context.py,sha256=F8abWAJRY1Ewx1sErINKN1ltWerXzeCcJgjTvLvucKE,5696
-dpdispatcher/contexts/local_context.py,sha256=AsIfOT24FV0_bNlD2xU-pqAJy-XHZ6XTsbll4Vt6bMM,14065
-dpdispatcher/contexts/openapi_context.py,sha256=DXaMS10SXN3VKEeEdzQyfOgRwUyHRJVCJHd2fKKdsmA,9499
-dpdispatcher/contexts/ssh_context.py,sha256=baMiD_1KlrksqNKCkpx7apovLW_qdfU9U1KRDNTjCz0,38578
+dpdispatcher/contexts/dp_cloud_server_context.py,sha256=PGRMef3q2hfK-o5dNIWWvzPca2NK1HrWEgungM4L9Go,12420
+dpdispatcher/contexts/hdfs_context.py,sha256=mYQzXMZ4A9EjjWBAH3Ba6HOErUhMMwCsKxOjpd5R57Y,9105
+dpdispatcher/contexts/lazy_local_context.py,sha256=FAClbLD2F4LizUqFzMOg3t0Z6NLeTDLJy7NkRcDELFs,5070
+dpdispatcher/contexts/local_context.py,sha256=VbaSXGAc_EDMT0K5WV_flBF0bX87ntrwO_hq_Bkcb04,14590
+dpdispatcher/contexts/openapi_context.py,sha256=M7L9axpjOrzvdTpLMDuEzZqe4ZuKIxjS0bzZUv8W2IQ,9674
+dpdispatcher/contexts/ssh_context.py,sha256=s0K-gSKPSykq2PyOzAt4yNEczAdsVGvQ1QmPJpZ4_Vo,37648
 dpdispatcher/dpcloudserver/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dpdispatcher/dpcloudserver/client.py,sha256=k1niKjG6zFnMtHn_UuCjYoOcMju3o3PV-GdyVLr5-KM,165
 dpdispatcher/entrypoints/__init__.py,sha256=exKSFT3j2oCerGwtI8WbHQK-D0K-CyifocRji1xntT4,20
 dpdispatcher/entrypoints/gui.py,sha256=29lMXqbmSRbLj4rfBv7Jnw89NLU9syTB88IUP6IRJsU,830
 dpdispatcher/entrypoints/run.py,sha256=tRkHfeAktV6gF31yb2MVOSTlpNGZFw3N0jHBmM1YfIg,175
 dpdispatcher/entrypoints/submission.py,sha256=ikVwIZAQL0SsYO5xaMIdKXgO6qtc05w1vqmvtG7Nk5M,3401
-dpdispatcher/machines/JH_UniScheduler.py,sha256=f7Vs9_m4Th1GVSgsJTy9_nMAY8g9n0ZewnPY2DFECfI,5795
+dpdispatcher/machines/JH_UniScheduler.py,sha256=B-LGldr9H8qPQYdCYoEaXFCEFBPmjFEi0fwEWp0wdR0,5783
 dpdispatcher/machines/__init__.py,sha256=tOQuPUlW1Ab4qcC0oSAIyDjZA_WyE67h_EIxPCWGhys,336
-dpdispatcher/machines/distributed_shell.py,sha256=LvWl6ktPlgmJ7rk90VWxp4douve8hYmuRf-B0saFBds,7534
+dpdispatcher/machines/distributed_shell.py,sha256=TVnXFNqQmBgWk3s34rKSZo0S5N5KPZVmAG3Xbu_kuBo,7535
 dpdispatcher/machines/dp_cloud_server.py,sha256=SR69gsFb2BvOQCW1QnWfP3cQvu_qHLJNsycp5wzosJU,11706
 dpdispatcher/machines/fugaku.py,sha256=oY2hD2ldL2dztwtJ9WNisdsfPnaX-5yTRXewIT9r60I,4314
-dpdispatcher/machines/lsf.py,sha256=Q6IE4nCkNEKcW0AdBTKPOYgmCJAeXWmUVxZ9sQFkxos,7932
+dpdispatcher/machines/lsf.py,sha256=fOZoOTpFn1nKx79lYkvZQOhNwz39YAIEytxICd56AFU,7920
 dpdispatcher/machines/openapi.py,sha256=Gzzbo8YOAybXGTrgMutexErcaEi3ts7uTUNvOhThFS8,8858
-dpdispatcher/machines/pbs.py,sha256=KjJcLpQr748ZgOwFfWmJ_LG1q6Jm1UF24YCSLiDfcac,10308
-dpdispatcher/machines/shell.py,sha256=DnqMNb2nmBc3gVx8tA8oiUWdnWHKJwpIPs660i3Eq7A,4703
-dpdispatcher/machines/slurm.py,sha256=YM2Mv55jAFtDIiJoJLkD6p1Wi1ujjH6t4WlU8EtlbCw,15592
+dpdispatcher/machines/pbs.py,sha256=XeeFQMZoH9DscsrJ_Ykv6fNUtc9TBp4epuFqbUyr3dk,12531
+dpdispatcher/machines/shell.py,sha256=ONaUJpszsCwCcbyVLvC6VoJ-ig2QTU9JQdA-nlgXnu8,4845
+dpdispatcher/machines/slurm.py,sha256=HLYk9E1dChnTeHjOOWNG854AWdlUJVYYmgwaiVswPQ8,15560
 dpdispatcher/utils/__init__.py,sha256=fwvwkMf7DFNQkNBiIce8Y8gRA6FhICwKjkKiXu_BEJg,13
 dpdispatcher/utils/hdfs_cli.py,sha256=n3EIfFIralsISlaEewawD35f0P8mabo-u8D8UW3k_7Y,5308
 dpdispatcher/utils/job_status.py,sha256=Eszs4TPLfszCuf6zLaFonf25feXDUguF28spYOjJpQE,233
@@ -41,9 +41,9 @@ dpdispatcher/utils/dpcloudserver/client.py,sha256=CLfXswvzI4inDrW2bYkfMQ6gQJFcZO
 dpdispatcher/utils/dpcloudserver/config.py,sha256=NteQzf1OeEkz2UbkXHHQ0B72cUu23zLVzpM9Yh4v1Cc,559
 dpdispatcher/utils/dpcloudserver/retcode.py,sha256=1qAF8gFZx55u2sO8KbtYSIIrjcO-IGufEUlwbkSfC1g,721
 dpdispatcher/utils/dpcloudserver/zip_file.py,sha256=f9WrlktwHW0YipaWg5Y0kxjMZlhD1cJYa6EUpvu4Cro,2611
-dpdispatcher-0.6.5.dist-info/LICENSE,sha256=46mU2C5kSwOnkqkw9XQAJlhBL2JAf1_uCD8lVcXyMRg,7652
-dpdispatcher-0.6.5.dist-info/METADATA,sha256=eLIZlw1J6l08VjltNG2O3Z7kWK_TNVJR08aaGlfwESc,12821
-dpdispatcher-0.6.5.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-dpdispatcher-0.6.5.dist-info/entry_points.txt,sha256=NRHUV0IU_u7_XtcmmEDnVzAcUmurhiEAGwENckrajo4,233
-dpdispatcher-0.6.5.dist-info/top_level.txt,sha256=35jAQoXY-b-e9fJ1_mxhZUiaCoJNt1ZI7mpFRf07Qjs,13
-dpdispatcher-0.6.5.dist-info/RECORD,,
+dpdispatcher-0.6.7.dist-info/LICENSE,sha256=46mU2C5kSwOnkqkw9XQAJlhBL2JAf1_uCD8lVcXyMRg,7652
+dpdispatcher-0.6.7.dist-info/METADATA,sha256=lNsC7Ruo7GmUOQl1TadlThoLvrOETbZ1s0-sXmrRYL4,12828
+dpdispatcher-0.6.7.dist-info/WHEEL,sha256=UvcQYKBHoFqaQd6LKyqHw9fxEolWLQnlzP0h_LgJAfI,91
+dpdispatcher-0.6.7.dist-info/entry_points.txt,sha256=NRHUV0IU_u7_XtcmmEDnVzAcUmurhiEAGwENckrajo4,233
+dpdispatcher-0.6.7.dist-info/top_level.txt,sha256=35jAQoXY-b-e9fJ1_mxhZUiaCoJNt1ZI7mpFRf07Qjs,13
+dpdispatcher-0.6.7.dist-info/RECORD,,

{dpdispatcher-0.6.5.dist-info → dpdispatcher-0.6.7.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.43.0)
+Generator: setuptools (74.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any