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

dpdispatcher/contexts/ssh_context.py

@@ -44,6 +44,8 @@ class SSHSession:
         totp_secret=None,
         tar_compress=True,
         look_for_keys=True,
+        execute_command=None,
+        proxy_command=None,
     ):
         self.hostname = hostname
         self.username = username
@@ -56,6 +58,8 @@ class SSHSession:
         self.ssh = None
         self.tar_compress = tar_compress
         self.look_for_keys = look_for_keys
+        self.execute_command = execute_command
+        self.proxy_command = proxy_command
         self._keyboard_interactive_auth = False
         self._setup_ssh()
 
@@ -88,8 +92,7 @@ class SSHSession:
         while not self._check_alive():
             if count == max_check:
                 raise RuntimeError(
-                    "cannot connect ssh after %d failures at interval %d s"
-                    % (max_check, sleep_time)
+                    f"cannot connect ssh after {max_check} failures at interval {sleep_time} s"
                 )
             dlog.info("connection check failed, try to reconnect to " + self.hostname)
             self._setup_ssh()
@@ -141,7 +144,12 @@ class SSHSession:
         # transport = self.ssh.get_transport()
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         sock.settimeout(self.timeout)
-        sock.connect((self.hostname, self.port))
+
+        # Use ProxyCommand if configured (either directly or via jump host parameters)
+        if self.proxy_command is not None:
+            sock = paramiko.ProxyCommand(self.proxy_command)
+        else:
+            sock.connect((self.hostname, self.port))
 
         # Make a Paramiko Transport object using the socket
         ts = paramiko.Transport(sock)
@@ -162,7 +170,6 @@ class SSHSession:
             if os.path.exists(key_path):
                 for pkey_class in (
                     paramiko.RSAKey,
-                    paramiko.DSSKey,
                     paramiko.ECDSAKey,
                     paramiko.Ed25519Key,
                 ):
@@ -180,7 +187,6 @@ class SSHSession:
         elif self.look_for_keys:
             for keytype, name in [
                 (paramiko.RSAKey, "rsa"),
-                (paramiko.DSSKey, "dsa"),
                 (paramiko.ECDSAKey, "ecdsa"),
                 (paramiko.Ed25519Key, "ed25519"),
             ]:
@@ -237,6 +243,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,12 +303,16 @@ class SSHSession:
         assert self.ssh is not None
         try:
             return self.ssh.exec_command(cmd)
-        except (paramiko.ssh_exception.SSHException, socket.timeout) as e:
+        except (
+            paramiko.ssh_exception.SSHException,
+            socket.timeout,
+            EOFError,
+        ) as e:
             # SSH session not active
             # retry for up to 3 times
             # ensure alive
             self.ensure_alive()
-            raise RetrySignal("SSH session not active in calling %s" % cmd) from e
+            raise RetrySignal(f"SSH session not active in calling {cmd}") from e
 
     @property
     def sftp(self):
@@ -334,6 +346,10 @@ 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."
+        doc_proxy_command = (
+            "ProxyCommand to use for SSH connection through intermediate servers."
+        )
         ssh_remote_profile_args = [
             Argument("hostname", str, optional=False, doc=doc_hostname),
             Argument("username", str, optional=False, doc=doc_username),
@@ -355,10 +371,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 +391,20 @@ class SSHSession:
                 default=True,
                 doc=doc_look_for_keys,
             ),
