diff --git a/docs/guides/executors/slurm.md b/docs/guides/executors/slurm.md index c49d3c20..a23b7e09 100644 --- a/docs/guides/executors/slurm.md +++ b/docs/guides/executors/slurm.md @@ -6,6 +6,7 @@ Launch tasks on a Slurm HPC cluster, optionally from your local machine over SSH - Access to a Slurm cluster with Pyxis installed - SSH key authentication set up (for remote launch via `SSHTunnel`) +- OpenSSH `ssh` and `scp` executables when using persistent connection multiplexing - A container image accessible from the cluster (e.g. on a shared registry or pulled to the nodes) ## Executor configuration @@ -20,6 +21,11 @@ ssh_tunnel = run.SSHTunnel( user="your-username", job_dir="/scratch/your-username/nemo-runs", # where NeMo-Run stores metadata on the cluster identity="~/.ssh/id_ed25519", # optional SSH key path + # Opt into OpenSSH and allow NeMo Run to create a persistent master. + use_openssh=True, + control_persist="10m", + # Optional override; creation mode defaults to ~/.nemo_run/.ssh/control-%C. + # control_path="~/.ssh/nemo-run-%C", ) executor = run.SlurmExecutor( @@ -38,6 +44,44 @@ executor = run.SlurmExecutor( Use `run.LocalTunnel()` instead of `SSHTunnel` when launching from a login node directly. +### Persistent SSH multiplexing + +`use_openssh` is configured on the `SSHTunnel` passed to `SlurmExecutor`. In creation mode, +`control_persist` accepts an OpenSSH duration such as `"10m"`; NeMo Run reuses a compatible master +or starts one with that lifetime. The master is a separate process, so later NeMo Run invocations +can reuse it until it has had no clients for the configured duration. + +For MFA-protected hosts, create the authenticated master yourself and prevent NeMo Run from opening +a new connection that could prompt unexpectedly: + +```python +ssh_tunnel = run.SSHTunnel( + host="login-ptyche", + user="your-username", + job_dir="/scratch/your-username/nemo-runs", + use_openssh=True, + require_existing_master=True, +) +``` + +This mode runs `ssh -O check` and only reuses an existing master. It never creates one. Configure +`ControlMaster`, `ControlPath`, and `ControlPersist` in `~/.ssh/config`, then start the master (for +example, `ssh -fN login-ptyche`) before launching NeMo Run. If no master exists, NeMo Run fails with +an actionable startup command. + +Existing-master mode reads `ControlPath`, `ControlPersist`, and other connection settings from +OpenSSH configuration unless explicitly overridden. Creation mode defaults to the stable +`~/.nemo_run/.ssh/control-%C` path when `control_path` is omitted. Keep an override stable, +include a token such as `%C` for multiple destinations, and place it in a directory owned by the +current user that is not group/world-writable and does not traverse symlinks. + +This mode requires working `ssh` and `scp` executables and key-based, agent-based, or otherwise +non-interactive OpenSSH authentication. Configuration errors and OpenSSH connection failures are +reported directly; NeMo Run does **not** dynamically fall back to Paramiko after multiplexing is +selected. Omit `use_openssh`, `control_persist`, and `require_existing_master` to retain the existing +in-process Fabric/Paramiko behavior. For backward compatibility, setting `control_persist` also +selects OpenSSH creation mode. + Key parameters: | Parameter | Description | @@ -50,6 +94,10 @@ Key parameters: | `container_image` | Container image URI | | `time` | Wall-time limit (`"HH:MM:SS"`) | | `tunnel` | `SSHTunnel` (remote) or `LocalTunnel` (on-cluster) | +| `use_openssh` | Select the OpenSSH backend instead of Fabric/Paramiko | +| `require_existing_master` | Reuse a pre-authenticated master and never create a connection | +| `control_persist` | Lifetime applied only when NeMo Run creates a master | +| `control_path` | Optional `ControlPath` override | | `packager` | Code packaging strategy | ## E2E workflow diff --git a/nemo_run/core/tunnel/client.py b/nemo_run/core/tunnel/client.py index 750defa4..35f93c0c 100644 --- a/nemo_run/core/tunnel/client.py +++ b/nemo_run/core/tunnel/client.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024-2025, NVIDIA CORPORATION. +# Copyright (c) 2024-2026, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -15,6 +15,7 @@ import getpass import logging import os +import shlex import shutil import socket import subprocess @@ -166,6 +167,178 @@ def cleanup(self): self.session.clear() +class _OpenSSHSession: + """Fabric-compatible subset backed by a persistent OpenSSH control master.""" + + def __init__( + self, + *, + host: str, + user: str, + port: Optional[int], + identity: Optional[str], + control_persist: Optional[str], + control_path: Optional[str], + require_existing_master: bool, + ): + self.host = host + self.user = user + self.port = port + self.connect_kwargs = {"key_filename": [identity]} if identity else {} + self.control_persist = control_persist + self.control_path = os.path.expanduser(control_path) if control_path else None + self.require_existing_master = require_existing_master + self._context = Context() + + @property + def _target(self) -> str: + return f"{self.user}@{self.host}" + + @property + def _scp_target(self) -> str: + host = f"[{self.host}]" if self.host.count(":") > 1 else self.host + return f"{self.user}@{host}" + + @property + def _control_options(self) -> list[str]: + options: list[str] = [] + if self.control_persist: + options.extend( + ["-o", "ControlMaster=auto", "-o", f"ControlPersist={self.control_persist}"] + ) + if self.control_path: + options.extend(["-o", f"ControlPath={self.control_path}"]) + if self.require_existing_master: + # A vanished master must fail instead of opening a direct connection or prompting. + options.extend( + [ + "-o", + "ControlMaster=no", + "-o", + "BatchMode=yes", + "-o", + "ProxyCommand=false", + ] + ) + return options + + def _connection_options(self, executable: str) -> list[str]: + options = [*self._control_options] + if self.port is not None: + port_flag = "-P" if executable == "scp" else "-p" + options.extend([port_flag, str(self.port)]) + if self.connect_kwargs: + options.extend(["-i", self.connect_kwargs["key_filename"][0]]) + return options + + @property + def ssh_options(self) -> str: + """Options which let rsync reuse this session's control master.""" + return shlex.join(self._control_options) + + @property + def is_connected(self) -> bool: + result = self._context.run( + self._command("ssh", "-O", "check", self._target), hide=True, warn=True + ) + return result.ok + + def _command(self, executable: str, *args: str) -> str: + return shlex.join([executable, *self._connection_options(executable), *args]) + + def _master_start_command(self) -> str: + options: list[str] = [] + if self.control_path: + options.extend(["-o", f"ControlPath={self.control_path}"]) + if self.port is not None: + options.extend(["-p", str(self.port)]) + if self.connect_kwargs: + options.extend(["-i", self.connect_kwargs["key_filename"][0]]) + return shlex.join(["ssh", *options, "-o", "ControlMaster=yes", "-fN", self._target]) + + def _prepare_control_directory(self) -> None: + assert self.control_path + parent = Path(self.control_path).parent + parent.mkdir(mode=0o700, parents=True, exist_ok=True) + resolved_parent = parent.resolve(strict=True) + if parent.absolute() != resolved_parent: + raise RuntimeError( + f"OpenSSH control socket directory must not contain symlinks: {parent}" + ) + directory_mode = resolved_parent.stat() + if directory_mode.st_uid != os.getuid() or directory_mode.st_mode & 0o022: + raise RuntimeError( + f"OpenSSH control socket directory must be owned by the current user and not " + f"group/world-writable: {resolved_parent}" + ) + + def open(self) -> None: + if self.is_connected: + return + if self.require_existing_master: + raise RuntimeError( + f"No existing OpenSSH control master found for {self._target}. " + f"Start one first (for example, `{self._master_start_command()}`) and retry." + ) + assert self.control_persist + if self.control_path: + self._prepare_control_directory() + self._context.run(self._command("ssh", "-fN", self._target), hide=False) + + def run(self, command: str, hide: bool = True, warn: bool = False, **kwargs) -> RunResult: + self.open() + return self._context.run( + self._command("ssh", self._target, command), hide=hide, warn=warn, **kwargs + ) + + def local(self, command: str, hide: bool = True, **kwargs) -> RunResult: + return self._context.run(command, hide=hide, **kwargs) + + def put(self, local_path: str, remote_path: str) -> None: + self.open() + self._context.run( + self._command("scp", local_path, f"{self._scp_target}:{remote_path}"), hide=True + ) + + def get(self, remote_path: str, local_path: str) -> None: + self.open() + self._context.run( + self._command("scp", f"{self._scp_target}:{remote_path}", local_path), hide=True + ) + + def forward_local( + self, + local_port: int, + remote_port: Optional[int] = None, + remote_host: str = "localhost", + local_host: str = "localhost", + ): + self.open() + forward = f"{local_host}:{local_port}:{remote_host}:{remote_port or local_port}" + start = self._command("ssh", "-O", "forward", "-L", forward, self._target) + stop = self._command("ssh", "-O", "cancel", "-L", forward, self._target) + + class ForwardContext: + def __enter__(self): + self._session.open() + self._session._context.run(start, hide=True) + return self + + def __exit__(self, *_): + self._session.open() + self._session._context.run(stop, hide=True, warn=True) + + def __init__(self, session): + self._session = session + + return ForwardContext(self) + + def close(self) -> None: + # The control master intentionally outlives this Python process. OpenSSH exits it after + # ControlPersist has elapsed without clients; later operations probe the socket first. + pass + + @dataclass(kw_only=True) class SSHTunnel(Tunnel): """ @@ -173,6 +346,11 @@ class SSHTunnel(Tunnel): Currently only supports SlurmExecutor. Uses key based authentication if *identity* is provided else password authentication. + Set *use_openssh* to multiplex commands and transfers through an OpenSSH control master. + Set *require_existing_master* to reuse a master configured and started outside NeMo Run without + ever creating a connection. Otherwise, *control_persist* specifies the lifetime of a master + which NeMo Run may create. Without *use_openssh* or *control_persist*, the existing in-process + Fabric/Paramiko connection is used. Examples -------- @@ -188,7 +366,9 @@ class SSHTunnel(Tunnel): host=os.environ["ANOTHER_SSH_HOST"], user=os.environ["ANOTHER_SSH_USER"], job_dir=os.environ["ANOTHER_REMOTE_JOBDIR"], - identity="path_to_private_key" + identity="path_to_private_key", + use_openssh=True, + control_persist="10m", ) """ @@ -199,8 +379,28 @@ class SSHTunnel(Tunnel): identity: Optional[str] = None shell: Optional[str] = None pre_command: Optional[str] = None + use_openssh: bool = False + require_existing_master: bool = False + control_persist: Optional[str] = None + control_path: Optional[str] = None def __post_init__(self): + if self.control_persist: + self.use_openssh = True + if self.require_existing_master and not self.use_openssh: + raise ValueError("require_existing_master requires use_openssh") + if self.require_existing_master and self.control_persist: + raise ValueError("require_existing_master cannot be combined with control_persist") + if self.use_openssh and not self.require_existing_master and not self.control_persist: + raise ValueError("OpenSSH master creation requires control_persist") + if self.control_path and not self.use_openssh: + raise ValueError("control_path requires use_openssh") + if self.control_persist == "": + raise ValueError("control_persist must not be empty") + if self.use_openssh and not shutil.which("ssh"): + raise RuntimeError("OpenSSH multiplexing requires the ssh executable") + if self.use_openssh and not shutil.which("scp"): + raise RuntimeError("OpenSSH multiplexing requires the scp executable") self.console = CONSOLE self.session = None self.auth_handler: Callable = authentication_handler @@ -224,8 +424,23 @@ def _create_job_dir(self, tunnel: Tunnel): tunnel.run(command) def connect(self): + if self.use_openssh and not self.session: + self.session = _OpenSSHSession( + host=self.host, + user=self.user, + port=self.port, + identity=self.identity, + control_persist=self.control_persist, + control_path=self.control_path + or ( + None + if self.require_existing_master + else os.path.join(get_nemorun_home(), ".ssh", "control-%C") + ), + require_existing_master=self.require_existing_master, + ) if not (self.session and self.session.is_connected): - self._authenticate() + self.session.open() if self.use_openssh else self._authenticate() def _check_connect(self): if not (self.session and self.session.is_connected): diff --git a/nemo_run/core/tunnel/rsync.py b/nemo_run/core/tunnel/rsync.py index 22d0e86c..360483d8 100644 --- a/nemo_run/core/tunnel/rsync.py +++ b/nemo_run/core/tunnel/rsync.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,6 +14,7 @@ # limitations under the License. import logging +import shlex from typing import Iterable from fabric import Connection @@ -38,56 +39,44 @@ def rsync( # Turn single-string exclude into a one-item list for consistency if isinstance(exclude, str): exclude = [exclude] - # Create --exclude options from exclude list - exclude_opts = ' --exclude "{}"' * len(exclude) - # Double-backslash-escape - exclusions = tuple([str(s).replace('"', '\\\\"') for s in exclude]) - # Honor SSH key(s) - key_string = "" - # TODO: seems plausible we need to look in multiple places if there's too - # much deferred evaluation going on in how we eg source SSH config files - # and so forth, re: connect_kwargs - # TODO: we could get VERY fancy here by eg generating a tempfile from any - # in-memory-only keys...but that's also arguably a security risk, so... + + ssh_command = ["ssh"] keys = c.connect_kwargs.get("key_filename", []) - # TODO: would definitely be nice for Connection/FabricConfig to expose an - # always-a-list, always-up-to-date-from-all-sources attribute to save us - # from having to do this sort of thing. (may want to wait for Paramiko auth - # overhaul tho!) if isinstance(keys, str): keys = [keys] - if keys: - key_string = "-i " + " -i ".join(keys) - # Get base cxn params + for key in keys: + ssh_command.extend(["-i", key]) + user, host, port = c.user, c.host, c.port - port_string = "-p {}".format(port) - # Remote shell (SSH) options - rsh_string = "" - # Strict host key checking - disable_keys = "-o StrictHostKeyChecking=no" - if not strict_host_keys and disable_keys not in ssh_opts: - ssh_opts += " {}".format(disable_keys) - rsh_parts = [key_string, port_string, ssh_opts] - if any(rsh_parts): - rsh_string = "--rsh='ssh {}'".format(" ".join(rsh_parts)) - # Set up options part of string - options_map = { - "delete": "--delete" if delete else "", - "exclude": exclude_opts.format(*exclusions), - "rsh": rsh_string, - "extra": rsync_opts, - } - options = "{delete}{exclude} -pthrvz {extra} {rsh}".format(**options_map) - # Create and run final command string - # TODO: richer host object exposing stuff like .address_is_ipv6 or whatever + if port is not None: + ssh_command.extend(["-p", str(port)]) + + session_ssh_opts = getattr(c, "ssh_options", "") + if not isinstance(session_ssh_opts, str): + session_ssh_opts = "" + ssh_option_tokens = shlex.split(session_ssh_opts) + shlex.split(ssh_opts) + if not strict_host_keys and "StrictHostKeyChecking=no" not in ssh_option_tokens: + ssh_option_tokens.extend(["-o", "StrictHostKeyChecking=no"]) + ssh_command.extend(ssh_option_tokens) + + command = ["rsync"] + if delete: + command.append("--delete") + for pattern in exclude: + command.extend(["--exclude", str(pattern)]) + command.append("-pthrvz") + command.extend(shlex.split(rsync_opts)) + if len(ssh_command) > 1: + command.extend(["--rsh", shlex.join(ssh_command)]) + if host.count(":") > 1: - # Square brackets are mandatory for IPv6 rsync address, - # even if port number is not specified - cmd = "rsync {} {} [{}@{}]:{}" + destination = f"[{user}@{host}]:{target}" else: - cmd = "rsync {} {} {}@{}:{}" - cmd = cmd.format(options, source, user, host, target) - c.run(f"mkdir -p {target}", hide=hide_output) + destination = f"{user}@{host}:{target}" + command.extend([source, destination]) + cmd = shlex.join(command) + + c.run(f"mkdir -p {shlex.quote(target)}", hide=hide_output) result = c.local(cmd, hide=hide_output) if result: logger.info(f"Successfully ran `{result.command}`") diff --git a/test/core/tunnel/test_client.py b/test/core/tunnel/test_client.py index e2f5f50e..42c4db92 100644 --- a/test/core/tunnel/test_client.py +++ b/test/core/tunnel/test_client.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024-2025, NVIDIA CORPORATION. +# Copyright (c) 2024-2026, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,7 +13,7 @@ # limitations under the License. from pathlib import Path -from unittest.mock import MagicMock, call, mock_open, patch +from unittest.mock import MagicMock, PropertyMock, call, mock_open, patch import pytest @@ -23,6 +23,7 @@ PackagingJob, SSHConfigFile, SSHTunnel, + _OpenSSHSession, authentication_handler, delete_tunnel_dir, ) @@ -135,6 +136,8 @@ def test_init(self, ssh_tunnel): assert ssh_tunnel.user == "test_user" assert ssh_tunnel.job_dir == "/remote/job" assert ssh_tunnel.identity is None + assert ssh_tunnel.control_persist is None + assert ssh_tunnel.control_path is None assert ssh_tunnel.session is None def test_set_job_dir(self, ssh_tunnel): @@ -252,6 +255,318 @@ def test_setup(self, ssh_tunnel): mock_run.assert_called_once_with(f"mkdir -p {ssh_tunnel.job_dir}") +class TestOpenSSHSession: + @pytest.fixture + def context(self): + with patch("nemo_run.core.tunnel.client.Context") as mock_context: + yield mock_context.return_value + + @pytest.fixture + def session(self, context, tmp_path): + return _OpenSSHSession( + host="test.host", + user="test_user", + port=2222, + identity="/path/to/key", + control_persist="10m", + control_path=str(tmp_path / "control-%C"), + require_existing_master=False, + ) + + def test_connect_starts_control_master(self, session, context): + context.run.return_value.ok = False + + session.open() + + assert context.run.call_count == 2 + check_command = context.run.call_args_list[0].args[0] + open_command = context.run.call_args_list[1].args[0] + assert "ControlMaster=auto" in check_command + assert "ControlPersist=10m" in check_command + assert "ControlPath=" in check_command + assert "-p 2222" in open_command + assert "-i /path/to/key" in open_command + assert "-fN test_user@test.host" in open_command + + def test_new_session_reuses_existing_control_master(self, session, context): + context.run.return_value.ok = True + second_session = _OpenSSHSession( + host=session.host, + user=session.user, + port=session.port, + identity=session.connect_kwargs["key_filename"][0], + control_persist=session.control_persist, + control_path=session.control_path, + require_existing_master=False, + ) + + second_session.open() + + context.run.assert_called_once() + + def test_run_uses_control_master(self, session, context): + context.run.side_effect = [MagicMock(ok=True), MagicMock()] + + session.run("squeue", hide=False, warn=True) + + check_command = context.run.call_args_list[0].args[0] + command = context.run.call_args_list[1].args[0] + assert "-O check" in check_command + assert command.endswith("test_user@test.host squeue") + context.run.assert_any_call(command, hide=False, warn=True) + + def test_put_and_get_use_control_master(self, session, context): + context.run.return_value.ok = True + + session.put("local file", "/remote/file") + put_command = context.run.call_args.args[0] + assert put_command.startswith("scp ") + assert "ControlPath=" in put_command + assert "-P 2222" in put_command + assert "'local file' test_user@test.host:/remote/file" in put_command + + context.reset_mock() + session.get("/remote/file", "local file") + get_command = context.run.call_args.args[0] + assert get_command.startswith("scp ") + assert "test_user@test.host:/remote/file 'local file'" in get_command + + def test_forward_local_adds_and_cancels_forward(self, session, context): + context.run.return_value.ok = True + + with session.forward_local(7000, remote_host="compute"): + pass + + commands = [call.args[0] for call in context.run.call_args_list] + assert sum("-O check" in command for command in commands) == 3 + assert any("-O forward -L localhost:7000:compute:7000" in command for command in commands) + assert any("-O cancel -L localhost:7000:compute:7000" in command for command in commands) + + def test_expired_master_is_restarted_between_operations(self, session, context): + context.run.side_effect = [ + MagicMock(ok=True), + MagicMock(), + MagicMock(ok=False), + MagicMock(), + MagicMock(), + ] + + session.run("squeue") + session.run("sacct") + + assert "-O check" in context.run.call_args_list[0].args[0] + assert "-O check" in context.run.call_args_list[2].args[0] + assert "-fN test_user@test.host" in context.run.call_args_list[3].args[0] + assert context.run.call_args_list[4].args[0].endswith("test_user@test.host sacct") + + def test_ipv6_put_and_get_bracket_host(self, context, tmp_path): + session = _OpenSSHSession( + host="2001:db8::1", + user="test_user", + port=22, + identity=None, + control_persist="10m", + control_path=str(tmp_path / "control-%C"), + require_existing_master=False, + ) + context.run.return_value.ok = True + + session.put("local", "/remote/file") + assert "test_user@[2001:db8::1]:/remote/file" in context.run.call_args.args[0] + + context.reset_mock() + session.get("/remote/file", "local") + assert "test_user@[2001:db8::1]:/remote/file" in context.run.call_args.args[0] + + def test_creation_rejects_unsafe_control_directory(self, context, tmp_path): + unsafe_directory = tmp_path / "unsafe" + unsafe_directory.mkdir(mode=0o777) + unsafe_directory.chmod(0o777) + session = _OpenSSHSession( + host="test.host", + user="test_user", + port=None, + identity=None, + control_persist="10m", + control_path=str(unsafe_directory / "control-%C"), + require_existing_master=False, + ) + context.run.return_value.ok = False + + with pytest.raises(RuntimeError, match="not group/world-writable"): + session.open() + + def test_creation_rejects_symlinked_control_directory(self, context, tmp_path): + safe_directory = tmp_path / "safe" + safe_directory.mkdir(mode=0o700) + symlink = tmp_path / "linked" + symlink.symlink_to(safe_directory, target_is_directory=True) + session = _OpenSSHSession( + host="test.host", + user="test_user", + port=None, + identity=None, + control_persist="10m", + control_path=str(symlink / "control-%C"), + require_existing_master=False, + ) + context.run.return_value.ok = False + + with pytest.raises(RuntimeError, match="must not contain symlinks"): + session.open() + + def test_close_leaves_control_master_running(self, session, context): + session.close() + + context.run.assert_not_called() + + @patch("nemo_run.core.tunnel.client.shutil.which", return_value="/usr/bin/ssh") + def test_tunnel_connect_uses_openssh_session(self, _): + tunnel = SSHTunnel( + host="test.host", + user="test_user", + job_dir="/remote/job", + control_persist="10m", + ) + + with patch.object(_OpenSSHSession, "is_connected", new_callable=PropertyMock) as connected: + connected.return_value = False + with patch.object(_OpenSSHSession, "open") as open_session: + tunnel.connect() + + assert isinstance(tunnel.session, _OpenSSHSession) + assert tunnel.session.control_path.endswith("/.ssh/control-%C") + open_session.assert_called_once() + + @patch("nemo_run.core.tunnel.client.shutil.which", return_value="/usr/bin/ssh") + def test_tunnel_reuses_master_from_ssh_config(self, _): + tunnel = SSHTunnel( + host="login-ptyche", + user="test_user", + job_dir="/remote/job", + use_openssh=True, + require_existing_master=True, + ) + + with patch.object(_OpenSSHSession, "is_connected", new_callable=PropertyMock) as connected: + connected.return_value = True + with patch.object(_OpenSSHSession, "open") as open_session: + tunnel.connect() + + assert isinstance(tunnel.session, _OpenSSHSession) + assert tunnel.session.control_path is None + assert tunnel.session.control_persist is None + open_session.assert_not_called() + + def test_reuse_only_operations_cannot_fall_back_or_prompt(self, context): + session = _OpenSSHSession( + host="login-ptyche", + user="test_user", + port=None, + identity=None, + control_persist=None, + control_path=None, + require_existing_master=True, + ) + context.run.return_value.ok = True + + session.run("squeue") + run_command = context.run.call_args_list[-1].args[0] + session.put("local", "/remote/file") + put_command = context.run.call_args_list[-1].args[0] + + for command in (run_command, put_command, session.ssh_options): + assert "ControlMaster=no" in command + assert "BatchMode=yes" in command + assert "ProxyCommand=false" in command + + def test_recovery_command_includes_explicit_overrides(self, context, tmp_path): + control_path = tmp_path / "socket path-%C" + session = _OpenSSHSession( + host="login-ptyche", + user="test user", + port=2222, + identity="/key path/id_ed25519", + control_persist=None, + control_path=str(control_path), + require_existing_master=True, + ) + context.run.return_value.ok = False + + with pytest.raises(RuntimeError, match="No existing OpenSSH control master") as error: + session.open() + + message = str(error.value) + assert "ControlPath=" in message + assert str(control_path) in message + assert "-p 2222" in message + assert "/key path/id_ed25519" in message + assert "ControlMaster=yes" in message + assert "test user@login-ptyche" in message + + def test_existing_master_mode_does_not_create_connection(self, context): + session = _OpenSSHSession( + host="login-ptyche", + user="test_user", + port=None, + identity=None, + control_persist=None, + control_path=None, + require_existing_master=True, + ) + context.run.return_value.ok = False + + with pytest.raises(RuntimeError, match=r"ControlMaster=yes.*test_user@login-ptyche"): + session.open() + + context.run.assert_called_once() + check_command = context.run.call_args.args[0] + assert "-O check" in check_command + assert "ControlPath=" not in check_command + assert "ControlPersist=" not in check_command + assert " -p " not in check_command + assert " -i " not in check_command + + def test_control_path_requires_openssh(self): + with pytest.raises(ValueError, match="control_path requires use_openssh"): + SSHTunnel( + host="test.host", + user="test_user", + job_dir="/remote/job", + control_path="/tmp/control-%C", + ) + + def test_existing_master_requires_openssh(self): + with pytest.raises(ValueError, match="require_existing_master requires use_openssh"): + SSHTunnel( + host="test.host", + user="test_user", + job_dir="/remote/job", + require_existing_master=True, + ) + + @patch("nemo_run.core.tunnel.client.shutil.which", return_value="/usr/bin/ssh") + def test_creation_mode_requires_control_persist(self, _): + with pytest.raises(ValueError, match="creation requires control_persist"): + SSHTunnel(host="test.host", user="test_user", job_dir="/remote/job", use_openssh=True) + + @patch("nemo_run.core.tunnel.client.shutil.which", return_value="/usr/bin/ssh") + def test_existing_master_rejects_control_persist(self, _): + with pytest.raises(ValueError, match="cannot be combined"): + SSHTunnel( + host="test.host", + user="test_user", + job_dir="/remote/job", + use_openssh=True, + require_existing_master=True, + control_persist="1d", + ) + + def test_empty_control_persist_is_rejected(self): + with pytest.raises(ValueError, match="control_persist must not be empty"): + SSHTunnel(host="test.host", user="test_user", job_dir="/remote/job", control_persist="") + + class TestSSHConfigFile: def test_init_default_path(self): with patch("os.path.expanduser", return_value="/home/user/.ssh/config"): diff --git a/test/core/tunnel/test_rsync.py b/test/core/tunnel/test_rsync.py index ddc0ba41..8bce2971 100644 --- a/test/core/tunnel/test_rsync.py +++ b/test/core/tunnel/test_rsync.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -15,6 +15,7 @@ """Tests for the rsync module.""" +import shlex import unittest from unittest.mock import Mock, call, patch @@ -68,7 +69,8 @@ def test_rsync_with_exclude_string(self): rsync(self.mock_connection, self.source, self.target, exclude=exclude_pattern) cmd = self.mock_connection.local.call_args[0][0] - self.assertIn(f'--exclude "{exclude_pattern}"', cmd) + tokens = shlex.split(cmd) + self.assertEqual(tokens[tokens.index("--exclude") + 1], exclude_pattern) def test_rsync_with_exclude_list(self): """Test rsync with a list of exclude patterns.""" @@ -76,8 +78,9 @@ def test_rsync_with_exclude_list(self): rsync(self.mock_connection, self.source, self.target, exclude=exclude_patterns) cmd = self.mock_connection.local.call_args[0][0] - for pattern in exclude_patterns: - self.assertIn(f'--exclude "{pattern}"', cmd) + tokens = shlex.split(cmd) + excluded = [tokens[index + 1] for index, token in enumerate(tokens) if token == "--exclude"] + self.assertEqual(excluded, exclude_patterns) def test_rsync_with_exclude_generator(self): """Test rsync with a generator of exclude patterns.""" @@ -86,8 +89,9 @@ def test_rsync_with_exclude_generator(self): rsync(self.mock_connection, self.source, self.target, exclude=exclude_patterns) cmd = self.mock_connection.local.call_args[0][0] - for pattern in ["*.log", "*.tmp", ".git/"]: - self.assertIn(f'--exclude "{pattern}"', cmd) + tokens = shlex.split(cmd) + excluded = [tokens[index + 1] for index, token in enumerate(tokens) if token == "--exclude"] + self.assertEqual(excluded, exclude_patterns) def test_rsync_with_delete(self): """Test rsync with delete flag enabled.""" @@ -111,6 +115,57 @@ def test_rsync_with_custom_ssh_opts(self): cmd = self.mock_connection.local.call_args[0][0] self.assertIn(ssh_opts, cmd) + def test_rsync_without_explicit_port_uses_ssh_config(self): + self.mock_connection.port = None + + rsync(self.mock_connection, "source", "target") + + cmd = self.mock_connection.local.call_args[0][0] + self.assertNotIn("-p None", cmd) + + def test_rsync_with_session_ssh_options(self): + self.mock_connection.ssh_options = "-o ControlMaster=auto -o ControlPath=/tmp/control-%C" + + rsync(self.mock_connection, self.source, self.target) + + cmd = self.mock_connection.local.call_args[0][0] + self.assertIn(self.mock_connection.ssh_options, cmd) + + def test_rsync_quotes_remote_shell_once(self): + control_path = "/tmp/socket path/'quote';$(touch injected)-%C" + self.mock_connection.ssh_options = shlex.join( + [ + "-o", + f"ControlPath={control_path}", + "-o", + "ControlMaster=no", + "-o", + "BatchMode=yes", + "-o", + "ProxyCommand=false", + ] + ) + + rsync(self.mock_connection, self.source, self.target) + + command = self.mock_connection.local.call_args.args[0] + tokens = shlex.split(command) + remote_shell = tokens[tokens.index("--rsh") + 1] + assert shlex.split(remote_shell) == [ + "ssh", + "-p", + "22", + "-o", + f"ControlPath={control_path}", + "-o", + "ControlMaster=no", + "-o", + "BatchMode=yes", + "-o", + "ProxyCommand=false", + ] + assert "touch injected" not in tokens + def test_rsync_with_custom_rsync_opts(self): """Test rsync with custom rsync options.""" rsync_opts = "--checksum"