dpdispatcher 0.6.5__tar.gz → 0.6.6__tar.gz

{dpdispatcher-0.6.5 → dpdispatcher-0.6.6}/.git_archival.txt

@@ -1,4 +1,3 @@
 node: $Format:%H$
 node-date: $Format:%cI$
 describe-name: $Format:%(describe:tags=true,match=*[0-9]*)$
-ref-names: $Format:%D$

{dpdispatcher-0.6.5 → dpdispatcher-0.6.6}/.pre-commit-config.yaml

@@ -18,7 +18,7 @@ repos:
   # Python
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.4.4
+    rev: v0.5.4
     hooks:
       - id: ruff
         args: ["--fix"]
@@ -31,7 +31,7 @@ repos:
         args: ["--write"]
   # Python inside docs
   - repo: https://github.com/asottile/blacken-docs
-    rev: 1.16.0
+    rev: 1.18.0
     hooks:
       - id: blacken-docs
   # markdown, yaml

{dpdispatcher-0.6.5 → dpdispatcher-0.6.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dpdispatcher
-Version: 0.6.5
+Version: 0.6.6
 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
@@ -199,7 +199,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: cloudserver
 Requires-Dist: oss2; extra == "cloudserver"
 Requires-Dist: tqdm; extra == "cloudserver"

{dpdispatcher-0.6.5 → dpdispatcher-0.6.6}/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.6'
+__version_tuple__ = version_tuple = (0, 6, 6)

{dpdispatcher-0.6.5 → dpdispatcher-0.6.6}/dpdispatcher/contexts/ssh_context.py

@@ -295,7 +295,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
@@ -355,10 +359,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",
@@ -603,7 +615,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 +636,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,7 +751,8 @@ 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)
+                file_list,
+                tar_compress=self.remote_profile.get("tar_compress", None),
             )
 
     def block_checkcall(self, cmd, asynchronously=False, stderr_whitelist=None):
@@ -793,18 +809,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 +966,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-0.6.5 → dpdispatcher-0.6.6}/dpdispatcher/machine.py

@@ -261,8 +261,7 @@ 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

{dpdispatcher-0.6.5 → dpdispatcher-0.6.6}/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
@@ -181,10 +184,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 +211,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 +243,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)
@@ -245,8 +267,7 @@ class SGE(PBS):
         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 qstat 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 +280,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 +301,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-0.6.5 → dpdispatcher-0.6.6}/dpdispatcher.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dpdispatcher
-Version: 0.6.5
+Version: 0.6.6
 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
@@ -199,7 +199,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: cloudserver
 Requires-Dist: oss2; extra == "cloudserver"
 Requires-Dist: tqdm; extra == "cloudserver"

{dpdispatcher-0.6.5 → dpdispatcher-0.6.6}/dpdispatcher.egg-info/requires.txt

@@ -27,7 +27,7 @@ sphinx_rtd_theme>=1.0.0rc1
 numpydoc
 deepmodeling_sphinx>=0.1.1
 dargs>=0.3.1
-sphinx-argparse
+sphinx-argparse<0.5.0
 
 [gui]
 dpgui

{dpdispatcher-0.6.5 → dpdispatcher-0.6.6}/pyproject.toml

@@ -55,7 +55,7 @@ docs = [
     'numpydoc',
     'deepmodeling_sphinx>=0.1.1',
     'dargs>=0.3.1',
-    'sphinx-argparse',
+    'sphinx-argparse<0.5.0',
 ]
 cloudserver = ["oss2", "tqdm", "bohrium-sdk"]
 bohrium = ["oss2", "tqdm", "bohrium-sdk"]

{dpdispatcher-0.6.5 → dpdispatcher-0.6.6}/tests/test_lsf_script_generation.py

@@ -108,8 +108,8 @@ class TestLSFScriptGeneration(unittest.TestCase):
             module load use.own
             module load deepmd/1.3
 
-            {{ source /data/home/ypliu/scripts/avail_gpu.sh; }} 
-            {{ source /data/home/ypliu/dprun/tf_envs.sh; }} 
+            source /data/home/ypliu/scripts/avail_gpu.sh
+            source /data/home/ypliu/dprun/tf_envs.sh
 
             export DP_DISPATCHER_EXPORT=test_foo_bar_baz