Source code for romtools.hpc.dispatchers.remote_dispatcher

import posixpath as ppath
import shlex
from typing import Optional

from romtools.hpc.logger import Logger
from romtools.hpc.connection import Connection, Result
from romtools.hpc.dispatchers.base_dispatcher import BaseDispatcher
from romtools.hpc.components.caller import RemoteCaller
from romtools.hpc.components.file_manager import RemoteFileManager
from romtools.hpc.components.slurm_job_manager import SlurmJobManager
from romtools.hpc.components.transfer_manager import RemoteTransferManager


def _normalize_campaign_directory(campaign_directory: str) -> str:
    """The campaign directory as a clean relative path, mirrored on both hosts."""
    cleaned = ppath.normpath((campaign_directory or "").strip())

    if not cleaned or cleaned == ".":
        raise ValueError("The campaign directory must name a directory.")

    if ppath.isabs(cleaned) or cleaned == ".." or cleaned.startswith("../"):
        raise ValueError(
            f"The campaign directory must be relative to the remote root (received: "
            f"{campaign_directory}). This workflow also creates the same directory "
            "on the local machine."
        )

    return cleaned




class RemoteDispatcher(BaseDispatcher):
    """
    Runs ROM workflows on a remote host over SSH, connecting as user@remote:port.

    Coordinates composed helpers rather than doing the work itself: files go through
    RemoteFileManager, Python calls through RemoteCaller, batch jobs through
    SlurmJobManager, and archives through RemoteTransferManager. Paths are taken relative
    to the configured remote root.

    Arguments:
        logger: An instance of the Logger class for logging
        campaign_directory: An optional string naming the directory this campaign
            runs in, relative to the remote root. It is mirrored locally and
            under the remote root.
        config: A YAML path or a Configuration. Defaults to reading the command line.
    """
    def __init__(self, campaign_directory: str = "hpctools", logger: Logger = None,
                 connection: Optional[Connection] = None, config=None):
        # Normalized before super() so every component is built with the same value
        super().__init__(_normalize_campaign_directory(campaign_directory), logger, config=config)

        # Core members
        self.conn : Optional[Connection] = None

        # Confirm that connection is possible
        if not self.config.get("remote") or not self.config.get("user"):
            raise ValueError("Remote host and user must be specified in the configuration to use RemoteDispatcher.")

        # Establish connection, or use the one provided (e.g. by tests)
        if connection is not None:
            self.conn = connection
            self.logger.set_hostname(self.conn.host)
        else:
            self._connect_to_remote()

        # Composed sub-classes
        self.files = RemoteFileManager(
            connection=self.conn,
            config=self.config,
            logger=self.logger)

        self.caller = RemoteCaller(
            connection=self.conn,
            config=self.config,
            logger=self.logger,
            files=self.files)

        self.slurm = SlurmJobManager(
            run_cmd=self.conn.run,
            config=self.config,
            logger=self.logger,
            campaign_directory=self.campaign_directory,
            files=self.files)

        self.transfer = RemoteTransferManager(
            connection=self.conn,
            config=self.config,
            logger=self.logger,
            campaign_directory=self.campaign_directory,
            files=self.files)

    # ------------------------------------------------------------------
    # Initialization and setup
    # ------------------------------------------------------------------

    def _connect_to_remote(self) -> None:
        """Establish the SSH connection described by the configuration."""
        try:
            self.conn = Connection(host=self.config.get("remote"), user=self.config.get("user"), port=self.config.get("port"))
        except Exception as e:
            raise RuntimeError(f"Failed to establish SSH connection: {e}") from e
        self.logger.set_hostname(self.conn.host)
        self.logger.log(f"Connection established with {self.conn.host}.", local=True)

    # ------------------------------------------------------------------
    # Resource management
    # ------------------------------------------------------------------

    def close(self):
        self.conn.close()
        self.logger.log("Connection closed.", local=True)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def require_relative_path(self, path: str) -> None:
        if ppath.isabs(path):
            raise ValueError(
                f"You must provide a path relative to the remote root (received: {path}). "
                "This workflow also creates the same directory on the local machine."
            )

    def require_supported_concurrency(self, concurrency: int) -> None:
        if concurrency != 1:
            raise ValueError(
                f"Concurrency > 1 is not supported with a RemoteDispatcher (received: {concurrency}). "
                "Use a concurrency of 1 and let SLURM provide the parallelism."
            )



    def run(self, cmd: str, run_directory: str = None) -> Result:
        """
        Run a command directly on the remote host, without SLURM.

        This lands wherever the SSH session does, which on a cluster is the login
        node, so it suits quick staging or preprocessing; anything long or
        parallel belongs in submit_job(). run_directory is relative to the remote
        root, and defaults to the root itself.

        Returns a Result object (with stdout, stderr, exit_code, ok)
        """
        remote_cmd = f"cd {shlex.quote(self.files.resolve_path(run_directory))} && {cmd}"
        res = self.conn.run(remote_cmd)
        if res.ok:
            self.logger.debug(f"Executed command on remote host: {cmd}")
        else:
            self.logger.log(f"Command failed ({cmd}), exit code {res.exit_code}: {res.stderr}")
        return res