+            Argument(
+                "execute_command",
+                str,
+                optional=True,
+                default=None,
+                doc=doc_execute_command,
+            ),
+            Argument(
+                "proxy_command",
+                [str, type(None)],
+                optional=True,
+                default=None,
+                doc=doc_proxy_command,
+            ),
         ]
         ssh_remote_profile_format = Argument(
             "ssh_session", dict, ssh_remote_profile_args
@@ -375,23 +413,37 @@ class SSHSession:
 
     def put(self, from_f, to_f):
         if self.rsync_available:
+            # For rsync, we need to use %h:%p placeholders for target host/port
+            proxy_cmd_rsync = None
+            if self.proxy_command is not None:
+                proxy_cmd_rsync = self.proxy_command.replace(
+                    f"{self.hostname}:{self.port}", "%h:%p"
+                )
             return rsync(
                 from_f,
                 self.remote + ":" + to_f,
                 port=self.port,
                 key_filename=self.key_filename,
                 timeout=self.timeout,
+                proxy_command=proxy_cmd_rsync,
             )
         return self.sftp.put(from_f, to_f)
 
     def get(self, from_f, to_f):
         if self.rsync_available:
+            # For rsync, we need to use %h:%p placeholders for target host/port
+            proxy_cmd_rsync = None
+            if self.proxy_command is not None:
+                proxy_cmd_rsync = self.proxy_command.replace(
+                    f"{self.hostname}:{self.port}", "%h:%p"
+                )
             return rsync(
                 self.remote + ":" + from_f,
                 to_f,
                 port=self.port,
                 key_filename=self.key_filename,
                 timeout=self.timeout,
+                proxy_command=proxy_cmd_rsync,
             )
         return self.sftp.get(from_f, to_f)
 
@@ -426,7 +478,9 @@ class SSHContext(BaseContext):
         self.init_local_root = local_root
         self.init_remote_root = remote_root
         self.temp_local_root = os.path.abspath(local_root)
-        assert os.path.isabs(remote_root), "remote_root must be a abspath"
+        assert os.path.isabs(os.path.realpath(remote_root)), (
+            "remote_root must be a abspath"
+        )
         self.temp_remote_root = remote_root
         self.remote_profile = remote_profile
         self.remote_root = None
@@ -569,7 +623,7 @@ class SSHContext(BaseContext):
                     rel_file_list, work_path, file_list, directory_list
                 )
             else:
-                raise RuntimeError(f"cannot find upload file {work_path} {jj}")
+                raise FileNotFoundError(f"cannot find upload file {work_path} {jj}")
 
     def upload(
         self,
@@ -603,7 +657,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,15 +678,14 @@ 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
             _, stdout, _ = self.block_checkcall(
-                "sha256sum -c %s --quiet >.sha256sum_stdout 2>/dev/null || :"
-                % shlex.quote(sha256_file)
+                f"sha256sum -c {shlex.quote(sha256_file)} --quiet >.sha256sum_stdout 2>/dev/null || :"
             )
             self.sftp.remove(sha256_file)
             # regenerate file list
@@ -708,7 +764,7 @@ class SSHContext(BaseContext):
                                 os.path.join(
                                     self.local_root,
                                     ii.task_work_path,
-                                    "tag_failure_download_%s" % jj,
+                                    f"tag_failure_download_{jj}",
                                 ),
                                 "w",
                             ) as fp:
