improve _get_files in SSH context (#473)

Author: thangckt

dpdispatcher/contexts/ssh_context.py

@@ -295,7 +295,11 @@ def exec_command(self, cmd):
         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 @@ def arginfo():
             ),
             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 @@ def upload(
                 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
@@ -736,7 +751,8 @@ def download(
         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 @@ def write_file(self, fname, write_str):
         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 @@ def _get_files(self, files, tar_compress=True):
         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/machines/pbs.py

@@ -215,12 +215,13 @@ def gen_script_header(self, job):
         # resources.number_node is not used in SGE
         resources = job.resources
         job_name = resources.kwargs.get("job_name", "wDPjob")
-        sge_pe_name = resources.kwargs.get("sge_pe_name", "mpi")
+        pe_name = resources.kwargs.get("pe_name", "mpi")
         sge_script_header_dict = {}
         sge_script_header_dict["select_node_line"] = f"#$ -N {job_name}\n"
         sge_script_header_dict["select_node_line"] += (
-            f"#$ -pe {sge_pe_name} {resources.cpu_per_node}\n"
+            f"#$ -pe {pe_name} {resources.cpu_per_node}\n"
         )
+
         if resources.queue_name != "":
             sge_script_header_dict["select_node_line"] += (
                 f"#$ -q {resources.queue_name}"
@@ -266,8 +267,7 @@ def check_status(self, job):
         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:
@@ -280,8 +280,7 @@ def check_status(self, job):
                 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
@@ -307,15 +306,15 @@ def check_finish_tag(self, job):
     def resources_subfields(cls) -> List[Argument]:
         """Generate the resources subfields.
 
-            sge_pe_name : str
+            pe_name : str
         The parallel environment name of SGE.
 
         Returns
         -------
         list[Argument]
             resources subfields
         """
-        doc_sge_pe_name = "The parallel environment name of SGE."
+        doc_pe_name = "The parallel environment name of SGE system."
         doc_job_name = "The name of SGE's job."
 
         return [
@@ -324,11 +323,12 @@ def resources_subfields(cls) -> List[Argument]:
                 dict,
                 [
                     Argument(
-                        "sge_pe_name",
+                        "pe_name",
                         str,
                         optional=True,
                         default="mpi",
-                        doc=doc_sge_pe_name,
+                        doc=doc_pe_name,
+                        alias=["sge_pe_name"],
                     ),
                     Argument(
                         "job_name",