feat: transfer data with mscp

src/anemoi/utils/remote/__init__.py

@@ -614,9 +614,9 @@ def _find_transfer_class(source: str, target: str) -> type:
     assert sum([from_ssh, from_local, from_s3]) == 1, (from_ssh, from_local, from_s3)
 
     if from_local and into_ssh:  # local -> ssh
-        from .ssh import RsyncUpload
+        from .ssh import SshUpload
 
-        return RsyncUpload
+        return SshUpload
 
     if from_s3 and into_local:  # local <- S3
         from .s3 import S3Download

src/anemoi/utils/remote/ssh.py

@@ -17,6 +17,12 @@
 LOGGER = logging.getLogger(__name__)
 
 
+def _isProgramOnPath(program_name):
+    import shutil
+
+    return shutil.which(program_name) is not None
+
+
 def call_process(*args: str) -> str:
     """Execute a subprocess with the given arguments and return its output.
 
@@ -133,6 +139,95 @@ def delete_target(self, target: str) -> None:
         # call_process("ssh", hostname, "rm", "-rf", shlex.quote(path))
 
 
+class MscpUpload(SshBaseUpload):
+
+    def copy(self, source: str, target: str, **kwargs) -> None:
+        """Copy a file or a folder from the source to the target location.
+
+        Parameters
+        ----------
+        source : str
+            The source location.
+        target : str
+            The target location.
+        kwargs : dict
+            Additional arguments for the transfer.
+        """
+        self.transfer_file(source=source, target=target, **kwargs)
+
+    def _strip_zarr_path(self, path: str) -> str:
+        """Removes last level from path string if it ends in ".zarr/" or ".zarr"
+
+        Needed for mscp because i always creates a dir called dataset_name.zarr under the given dest
+        so if the dest is given as "/path/to/dataset_name.zarr" you would end up with
+        "/path/to/dataset_name.zarr/dataset_name.zarr".
+
+        """
+        path = path.rstrip("/")  # Remove trailing slash if any
+        dirname, basename = os.path.split(path)
+
+        if basename.endswith(".zarr"):
+            return os.path.join(dirname, "")
+        else:
+            return path + "/" if not path.endswith("/") else path
+
+    def _transfer_file(
+        self, source: str, target: str, overwrite: bool, resume: bool, verbosity: int, threads: int, config: dict = None
+    ) -> int:
+        """Transfer a file using mscp.
+
+        Parameters
+        ----------
+        source : str
+            The source file path.
+        target : str
+            The target file path.
+        overwrite : bool
+            Whether to overwrite the target if it exists.
+        resume : bool
+            Whether to resume the transfer if possible.
+        verbosity : int
+            The verbosity level.
+        threads : int
+            The number of threads to use.
+        config : dict, optional
+            Additional configuration options.
+
+        Returns
+        -------
+        int
+            The size of the transferred file.
+        """
+        hostname, path = self._parse_target(target)
+
+        size = os.path.getsize(source)
+
+        # remove dataset name from dest path if its included
+        # mscp will add it regardless, so we dont want it twice
+        target = self._strip_zarr_path(target)
+        LOGGER.debug(f"Copying {source} to {target} with Mscp")
+
+        if verbosity > 0:
+            LOGGER.info(f"{self.action} {source} to {target} ({bytes_to_human(size)})")
+
+        call_process("ssh", hostname, "mkdir", "-p", shlex.quote(os.path.dirname(path)))
+        if threads > 1:
+            call_process(
+                "mscp",
+                "-n",
+                str(threads),
+                source,
+                f"{hostname}:{path}",
+            )
+        else:  # if threads not specified, use the default number of cores from mscp "floor(log(#cores)) + 1"
+            call_process(
+                "mscp",
+                source,
+                f"{hostname}:{path}",
+            )
+        return size
+
+
 class RsyncUpload(SshBaseUpload):
 
     def _transfer_file(
@@ -249,6 +344,29 @@ def _transfer_file(
         return size
 
 
+def _pickTransferTool():
+    tools = {"mscp": MscpUpload, "rsync": RsyncUpload, "scp": ScpUpload}
+
+    from anemoi.utils.config import load_config
+
+    tool = load_config().get("utils", {}).get("transfer_tool", None)
+    if tool is not None:
+        # check if the tool listed in the config can be found
+        if tool in tools and _isProgramOnPath(tool):
+            LOGGER.info(f"Using {tool} to transfer as specified in the anemoi utils config")
+            return tools[tool]
+
+    # Loops through this list in order until it finds a tool
+    for tool in tools:
+        if _isProgramOnPath(tool):
+            LOGGER.info(f"Using {tool} to transfer")
+            return tools[tool]
+    raise RuntimeError(f"No suitable transfer tool found. Looked for the following: {tools}")
+
+
+SshUpload = _pickTransferTool()
+
+
 def upload(source: str, target: str, **kwargs) -> None:
     """Upload a file or folder to the target location using rsync.
 
@@ -261,9 +379,6 @@ def upload(source: str, target: str, **kwargs) -> None:
     kwargs : dict
         Additional arguments for the transfer.
     """
-    uploader = RsyncUpload()
+    uploader = SshUpload()
 
-    if os.path.isdir(source):
-        uploader.transfer_folder(source=source, target=target, **kwargs)
-    else:
-        uploader.transfer_file(source=source, target=target, **kwargs)
+    uploader.transfer_file(source=source, target=target, **kwargs)