@@ -737,49 +793,15 @@ 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 = "nohup %s >/dev/null &" % cmd
-        stdin, stdout, stderr = self.ssh_session.exec_command(
-            ("cd %s ;" % 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
         self.ssh_session.ensure_alive()
         stdin, stdout, stderr = self.ssh_session.exec_command(
-            ("cd %s ;" % shlex.quote(self.remote_root)) + cmd
+            (f"cd {shlex.quote(self.remote_root)} ;") + cmd
         )
         exit_status = stdout.channel.recv_exit_status()
         return exit_status, stdin, stdout, stderr
@@ -794,18 +816,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
@@ -829,8 +856,8 @@ class SSHContext(BaseContext):
         # print(pid)
         return {"stdin": stdin, "stdout": stdout, "stderr": stderr}
 
-    def check_finish(self, cmd_pipes):
-        return cmd_pipes["stdout"].channel.exit_status_ready()
+    def check_finish(self, proc):
+        return proc["stdout"].channel.exit_status_ready()
 
     def get_return(self, cmd_pipes):
         if not self.check_finish(cmd_pipes):
@@ -846,12 +873,12 @@ class SSHContext(BaseContext):
         # Thus, it's better to use system's `rm` to remove a directory, which may
         # save a lot of time.
         if verbose:
-            dlog.info("removing %s" % remotepath)
+            dlog.info(f"removing {remotepath}")
         # In some supercomputers, it's very slow to remove large numbers of files
         # (e.g. directory containing trajectory) due to bad I/O performance.
         # So an asynchronously option is provided.
         self.block_checkcall(
-            "rm -rf %s" % shlex.quote(remotepath),
+            f"rm -rf {shlex.quote(remotepath)}",
             asynchronously=self.clean_asynchronously,
         )
 
@@ -892,11 +919,11 @@ class SSHContext(BaseContext):
         # local tar
         if os.path.isfile(os.path.join(self.local_root, of)):
             os.remove(os.path.join(self.local_root, of))
-        with tarfile.open(
+        with tarfile.open(  # type: ignore[reportCallIssue, reportArgumentType]
             os.path.join(self.local_root, of),
-            tarfile_mode,
-            dereference=dereference,
-            **kwargs,
+            mode=tarfile_mode,  # type: ignore[reportArgumentType]
+            dereference=dereference,  # type: ignore[reportArgumentType]
+            **kwargs,  # type: ignore[reportArgumentType]
         ) as tar:
             # avoid compressing duplicated files or directories
             for ii in set(files):
@@ -921,7 +948,7 @@ class SSHContext(BaseContext):
                 f"from {from_f} to {self.ssh_session.username} @ {self.ssh_session.hostname} : {to_f} Error!"
             )
         # remote extract
-        self.block_checkcall("tar xf %s" % of)
+        self.block_checkcall(f"tar xf {of}")
         # clean up
         os.remove(from_f)
         self.sftp.remove(to_f)
@@ -946,36 +973,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/dlog.py

@@ -6,21 +6,25 @@ import warnings
 dlog = logging.getLogger("dpdispatcher")
 dlog.propagate = False
 dlog.setLevel(logging.INFO)
+cwd_logfile_path = os.path.join(os.getcwd(), "dpdispatcher.log")
+dlogf = logging.FileHandler(cwd_logfile_path, delay=True)
 try:
-    dlogf = logging.FileHandler(
-        os.getcwd() + os.sep + "dpdispatcher" + ".log", delay=True
-    )
+    dlog.addHandler(dlogf)
+    dlog.info(f"LOG INIT:dpdispatcher log direct to {cwd_logfile_path}")
 except PermissionError:
+    dlog.removeHandler(dlogf)
     warnings.warn(
-        "dpdispatcher.log meet permission error. redirect the log to ~/dpdispatcher.log"
+        f"dump logfile dpdispatcher.log to {cwd_logfile_path} meet permission error. redirect the log to ~/dpdispatcher.log"
     )
     dlogf = logging.FileHandler(
         os.path.join(os.path.expanduser("~"), "dpdispatcher.log"), delay=True
     )
+    dlog.addHandler(dlogf)
+    dlog.info("LOG INIT:dpdispatcher log init at ~/dpdispatcher.log")
 
 dlogf_formatter = logging.Formatter("%(asctime)s - %(levelname)s : %(message)s")
 dlogf.setFormatter(dlogf_formatter)
-dlog.addHandler(dlogf)
+# dlog.addHandler(dlogf)
 
 dlog_stdout = logging.StreamHandler(sys.stdout)
 dlog_stdout.setFormatter(dlogf_formatter)

dpdispatcher/dpcloudserver/client.py

@@ -0,0 +1,7 @@
+"""Provide backward compatbility with dflow."""
+
+from dpdispatcher.utils.dpcloudserver.client import RequestInfoException
+
+__all__ = [
+    "RequestInfoException",
+]

dpdispatcher/dpdisp.py

@@ -3,6 +3,7 @@ import argparse
 from typing import List, Optional
 
 from dpdispatcher.entrypoints.gui import start_dpgui
+from dpdispatcher.entrypoints.run import run
 from dpdispatcher.entrypoints.submission import handle_submission
 
 
@@ -54,6 +55,11 @@ def main_parser() -> argparse.ArgumentParser:
         action="store_true",
         help="Clean submission.",
     )
+    parser_submission_action.add_argument(
+        "--reset-fail-count",
+        action="store_true",
+        help="Reset fail count of all jobs to zero.",
+    )
     ##########################################
     # gui
     parser_gui = subparsers.add_parser(
@@ -76,6 +82,18 @@ def main_parser() -> argparse.ArgumentParser:
             "to the network on both IPv4 and IPv6 (where available)."
         ),
     )
