diff --git a/cvs/lib/parallel/interfaces.py b/cvs/lib/parallel/interfaces.py new file mode 100644 index 00000000..f98c7e70 --- /dev/null +++ b/cvs/lib/parallel/interfaces.py @@ -0,0 +1,42 @@ +''' +Copyright 2025 Advanced Micro Devices, Inc. +All rights reserved. This notice is intended as a precaution against inadvertent publication and does not imply publication or any waiver of confidentiality. +The year included in the foregoing notice is the year of creation of the work. +All code contained here is Property of Advanced Micro Devices, Inc. +''' + +from abc import ABC, abstractmethod + + +class ShardableSshInterface(ABC): + """Abstract base class defining operations that MUST support sharding. + + Any class implementing this interface must provide sharded implementations + for all these methods. This prevents silent performance degradation where + operations fall back to non-sharded implementations. + """ + + @abstractmethod + def exec(self, cmd, timeout=None, print_console=True, detailed=False): + """Execute command - must support sharding for performance.""" + pass + + @abstractmethod + def exec_cmd_list(self, cmd_list, timeout=None, print_console=True): + """Execute command list - must support sharding for performance.""" + pass + + @abstractmethod + def upload_file(self, local_file, remote_file, recurse=False): + """Upload file via SFTP - must support sharding for performance.""" + pass + + @abstractmethod + def download_file(self, remote_file, local_file, recurse=False, suffix_separator='_'): + """Download file via SFTP - must support sharding for performance.""" + pass + + @abstractmethod + def reboot_connections(self): + """Reboot connections - must support sharding for performance.""" + pass diff --git a/cvs/lib/parallel/multiprocess_pssh.py b/cvs/lib/parallel/multiprocess_pssh.py index 8e2e8c4b..4d9ce9fb 100644 --- a/cvs/lib/parallel/multiprocess_pssh.py +++ b/cvs/lib/parallel/multiprocess_pssh.py @@ -9,12 +9,13 @@ from cvs.lib.parallel.pssh import Pssh from cvs.lib.parallel.config import ParallelConfig from cvs.lib.parallel.pssh_sharder import PsshSharder +from cvs.lib.parallel.interfaces import ShardableSshInterface from cvs.lib import globals global_log = globals.log -class MultiProcessPssh(Pssh): +class MultiProcessPssh(ShardableSshInterface): """ Multi-process parallel SSH with automatic host sharding for large host lists. @@ -47,10 +48,16 @@ def __init__( n = len(host_list) if host_list is not None else 0 use_mp = hosts_per_shard > 0 and n > hosts_per_shard - if not use_mp: - # Use single-process base class - super().__init__( - log, # Pass through log parameter to parent + if use_mp: + # Initialize for multi-process sharding - no Pssh instance needed + self._init_sharded(log, host_list, user, password, pkey, host_key_check, stop_on_errors, env_vars) + self.sharder = PsshSharder(self.config) + self._use_process_sharding = True + self.pssh = None # No Pssh instance needed for sharded mode + else: + # No sharding - create Pssh instance for delegation + self.pssh = Pssh( + log, host_list, user, password, @@ -63,12 +70,6 @@ def __init__( self._use_process_sharding = False # Ensure attributes needed by _shard_init_kwargs are available self.env_vars = env_vars - else: - # Initialize for multi-process sharding - self._init_sharded(log, host_list, user, password, pkey, host_key_check, stop_on_errors, env_vars) - - # Create the sharder - no registration needed with direct operations! - self.sharder = PsshSharder(self.config) def _init_sharded( self, @@ -116,9 +117,8 @@ def _shard_init_kwargs(self): def _merge_shard_returns(self, shard_returns, merge_unreachable=True): """ - Update reachable/unreachable host lists from shard workers and convert results to cmd_output dict. + Update reachable/unreachable host lists from shard workers and merge results to cmd_output dict. - Handles both state updates and SimpleHostOutput conversion in one place for cleaner architecture. Returns cmd_output dict ready for _print_merged_outputs. """ # 1. Update host lists (existing logic) @@ -132,18 +132,12 @@ def _merge_shard_returns(self, shard_returns, merge_unreachable=True): if u not in self.unreachable_hosts: self.unreachable_hosts.append(u) - # 2. Convert SimpleHostOutput objects to cmd_output dict + # 2. Merge results from shard workers (always dict format) merged = {} for shard in shard_returns: result = shard.get('result') or {} # treat None as {} (for scp/reboot) - - if isinstance(result, list): - # Handle SimpleHostOutput objects - convert using _process_output - temp_dict = self._process_output(result, print_console=False) - merged.update(temp_dict) - else: - # Handle old dict format (for scp/reboot) - merged.update(result) + # Shard workers always return processed dicts from _process_output + merged.update(result) # 3. Preserve original host order cmd_output = {} @@ -171,10 +165,10 @@ def _print_merged_outputs(self, cmd_output, cmd=None, cmd_list=None, print_conso for line in cmd_output[host].splitlines(): self.log.info("%s", line) - def exec(self, cmd, timeout=None, print_console=True): + def exec(self, cmd, timeout=None, print_console=True, detailed=False): """Execute command with automatic sharding if needed.""" - if not getattr(self, '_use_process_sharding', False): - return super().exec(cmd, timeout=timeout, print_console=print_console) + if not self._use_process_sharding: + return self.pssh.exec(cmd, timeout=timeout, print_console=print_console, detailed=detailed) if self.env_prefix: full_cmd = f"{self.env_prefix} ; {cmd}" @@ -195,7 +189,7 @@ def exec(self, cmd, timeout=None, print_console=True): # Use sharder for sharded execution host_chunks = list(self.sharder.chunk_hosts(self.reachable_hosts)) payloads = self.sharder.create_payloads( - 'exec', host_chunks, self._shard_init_kwargs(), cmd=cmd, timeout=timeout + 'exec', host_chunks, self._shard_init_kwargs(), cmd=cmd, timeout=timeout, detailed=detailed ) shard_returns = self.sharder.execute_sharded(payloads) @@ -213,8 +207,8 @@ def exec(self, cmd, timeout=None, print_console=True): def exec_cmd_list(self, cmd_list, timeout=None, print_console=True): """Execute command list with automatic sharding if needed.""" - if not getattr(self, '_use_process_sharding', False): - return super().exec_cmd_list(cmd_list, timeout=timeout, print_console=print_console) + if not self._use_process_sharding: + return self.pssh.exec_cmd_list(cmd_list, timeout=timeout, print_console=print_console) if len(cmd_list) != len(self.host_list): raise ValueError('cmd_list length must match host_list length') @@ -256,7 +250,7 @@ def exec_cmd_list(self, cmd_list, timeout=None, print_console=True): offset += k payloads.append( self.sharder.create_payloads( - 'cmd_list', [shard_hosts], self._shard_init_kwargs(), cmd_list=shard_commands, timeout=timeout + 'exec_cmd_list', [shard_hosts], self._shard_init_kwargs(), cmd_list=shard_commands, timeout=timeout )[0] ) @@ -275,15 +269,26 @@ def exec_cmd_list(self, cmd_list, timeout=None, print_console=True): return cmd_output def scp_file(self, local_file, remote_file, recurse=False): - """Copy file with automatic sharding if needed.""" - if not getattr(self, '_use_process_sharding', False): - return super().scp_file(local_file, remote_file, recurse=recurse) + """ + Backward-compatible alias for upload_file. + Kept so existing callers (and log-grep tooling looking for the legacy + "About to copy local file..." line) keep working. New code should call + upload_file directly. + """ self.log.info('About to copy local file {} to remote {} on all Hosts'.format(local_file, remote_file)) + return self.upload_file(local_file, remote_file, recurse=recurse) + + def upload_file(self, local_file, remote_file, recurse=False): + """Upload file with automatic sharding if needed.""" + if not self._use_process_sharding: + return self.pssh.upload_file(local_file, remote_file, recurse=recurse) + + self.log.info('SFTP upload %s -> %s on %s', local_file, remote_file, self.reachable_hosts) host_chunks = list(self.sharder.chunk_hosts(self.reachable_hosts)) payloads = self.sharder.create_payloads( - 'scp', + 'upload_file', host_chunks, self._shard_init_kwargs(), local_file=local_file, @@ -291,24 +296,48 @@ def scp_file(self, local_file, remote_file, recurse=False): recurse=recurse, ) shard_returns = self.sharder.execute_sharded(payloads) - return self._merge_shard_returns(shard_returns, merge_unreachable=False) + return self._merge_shard_returns(shard_returns) + + def download_file(self, remote_file, local_file, recurse=False, suffix_separator='_'): + """Download file with automatic sharding if needed.""" + if not self._use_process_sharding: + return self.pssh.download_file(remote_file, local_file, recurse=recurse, suffix_separator=suffix_separator) + + self.log.info('SFTP download %s -> %s from %s', remote_file, local_file, self.reachable_hosts) + + host_chunks = list(self.sharder.chunk_hosts(self.reachable_hosts)) + payloads = self.sharder.create_payloads( + 'download_file', + host_chunks, + self._shard_init_kwargs(), + remote_file=remote_file, + local_file=local_file, + recurse=recurse, + suffix_separator=suffix_separator, + ) + shard_returns = self.sharder.execute_sharded(payloads) + return self._merge_shard_returns(shard_returns) def reboot_connections(self): """Reboot connections with automatic sharding if needed.""" - if not getattr(self, '_use_process_sharding', False): - return super().reboot_connections() + if not self._use_process_sharding: + return self.pssh.reboot_connections() self.log.info('Rebooting Connections') host_chunks = list(self.sharder.chunk_hosts(self.reachable_hosts)) - payloads = self.sharder.create_payloads('reboot', host_chunks, self._shard_init_kwargs()) + payloads = self.sharder.create_payloads('reboot_connections', host_chunks, self._shard_init_kwargs()) shard_returns = self.sharder.execute_sharded(payloads) return self._merge_shard_returns(shard_returns, merge_unreachable=False) def destroy_clients(self): """Destroy clients - handle both sharded and non-sharded modes.""" - if getattr(self, '_use_process_sharding', False): - self.log.info('Destroying Current phdl connections ..') + if self._use_process_sharding: + self.log.info('Cleaning up sharded mode state ..') + # In sharded mode, no persistent connections to destroy self.client = None - return - super().destroy_clients() + else: + self.log.info('Destroying Current phdl connections ..') + # In non-sharded mode, properly destroy the Pssh instance connections + if self.pssh is not None: + self.pssh.destroy_clients() diff --git a/cvs/lib/parallel/pssh.py b/cvs/lib/parallel/pssh.py index 1354dd75..26a58203 100644 --- a/cvs/lib/parallel/pssh.py +++ b/cvs/lib/parallel/pssh.py @@ -17,17 +17,6 @@ global_log = globals.log -class SimpleHostOutput: - """Simple HostOutput-compatible object for consistent processing across process boundaries.""" - - def __init__(self, host, stdout_lines, stderr_lines, exception, exit_code=0): - self.host = host - self.stdout = iter(stdout_lines) # Convert list to iterator for _process_output compatibility - self.stderr = iter(stderr_lines) # Convert list to iterator for _process_output compatibility - self.exception = exception - self.exit_code = exit_code - - class Pssh: """ Single-process parallel SSH: one ParallelSSHClient (one gevent hub) over a host list. @@ -385,32 +374,6 @@ def reboot_connections(self): self.log.info('Rebooting Connections') self.client.run_command('reboot -f', stop_on_errors=self.stop_on_errors) - def _extract_simple_data(self, output): - """ - Extract essential data from pssh.output.HostOutput objects for safe IPC transfer. - - Converts non-picklable HostOutput objects (containing active SSH channels, - generators, and C extension objects) into picklable SimpleHostOutput objects - by consuming stdout/stderr generators into lists and extracting core attributes. - """ - - simple_outputs = [] - for item in output: - # Consume generators to lists for pickling - stdout_lines = list(item.stdout or []) - stderr_lines = list(item.stderr or []) - - simple_output = SimpleHostOutput( - host=item.host, - stdout_lines=stdout_lines, - stderr_lines=stderr_lines, - exception=item.exception, - exit_code=getattr(item, 'exit_code', 0), - ) - simple_outputs.append(simple_output) - - return simple_outputs - def destroy_clients(self): self.log.info('Destroying Current phdl connections ..') del self.client diff --git a/cvs/lib/parallel/pssh_sharder.py b/cvs/lib/parallel/pssh_sharder.py index b1f27994..3c3d47dc 100644 --- a/cvs/lib/parallel/pssh_sharder.py +++ b/cvs/lib/parallel/pssh_sharder.py @@ -5,10 +5,19 @@ All code contained here is Property of Advanced Micro Devices, Inc. ''' +import inspect import multiprocessing as mp from concurrent.futures import ProcessPoolExecutor, as_completed from cvs.lib.parallel.pssh import Pssh +from cvs.lib.parallel.interfaces import ShardableSshInterface + + +# Dynamically discover supported operations from ABC (computed once at import time) + +SUPPORTED_OPERATIONS = { + name for name, method in inspect.getmembers(ShardableSshInterface) if getattr(method, '__isabstractmethod__', False) +} class PsshSharder: @@ -74,6 +83,8 @@ def run_shard(payload): """ Run an SSH operation on a shard of hosts (must be picklable for multiprocessing). + Dynamically supports all abstract methods defined in ShardableSshInterface. + Args: payload: Dict with 'operation' (operation type), 'init' (SSH setup), and operation args @@ -84,29 +95,35 @@ def run_shard(payload): if not isinstance(operation, str): raise TypeError('payload["operation"] must be str, got %r' % (type(operation),)) + # Validate operation is supported by ABC + if operation not in SUPPORTED_OPERATIONS: + raise ValueError(f'Unknown operation: {operation}. Supported: {sorted(SUPPORTED_OPERATIONS)}') + # Create SSH client for this shard of hosts init_kwargs = payload['init'] init_kwargs['process_output'] = False # Force raw output mode for sharding shard = Pssh(**init_kwargs) try: - # Direct operation calls - no registry needed! - if operation == 'exec': - result = shard.exec(payload['cmd'], timeout=payload.get('timeout'), print_console=False) - elif operation == 'cmd_list': - result = shard.exec_cmd_list( - payload['cmd_list'], - timeout=payload.get('timeout'), - print_console=False, - ) - elif operation == 'scp': - shard.scp_file(payload['local_file'], payload['remote_file'], recurse=payload.get('recurse', False)) - result = None - elif operation == 'reboot': - shard.reboot_connections() + # Ensure method exists in Pssh + if not hasattr(shard, operation): + raise RuntimeError(f'Method {operation} not found in Pssh class') + + shard_method = getattr(shard, operation) + + # Extract operation arguments from payload (excluding 'operation' and 'init') + args = {k: v for k, v in payload.items() if k not in ['operation', 'init']} + + # Add common parameters for operations that support them + if operation in ['exec', 'exec_cmd_list']: + args['print_console'] = False # Always False for sharded operations + + # Call the method dynamically + result = shard_method(**args) + + # Handle operations that should return None (void operations) + if operation in ['upload_file', 'reboot_connections']: result = None - else: - raise ValueError(f'Unknown operation: {operation}') return { 'result': result, diff --git a/cvs/lib/parallel/unittests/test_multiprocess_pssh.py b/cvs/lib/parallel/unittests/test_multiprocess_pssh.py index eec7a8d2..a6c251d4 100644 --- a/cvs/lib/parallel/unittests/test_multiprocess_pssh.py +++ b/cvs/lib/parallel/unittests/test_multiprocess_pssh.py @@ -8,22 +8,30 @@ class TestMultiProcessPsshInitialization(unittest.TestCase): def setUp(self): self.mock_log = MagicMock() self.host_list = ["host1", "host2", "host3", "host4", "host5"] + # Mock Pssh.__init__ for all tests in this class + self.pssh_patcher = patch('cvs.lib.parallel.multiprocess_pssh.Pssh.__init__') + self.mock_pssh_init = self.pssh_patcher.start() + self.mock_pssh_init.return_value = None - @patch('cvs.lib.parallel.multiprocess_pssh.Pssh.__init__') - def test_init_no_sharding_small_host_list(self, mock_parent_init): + def tearDown(self): + self.pssh_patcher.stop() + + def test_init_no_sharding_small_host_list(self): """Test initialization without sharding for small host lists.""" - mock_parent_init.return_value = None small_host_list = ["host1", "host2"] config = ParallelConfig(hosts_per_shard=32) pssh = MultiProcessPssh(self.mock_log, small_host_list, user="test", config=config) - # Should call parent __init__ for small lists - mock_parent_init.assert_called_once_with( + # Always creates composed Pssh instance with ABC+composition + self.mock_pssh_init.assert_called_once_with( self.mock_log, small_host_list, "test", None, 'id_rsa', False, True, None, process_output=True ) - # For small lists, no sharder should be created + # For small lists, no sharder should be created (key difference!) self.assertFalse(hasattr(pssh, 'sharder')) + self.assertFalse(pssh._use_process_sharding) + # But pssh instance should always exist + self.assertTrue(hasattr(pssh, 'pssh')) @patch('cvs.lib.parallel.multiprocess_pssh.MultiProcessPssh._init_sharded') @patch('cvs.lib.parallel.multiprocess_pssh.PsshSharder') @@ -33,174 +41,120 @@ def test_init_with_sharding_large_host_list(self, mock_sharder_class, mock_init_ mock_sharder = MagicMock() mock_sharder_class.return_value = mock_sharder - MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + # In sharded mode, should NOT create Pssh instance (avoids bottleneck) + self.mock_pssh_init.assert_not_called() # Should call _init_sharded for large lists mock_init_sharded.assert_called_once_with( self.mock_log, self.host_list, "test", None, 'id_rsa', False, True, None ) # Verify sharder was created mock_sharder_class.assert_called_once_with(config) + self.assertTrue(pssh._use_process_sharding) + # In sharded mode: pssh is None, sharder exists + self.assertIsNone(pssh.pssh) + self.assertTrue(hasattr(pssh, 'sharder')) def test_init_with_custom_config(self): """Test initialization with custom configuration.""" config = ParallelConfig(hosts_per_shard=64, max_workers_per_cpu=8) - with patch('cvs.lib.parallel.multiprocess_pssh.Pssh.__init__') as mock_parent_init: - mock_parent_init.return_value = None - pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", password="pass", config=config) - - self.assertEqual(pssh.config, config) - mock_parent_init.assert_called_once() - - def test_init_default_config(self): - """Test initialization with default configuration.""" - with patch('cvs.lib.parallel.multiprocess_pssh.Pssh.__init__') as mock_parent_init: - mock_parent_init.return_value = None - pssh = MultiProcessPssh(self.mock_log, ["host1"], user="test") - - # Should create default config - self.assertIsNotNone(pssh.config) - self.assertIsInstance(pssh.config, ParallelConfig) - - -class TestMultiProcessPsshInitSharded(unittest.TestCase): - def setUp(self): - self.mock_log = MagicMock() - self.host_list = ["host1", "host2"] + pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", password="pass", config=config) - @patch('cvs.lib.parallel.multiprocess_pssh.PsshSharder') - def test_init_sharded_sets_attributes(self, mock_sharder_class): - """Test that _init_sharded sets all required attributes.""" - pssh = MultiProcessPssh.__new__(MultiProcessPssh) # Create without calling __init__ - # Set up config and sharder that _init_sharded expects - pssh.config = ParallelConfig(hosts_per_shard=2, max_workers_per_cpu=1) - mock_sharder = MagicMock() - pssh.sharder = mock_sharder - - pssh._init_sharded( - self.mock_log, self.host_list, "testuser", "testpass", "custom_key", True, False, {"VAR": "value"} - ) - - # Logger is now module-level, no longer stored in pssh.log - self.assertEqual(pssh.host_list, self.host_list) - self.assertEqual(pssh.reachable_hosts, self.host_list) - self.assertEqual(pssh.user, "testuser") - self.assertEqual(pssh.password, "testpass") - self.assertEqual(pssh.pkey, "custom_key") - self.assertTrue(pssh.host_key_check) - self.assertFalse(pssh.stop_on_errors) - self.assertEqual(pssh.env_vars, {"VAR": "value"}) - self.assertTrue(pssh._use_process_sharding) + self.assertEqual(pssh.config, config) class TestMultiProcessPsshExec(unittest.TestCase): def setUp(self): self.mock_log = MagicMock() self.host_list = ["host1", "host2"] + # Mock Pssh.__init__ for all tests in this class + self.pssh_patcher = patch('cvs.lib.parallel.multiprocess_pssh.Pssh.__init__') + self.mock_pssh_init = self.pssh_patcher.start() + self.mock_pssh_init.return_value = None - def test_exec_no_sharding(self): - """Test exec without sharding calls parent method.""" - with patch('cvs.lib.parallel.multiprocess_pssh.Pssh.__init__') as mock_parent_init: - with patch('cvs.lib.parallel.multiprocess_pssh.Pssh.exec') as mock_parent_exec: - mock_parent_init.return_value = None - mock_parent_exec.return_value = {"host1": "result"} + # Mock execute_sharded to avoid process spawning while testing real payload creation + self.execute_sharded_patcher = patch('cvs.lib.parallel.pssh_sharder.PsshSharder.execute_sharded') + self.mock_execute_sharded = self.execute_sharded_patcher.start() - pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test") - result = pssh.exec("uptime") - - mock_parent_exec.assert_called_once_with("uptime", timeout=None, print_console=True) - self.assertEqual(result, {"host1": "result"}) + def tearDown(self): + self.pssh_patcher.stop() + self.execute_sharded_patcher.stop() def test_exec_with_sharding(self): """Test exec with sharding uses sharder.""" config = ParallelConfig(hosts_per_shard=1, max_workers_per_cpu=1) - with patch('cvs.lib.parallel.multiprocess_pssh.MultiProcessPssh._init_sharded'): - with patch('cvs.lib.parallel.multiprocess_pssh.PsshSharder') as mock_sharder_class: - mock_sharder = MagicMock() - mock_sharder_class.return_value = mock_sharder + # Mock execute_sharded to return realistic shard results + self.mock_execute_sharded.return_value = [ + {'result': {"host1": "up1"}, 'reachable_hosts': ["host1"], 'unreachable_hosts': []}, + {'result': {"host2": "up2"}, 'reachable_hosts': ["host2"], 'unreachable_hosts': []}, + ] - # Set up the mock sharder behavior - mock_sharder.chunk_hosts.return_value = [["host1"], ["host2"]] - mock_sharder.create_payloads.return_value = [ - {"operation": "exec", "init": {}, "cmd": "uptime"}, - {"operation": "exec", "init": {}, "cmd": "uptime"}, - ] - mock_sharder.execute_sharded.return_value = [ - {"result": {"host1": "up1"}, "reachable_hosts": ["host1"], "unreachable_hosts": []}, - {"result": {"host2": "up2"}, "reachable_hosts": ["host2"], "unreachable_hosts": []}, - ] - mock_sharder.merge_results.return_value = {"host1": "up1", "host2": "up2"} + pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + result = pssh.exec("uptime", timeout=30, detailed=True) - pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + # Verify execute_sharded was called with real payloads from create_payloads + self.mock_execute_sharded.assert_called_once() + payloads = self.mock_execute_sharded.call_args[0][0] - # Manually set attributes that _init_sharded would set - pssh.env_prefix = None - pssh.host_list = self.host_list - pssh.reachable_hosts = self.host_list - pssh.sharder = mock_sharder - pssh._use_process_sharding = True # Enable sharding for this test - # Add attributes needed by _shard_init_kwargs() - pssh.user = "test" - pssh.password = None - pssh.pkey = "id_rsa" - pssh.host_key_check = False - pssh.stop_on_errors = True - pssh.env_vars = None + # Check that real payloads were created correctly + self.assertEqual(len(payloads), 2) # Two shards (hosts_per_shard=1) + first_payload = payloads[0] + self.assertEqual(first_payload['operation'], 'exec') + self.assertEqual(first_payload['cmd'], 'uptime') + self.assertEqual(first_payload['timeout'], 30) + self.assertEqual(first_payload['detailed'], True) + self.assertIn('init', first_payload) - result = pssh.exec("uptime") + # Verify real init kwargs are properly constructed by _shard_init_kwargs + init_kwargs = first_payload['init'] + self.assertEqual(init_kwargs['user'], 'test') + self.assertIn('host_list', init_kwargs) - # Verify sharder methods were called - mock_sharder.chunk_hosts.assert_called_once_with(self.host_list) - mock_sharder.execute_sharded.assert_called_once() - self.assertEqual(result, {"host1": "up1", "host2": "up2"}) + # Verify result is merged from shards + self.assertEqual(result, {"host1": "up1", "host2": "up2"}) def test_exec_cmd_list_with_sharding(self): """Test exec_cmd_list with sharding uses sharder.""" config = ParallelConfig(hosts_per_shard=1, max_workers_per_cpu=1) cmd_list = ["uptime", "date"] - with patch('cvs.lib.parallel.multiprocess_pssh.MultiProcessPssh._init_sharded'): - with patch('cvs.lib.parallel.multiprocess_pssh.PsshSharder') as mock_sharder_class: - mock_sharder = MagicMock() - mock_sharder_class.return_value = mock_sharder + # Mock execute_sharded to capture real payloads and return results + def capture_payloads_and_return_results(payloads): + self.captured_payloads = payloads + return [ + { + 'result': {"host1": "up1", "host2": "date2"}, + 'reachable_hosts': ["host1", "host2"], + 'unreachable_hosts': [], + } + ] - # Set up the mock sharder behavior - create_payloads returns a list, we need the [0] element - mock_sharder.create_payloads.return_value = [ - {"operation": "cmd_list", "init": {}, "cmd_list": cmd_list} - ] - mock_sharder.execute_sharded.return_value = [ - { - "result": {"host1": "up1", "host2": "date2"}, - "reachable_hosts": ["host1", "host2"], - "unreachable_hosts": [], - }, - ] - mock_sharder.merge_results.return_value = {"host1": "up1", "host2": "date2"} + self.mock_execute_sharded.side_effect = capture_payloads_and_return_results - pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + result = pssh.exec_cmd_list(cmd_list, timeout=60) - # Manually set attributes that _init_sharded would set - pssh.env_prefix = None - pssh.host_list = self.host_list - pssh.reachable_hosts = self.host_list - pssh.config = config - pssh.sharder = mock_sharder - pssh._use_process_sharding = True # Enable sharding for this test - # Additional attributes needed by _shard_init_kwargs - pssh.user = "test" - pssh.password = None - pssh.pkey = "id_rsa" - pssh.host_key_check = False - pssh.stop_on_errors = True - pssh.env_vars = None + # Verify execute_sharded was called with real payloads + self.mock_execute_sharded.assert_called_once() - result = pssh.exec_cmd_list(cmd_list) + # Check that real payloads were created correctly by create_payloads + self.assertEqual(len(self.captured_payloads), 2) # Two shards + first_payload = self.captured_payloads[0] + self.assertEqual(first_payload['operation'], 'exec_cmd_list') + # With hosts_per_shard=1, each shard gets one command from the list + self.assertEqual(first_payload['cmd_list'], ['uptime']) # First shard gets first command + self.assertEqual(first_payload['timeout'], 60) + self.assertIn('init', first_payload) - # Verify sharder methods were called - mock_sharder.execute_sharded.assert_called_once() - self.assertEqual(result, {"host1": "up1", "host2": "date2"}) + # Verify real init kwargs + init_kwargs = first_payload['init'] + self.assertEqual(init_kwargs['user'], 'test') + + # Verify result is merged from shards + self.assertEqual(result, {"host1": "up1", "host2": "date2"}) def test_exec_cmd_list_with_unreachable_hosts(self): """Test exec_cmd_list correctly maps commands when some hosts are unreachable.""" @@ -219,7 +173,11 @@ def test_exec_cmd_list_with_unreachable_hosts(self): # Mock sharder behavior mock_sharder.chunk_hosts.return_value = [["host1", "host3"], ["host4"]] mock_sharder.create_payloads.return_value = [ - {"operation": "cmd_list", "init": {}, "cmd_list": ["cmd1", "cmd3"]}, # Expected: correct commands + { + "operation": "exec_cmd_list", + "init": {}, + "cmd_list": ["cmd1", "cmd3"], + }, # Expected: correct commands ] mock_sharder.execute_sharded.return_value = [ { @@ -260,98 +218,76 @@ def test_exec_cmd_list_with_unreachable_hosts(self): self.assertEqual(result, {"host1": "result1", "host3": "result3", "host4": "result4"}) def test_scp_file_with_sharding(self): - """Test scp_file with sharding uses sharder.""" + """Test scp_file delegates to upload_file with sharding.""" config = ParallelConfig(hosts_per_shard=1, max_workers_per_cpu=1) - with patch('cvs.lib.parallel.multiprocess_pssh.MultiProcessPssh._init_sharded'): - with patch('cvs.lib.parallel.multiprocess_pssh.PsshSharder') as mock_sharder_class: - mock_sharder = MagicMock() - mock_sharder_class.return_value = mock_sharder + # Mock execute_sharded for void scp_file operation + self.mock_execute_sharded.return_value = [ + {'result': None, 'reachable_hosts': ["host1"], 'unreachable_hosts': []}, + {'result': None, 'reachable_hosts': ["host2"], 'unreachable_hosts': []}, + ] - # Set up the mock sharder behavior - mock_sharder.chunk_hosts.return_value = [["host1"], ["host2"]] - mock_sharder.create_payloads.return_value = [ - {"operation": "scp", "init": {}, "local_file": "test.txt", "remote_file": "/tmp/test.txt"}, - ] - mock_sharder.execute_sharded.return_value = [ - { - "result": {"host1": "scp1", "host2": "scp2"}, - "reachable_hosts": ["host1", "host2"], - "unreachable_hosts": [], - }, - ] - mock_sharder.merge_results.return_value = {"host1": "scp1", "host2": "scp2"} - - pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + result = pssh.scp_file("test.txt", "/tmp/test.txt", recurse=False) - # Manually set attributes that _init_sharded would set - pssh.host_list = self.host_list - pssh.reachable_hosts = self.host_list - pssh.sharder = mock_sharder - pssh._use_process_sharding = True # Enable sharding for this test - # Add attributes needed by _shard_init_kwargs() - pssh.user = "test" - pssh.password = None - pssh.pkey = "id_rsa" - pssh.host_key_check = False - pssh.stop_on_errors = True - pssh.env_vars = None + # Verify execute_sharded was called with real payloads + self.mock_execute_sharded.assert_called_once() + payloads = self.mock_execute_sharded.call_args[0][0] - result = pssh.scp_file("test.txt", "/tmp/test.txt") + # Check that real payloads were created for upload_file operation + self.assertEqual(len(payloads), 2) # Two shards + first_payload = payloads[0] + self.assertEqual(first_payload['operation'], 'upload_file') # scp_file delegates to upload_file + self.assertEqual(first_payload['local_file'], 'test.txt') + self.assertEqual(first_payload['remote_file'], '/tmp/test.txt') + self.assertEqual(first_payload['recurse'], False) - # Verify sharder methods were called - mock_sharder.chunk_hosts.assert_called_once_with(self.host_list) - mock_sharder.execute_sharded.assert_called_once() - self.assertEqual(result, {"host1": "scp1", "host2": "scp2"}) + # scp_file should return empty dict (merged void results) + self.assertEqual(result, {}) def test_reboot_connections_with_sharding(self): """Test reboot_connections with sharding uses sharder.""" config = ParallelConfig(hosts_per_shard=1, max_workers_per_cpu=1) - with patch('cvs.lib.parallel.multiprocess_pssh.MultiProcessPssh._init_sharded'): - with patch('cvs.lib.parallel.multiprocess_pssh.PsshSharder') as mock_sharder_class: - mock_sharder = MagicMock() - mock_sharder_class.return_value = mock_sharder + # Mock execute_sharded for void reboot_connections operation + self.mock_execute_sharded.return_value = [ + {'result': None, 'reachable_hosts': ["host1"], 'unreachable_hosts': []}, + {'result': None, 'reachable_hosts': ["host2"], 'unreachable_hosts': []}, + ] - # Set up the mock sharder behavior - mock_sharder.chunk_hosts.return_value = [["host1"], ["host2"]] - mock_sharder.create_payloads.return_value = [{"operation": "reboot", "init": {}}] - mock_sharder.execute_sharded.return_value = [ - { - "result": {"host1": "reboot1", "host2": "reboot2"}, - "reachable_hosts": ["host1", "host2"], - "unreachable_hosts": [], - }, - ] - mock_sharder.merge_results.return_value = {"host1": "reboot1", "host2": "reboot2"} + pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + result = pssh.reboot_connections() - pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + # Verify execute_sharded was called with real payloads + self.mock_execute_sharded.assert_called_once() + payloads = self.mock_execute_sharded.call_args[0][0] - # Manually set attributes that _init_sharded would set - pssh.host_list = self.host_list - pssh.reachable_hosts = self.host_list - pssh.sharder = mock_sharder - pssh._use_process_sharding = True # Enable sharding for this test - # Add attributes needed by _shard_init_kwargs() - pssh.user = "test" - pssh.password = None - pssh.pkey = "id_rsa" - pssh.host_key_check = False - pssh.stop_on_errors = True - pssh.env_vars = None + # Check that real payloads were created correctly + self.assertEqual(len(payloads), 2) # Two shards + first_payload = payloads[0] + self.assertEqual(first_payload['operation'], 'reboot_connections') + self.assertIn('init', first_payload) - result = pssh.reboot_connections() - - # Verify sharder methods were called - mock_sharder.chunk_hosts.assert_called_once_with(self.host_list) - mock_sharder.execute_sharded.assert_called_once() - self.assertEqual(result, {"host1": "reboot1", "host2": "reboot2"}) + # reboot_connections should return empty dict (merged void results) + self.assertEqual(result, {}) class TestMultiProcessPsshHelperMethods(unittest.TestCase): def setUp(self): self.mock_log = MagicMock() self.host_list = ["host1", "host2"] + # Mock Pssh.__init__ for all tests in this class + self.pssh_patcher = patch('cvs.lib.parallel.multiprocess_pssh.Pssh.__init__') + self.mock_pssh_init = self.pssh_patcher.start() + self.mock_pssh_init.return_value = None + + # Mock execute_sharded to avoid process spawning while testing real payload creation + self.execute_sharded_patcher = patch('cvs.lib.parallel.pssh_sharder.PsshSharder.execute_sharded') + self.mock_execute_sharded = self.execute_sharded_patcher.start() + + def tearDown(self): + self.pssh_patcher.stop() + self.execute_sharded_patcher.stop() def test_shard_init_kwargs(self): """Test _shard_init_kwargs creates correct initialization arguments.""" @@ -405,6 +341,61 @@ def test_merge_shard_returns(self): self.assertEqual(set(pssh.reachable_hosts), {"host1", "host2"}) self.assertEqual(pssh.unreachable_hosts, []) + def test_upload_file_with_sharding(self): + """Test upload_file with sharding uses sharder.""" + config = ParallelConfig(hosts_per_shard=1, max_workers_per_cpu=1) + + # Mock execute_sharded for void upload_file operation + self.mock_execute_sharded.return_value = [ + {'result': None, 'reachable_hosts': ["host1"], 'unreachable_hosts': []}, + {'result': None, 'reachable_hosts': ["host2"], 'unreachable_hosts': []}, + ] + + pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + result = pssh.upload_file("test.txt", "/tmp/test.txt", recurse=True) + + # Verify execute_sharded was called with real payloads + self.mock_execute_sharded.assert_called_once() + payloads = self.mock_execute_sharded.call_args[0][0] + + # Check that real payload was created correctly + self.assertEqual(len(payloads), 2) # Two shards + first_payload = payloads[0] + self.assertEqual(first_payload['operation'], 'upload_file') + self.assertEqual(first_payload['local_file'], 'test.txt') + self.assertEqual(first_payload['remote_file'], '/tmp/test.txt') + self.assertEqual(first_payload['recurse'], True) + self.assertIn('init', first_payload) + + # upload_file should return empty dict (merged void results) + self.assertEqual(result, {}) + + def test_download_file_with_sharding(self): + """Test download_file with sharding uses sharder and merges results.""" + config = ParallelConfig(hosts_per_shard=1, max_workers_per_cpu=1) -if __name__ == "__main__": - unittest.main() + # Mock execute_sharded to return realistic download results + self.mock_execute_sharded.return_value = [ + {'result': {"host1": "/tmp/test.txt_host1"}, 'reachable_hosts': ["host1"], 'unreachable_hosts': []}, + {'result': {"host2": "/tmp/test.txt_host2"}, 'reachable_hosts': ["host2"], 'unreachable_hosts': []}, + ] + + pssh = MultiProcessPssh(self.mock_log, self.host_list, user="test", config=config) + result = pssh.download_file("/remote/test.txt", "/tmp/test.txt", suffix_separator="_") + + # Verify execute_sharded was called with real payloads + self.mock_execute_sharded.assert_called_once() + payloads = self.mock_execute_sharded.call_args[0][0] + + # Check that real payload was created correctly + self.assertEqual(len(payloads), 2) # Two shards + first_payload = payloads[0] + self.assertEqual(first_payload['operation'], 'download_file') + self.assertEqual(first_payload['remote_file'], '/remote/test.txt') + self.assertEqual(first_payload['local_file'], '/tmp/test.txt') + self.assertEqual(first_payload['suffix_separator'], '_') + self.assertIn('init', first_payload) + + # download_file should return merged host->path dict + expected_result = {"host1": "/tmp/test.txt_host1", "host2": "/tmp/test.txt_host2"} + self.assertEqual(result, expected_result) diff --git a/cvs/lib/parallel/unittests/test_pssh_sharder.py b/cvs/lib/parallel/unittests/test_pssh_sharder.py index 8c8bbf89..ced6b6f3 100644 --- a/cvs/lib/parallel/unittests/test_pssh_sharder.py +++ b/cvs/lib/parallel/unittests/test_pssh_sharder.py @@ -78,12 +78,12 @@ def test_create_payloads_cmd_list_mode(self): shard_init_kwargs = {'user': 'testuser'} payloads = self.sharder.create_payloads( - 'cmd_list', host_chunks, shard_init_kwargs, cmd_list=['echo 1', 'echo 2'] + 'exec_cmd_list', host_chunks, shard_init_kwargs, cmd_list=['echo 1', 'echo 2'] ) expected = [ { - 'operation': 'cmd_list', + 'operation': 'exec_cmd_list', 'init': {'user': 'testuser', 'host_list': ['host1', 'host2']}, 'cmd_list': ['echo 1', 'echo 2'], } @@ -230,7 +230,7 @@ def test_pssh_shard_worker_exec_mode(self, mock_pssh_class): mock_pssh_class.assert_called_once_with( log=None, host_list=['host1', 'host2'], user='test', process_output=False ) - mock_shard.exec.assert_called_once_with('echo hello', timeout=30, print_console=False) + mock_shard.exec.assert_called_once_with(cmd='echo hello', timeout=30, print_console=False) mock_shard.destroy_clients.assert_called_once() @patch('cvs.lib.parallel.pssh_sharder.Pssh') @@ -243,7 +243,7 @@ def test_pssh_shard_worker_cmd_list_mode(self, mock_pssh_class): mock_shard.unreachable_hosts = [] payload = { - 'operation': 'cmd_list', + 'operation': 'exec_cmd_list', 'init': {'log': None, 'host_list': ['host1'], 'user': 'test'}, 'cmd_list': ['echo 1', 'echo 2'], } @@ -255,19 +255,19 @@ def test_pssh_shard_worker_cmd_list_mode(self, mock_pssh_class): # Verify direct operation call mock_pssh_class.assert_called_once_with(log=None, host_list=['host1'], user='test', process_output=False) - mock_shard.exec_cmd_list.assert_called_once_with(['echo 1', 'echo 2'], timeout=None, print_console=False) + mock_shard.exec_cmd_list.assert_called_once_with(cmd_list=['echo 1', 'echo 2'], print_console=False) mock_shard.destroy_clients.assert_called_once() @patch('cvs.lib.parallel.pssh_sharder.Pssh') - def test_pssh_shard_worker_scp_mode(self, mock_pssh_class): - """Test shard worker in scp mode with direct operation calls.""" + def test_pssh_shard_worker_upload_file_mode(self, mock_pssh_class): + """Test shard worker in upload_file mode with direct operation calls.""" mock_shard = MagicMock() mock_pssh_class.return_value = mock_shard mock_shard.reachable_hosts = ['host1'] mock_shard.unreachable_hosts = [] payload = { - 'operation': 'scp', + 'operation': 'upload_file', 'init': {'log': None, 'host_list': ['host1'], 'user': 'test'}, 'local_file': 'test.txt', 'remote_file': '/tmp/test.txt', @@ -280,7 +280,36 @@ def test_pssh_shard_worker_scp_mode(self, mock_pssh_class): self.assertEqual(result, expected) # Verify direct operation call - mock_shard.scp_file.assert_called_once_with('test.txt', '/tmp/test.txt', recurse=False) + mock_shard.upload_file.assert_called_once_with( + local_file='test.txt', remote_file='/tmp/test.txt', recurse=False + ) + + @patch('cvs.lib.parallel.pssh_sharder.Pssh') + def test_pssh_shard_worker_download_file_mode(self, mock_pssh_class): + """Test shard worker in download_file mode with direct operation calls.""" + mock_shard = MagicMock() + mock_pssh_class.return_value = mock_shard + mock_shard.reachable_hosts = ['host1'] + mock_shard.unreachable_hosts = [] + mock_shard.download_file.return_value = {'host1': '/local/test_host1.txt'} + + payload = { + 'operation': 'download_file', + 'init': {'log': None, 'host_list': ['host1'], 'user': 'test'}, + 'remote_file': '/tmp/test.txt', + 'local_file': 'test.txt', + 'recurse': False, + } + + result = PsshSharder.run_shard(payload) + + expected = {'result': {'host1': '/local/test_host1.txt'}, 'reachable_hosts': ['host1'], 'unreachable_hosts': []} + self.assertEqual(result, expected) + + # Verify direct operation call + mock_shard.download_file.assert_called_once_with( + remote_file='/tmp/test.txt', local_file='test.txt', recurse=False + ) @patch('cvs.lib.parallel.pssh_sharder.Pssh') def test_pssh_shard_worker_reboot_mode(self, mock_pssh_class): @@ -291,7 +320,7 @@ def test_pssh_shard_worker_reboot_mode(self, mock_pssh_class): mock_shard.unreachable_hosts = [] payload = { - 'operation': 'reboot', + 'operation': 'reboot_connections', 'init': {'log': None, 'host_list': ['host1'], 'user': 'test'}, } @@ -319,6 +348,25 @@ def test_pssh_shard_worker_unknown_operation(self, mock_pssh_class): self.assertIn('Unknown operation: unknown_operation', str(cm.exception)) + @patch('cvs.lib.parallel.pssh_sharder.Pssh') + def test_pssh_shard_worker_abc_enforcement(self, mock_pssh_class): + """Test shard worker properly enforces ABC operations.""" + mock_shard = MagicMock() + mock_pssh_class.return_value = mock_shard + + # Test with an operation not in ABC + payload = { + 'operation': 'fake_operation_not_in_abc', + 'init': {'log': None, 'host_list': ['host1'], 'user': 'test'}, + } + + with self.assertRaises(ValueError) as cm: + PsshSharder.run_shard(payload) + + # Verify proper error message + self.assertIn('Unknown operation: fake_operation_not_in_abc', str(cm.exception)) + self.assertIn('Supported:', str(cm.exception)) + if __name__ == "__main__": unittest.main()