Add upload_file and download_file support (#161)

* Add upload_file and download_file support

to MultiProcessPssh with operation registry

This commit addresses API incompatibility issues between Pssh and MultiProcessPssh
classes and implements a robust operation registry system to prevent future drift.

Key Changes:
- Add missing upload_file and download_file methods to MultiProcessPssh class
- Implement operation registry (SUPPORTED_OPERATIONS) as single source of truth
- Use method names directly as operation names to eliminate mapping errors
- Add missing detailed parameter support to MultiProcessPssh.exec()

API Parity Testing:
- Add comprehensive test suite to catch missing methods automatically
- Validate method signatures match between Pssh and MultiProcessPssh
- Verify operation registry completeness and consistency
- Test runtime execution of all sharder operations with proper parameters

Operations Updated:
- exec -> exec (added detailed parameter support)
- upload_file -> upload_file (new operation, matches method name)
- download_file -> download_file (new operation, matches method name)

This ensures MultiProcessPssh maintains full API compatibility with Pssh
and prevents the class of issues that led to missing upload_file/download_file methods.

Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Ignatious Johnson <ichristo@amd.com>

* Refactor to ABC + composition + dynamic sharder architecture

This commit addresses all PR review comments by completely refactoring
the multiprocess SSH architecture to be cleaner, more maintainable, and
more robust.

Key Architectural Changes:
- Replace inheritance with ABC + composition pattern
- Eliminate registry + if/elif tree with dynamic method dispatch
- Create single source of truth via ShardableSshInterface ABC
- Move ABC to separate interfaces.py to avoid circular imports

Review Comments Addressed:
- #1-#2: Fixed operation name mismatches (upload/download)
- #3: Updated test expectations for new calling convention
- #4: Proper error handling with early ABC validation
- #5: Maintained custom download_file result merging
- #6: Implemented exact dynamic approach suggested by reviewer

Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Ignatious Johnson <ichristo@amd.com>

---------

Signed-off-by: Ignatious Johnson <ichristo@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: cijohnson

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

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,40 +269,75 @@ 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,
             remote_file=remote_file,
             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()

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

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,