+    ##########################################
+    # run
+    parser_run = subparsers.add_parser(
+        "run",
+        help="Run a Python script.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser_run.add_argument(
+        "filename",
+        type=str,
+        help="Python script to run. PEP 723 metadata should be contained in this file.",
+    )
     return parser
 
 
@@ -105,12 +123,15 @@ def main():
             download_terminated_log=args.download_terminated_log,
             download_finished_task=args.download_finished_task,
             clean=args.clean,
+            reset_fail_count=args.reset_fail_count,
         )
     elif args.command == "gui":
         start_dpgui(
             port=args.port,
             bind_all=args.bind_all,
         )
+    elif args.command == "run":
+        run(filename=args.filename)
     elif args.command is None:
         pass
     else:

dpdispatcher/entrypoints/run.py

@@ -0,0 +1,9 @@
+"""Run PEP 723 script."""
+
+from dpdispatcher.run import run_pep723
+
+
+def run(*, filename: str):
+    with open(filename) as f:
+        script = f.read()
+    run_pep723(script)

dpdispatcher/entrypoints/submission.py

@@ -12,6 +12,7 @@ def handle_submission(
     download_terminated_log: bool = False,
     download_finished_task: bool = False,
     clean: bool = False,
+    reset_fail_count: bool = False,
 ):
     """Handle terminated submission.
 
@@ -25,13 +26,21 @@ def handle_submission(
         Download finished tasks.
     clean : bool, optional
         Clean submission.
+    reset_fail_count : bool, optional
+        Reset fail count of all jobs to zero.
 
     Raises
     ------
     ValueError
         At least one action should be specified.
     """
-    if int(download_terminated_log) + int(download_finished_task) + int(clean) == 0:
+    if (
+        int(download_terminated_log)
+        + int(download_finished_task)
+        + int(clean)
+        + int(reset_fail_count)
+        == 0
+    ):
         raise ValueError("At least one action should be specified.")
 
     submission_file = record.get_submission(submission_hash)
@@ -42,7 +51,18 @@ def handle_submission(
     # TODO: for unclear reason, the submission_hash may be changed
     submission.submission_hash = submission_hash
     submission.machine.context.bind_submission(submission)
+    if reset_fail_count:
+        for job in submission.belonging_jobs:
+            job.fail_count = 0
+        # save to remote and local
+        submission.submission_to_json()
+        record.write(submission)
+    if int(download_terminated_log) + int(download_finished_task) + int(clean) == 0:
+        # if only reset_fail_count, no need to update submission state (expensive)
+        return
     submission.update_submission_state()
+    submission.submission_to_json()
+    record.write(submission)
 
     terminated_tasks = []
     finished_tasks = []

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):
@@ -224,7 +227,7 @@ class Machine(metaclass=ABCMeta):
         return if_recover
 
     @abstractmethod
-    def check_finish_tag(self, **kwargs):
+    def check_finish_tag(self, job):
         raise NotImplementedError(
             "abstract method check_finish_tag should be implemented by derived class"
         )
@@ -261,11 +264,19 @@ class Machine(metaclass=ABCMeta):
 
         source_list = job.resources.source_list
         for ii in source_list:
-            line = "{ source %s; } \n" % ii
-            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:
@@ -466,7 +477,7 @@ class Machine(metaclass=ABCMeta):
         job : Job
             job
         """
-        dlog.warning("Job %s should be manually killed" % job.job_id)
+        dlog.warning(f"Job {job.job_id} should be manually killed")
 
     def get_exit_code(self, job):
         """Get exit code of the job.