ttnn-visualizer 0.44.0__py3-none-any.whl → 0.45.0__py3-none-any.whl

ttnn_visualizer/csv_queries.py

@@ -4,7 +4,6 @@
 import csv
 import json
 import os
-import subprocess
 import tempfile
 from io import StringIO
 from pathlib import Path
@@ -13,62 +12,10 @@ from typing import Dict, List, Optional, Union
 import pandas as pd
 import zstd
 from tt_perf_report import perf_report
-from ttnn_visualizer.exceptions import (
-    AuthenticationException,
-    DataFormatError,
-    NoValidConnectionsError,
-    SSHException,
-)
+from ttnn_visualizer.exceptions import DataFormatError
 from ttnn_visualizer.models import Instance, RemoteConnection
 
 
-def handle_ssh_subprocess_error(
-    e: subprocess.CalledProcessError, remote_connection: RemoteConnection
-):
-    """
-    Convert subprocess SSH errors to appropriate SSH exceptions.
-
-    :param e: The subprocess.CalledProcessError
-    :param remote_connection: The RemoteConnection object for context
-    :raises: SSHException, AuthenticationException, or NoValidConnectionsError
-    """
-    stderr = e.stderr.lower() if e.stderr else ""
-
-    # Check for authentication failures
-    if any(
-        auth_err in stderr
-        for auth_err in [
-            "permission denied",
-            "authentication failed",
-            "publickey",
-            "password",
-            "host key verification failed",
-        ]
-    ):
-        raise AuthenticationException(f"SSH authentication failed: {e.stderr}")
-
-    # Check for connection failures
-    elif any(
-        conn_err in stderr
-        for conn_err in [
-            "connection refused",
-            "network is unreachable",
-            "no route to host",
-            "name or service not known",
-            "connection timed out",
-        ]
-    ):
-        raise NoValidConnectionsError(f"SSH connection failed: {e.stderr}")
-
-    # Check for general SSH protocol errors
-    elif "ssh:" in stderr or "protocol" in stderr:
-        raise SSHException(f"SSH protocol error: {e.stderr}")
-
-    # Default to generic SSH exception
-    else:
-        raise SSHException(f"SSH command failed: {e.stderr}")
-
-
 class LocalCSVQueryRunner:
     def __init__(self, file_path: str, offset: int = 0):
         self.file_path = file_path

ttnn_visualizer/exceptions.py

@@ -14,10 +14,12 @@ class RemoteConnectionException(Exception):
         message,
         status: ConnectionTestStates,
         http_status_code: Optional[HTTPStatus] = None,
+        detail: Optional[str] = None,
     ):
         super().__init__(message)
         self.message = message
         self.status = status
+        self.detail = detail
         self._http_status_code = http_status_code
 
     @property
@@ -37,12 +39,16 @@ class AuthenticationFailedException(RemoteConnectionException):
     """Exception for SSH authentication failures that should return HTTP 422"""
 
     def __init__(
-        self, message, status: ConnectionTestStates = ConnectionTestStates.FAILED
+        self,
+        message,
+        status: ConnectionTestStates = ConnectionTestStates.FAILED,
+        detail: Optional[str] = None,
     ):
         super().__init__(
             message=message,
             status=status,
             http_status_code=HTTPStatus.UNPROCESSABLE_ENTITY,  # 422
+            detail=detail,
         )
 
 

ttnn_visualizer/models.py

@@ -174,6 +174,7 @@ class RemoteConnection(SerializeableModel):
 class StatusMessage(SerializeableModel):
     status: ConnectionTestStates
     message: str
+    detail: Optional[str] = None
 
 
 class ActiveReports(SerializeableModel):

ttnn_visualizer/remote_sqlite_setup.py

@@ -3,100 +3,24 @@
 # SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
 
 import re
-import subprocess
 
 from ttnn_visualizer.decorators import remote_exception_handler
 from ttnn_visualizer.enums import ConnectionTestStates
-from ttnn_visualizer.exceptions import (
-    AuthenticationException,
-    NoValidConnectionsError,
-    RemoteSqliteException,
-    SSHException,
-)
+from ttnn_visualizer.exceptions import RemoteSqliteException, SSHException
 from ttnn_visualizer.models import RemoteConnection
-
-
-def handle_ssh_subprocess_error(
-    e: subprocess.CalledProcessError, remote_connection: RemoteConnection
-):
-    """
-    Convert subprocess SSH errors to appropriate SSH exceptions.
-
-    :param e: The subprocess.CalledProcessError
-    :param remote_connection: The RemoteConnection object for context
-    :raises: SSHException, AuthenticationException, or NoValidConnectionsError
-    """
-    stderr = e.stderr.lower() if e.stderr else ""
-
-    # Check for authentication failures
-    if any(
-        auth_err in stderr
-        for auth_err in [
-            "permission denied",
-            "authentication failed",
-            "publickey",
-            "password",
-            "host key verification failed",
-        ]
-    ):
-        raise AuthenticationException(f"SSH authentication failed: {e.stderr}")
-
-    # Check for connection failures
-    elif any(
-        conn_err in stderr
-        for conn_err in [
-            "connection refused",
-            "network is unreachable",
-            "no route to host",
-            "name or service not known",
-            "connection timed out",
-        ]
-    ):
-        raise NoValidConnectionsError(f"SSH connection failed: {e.stderr}")
-
-    # Check for general SSH protocol errors
-    elif "ssh:" in stderr or "protocol" in stderr:
-        raise SSHException(f"SSH protocol error: {e.stderr}")
-
-    # Default to generic SSH exception
-    else:
-        raise SSHException(f"SSH command failed: {e.stderr}")
-
+from ttnn_visualizer.ssh_client import SSHClient
 
 MINIMUM_SQLITE_VERSION = "3.38.0"
 
 
 def _execute_ssh_command(remote_connection: RemoteConnection, command: str) -> str:
     """Execute an SSH command and return the output."""
-    ssh_cmd = ["ssh", "-o", "PasswordAuthentication=no"]
-
-    # Handle non-standard SSH port
-    if remote_connection.port != 22:
-        ssh_cmd.extend(["-p", str(remote_connection.port)])
-
-    ssh_cmd.extend([f"{remote_connection.username}@{remote_connection.host}", command])
-
+    ssh_client = SSHClient(remote_connection)
     try:
-        result = subprocess.run(
-            ssh_cmd, capture_output=True, text=True, check=True, timeout=30
-        )
-        return result.stdout
-    except subprocess.CalledProcessError as e:
-        if e.returncode == 255:  # SSH protocol errors
-            handle_ssh_subprocess_error(e, remote_connection)
-            # This line should never be reached as handle_ssh_subprocess_error raises an exception
-            raise RemoteSqliteException(
-                message=f"SSH command failed: {e.stderr}",
-                status=ConnectionTestStates.FAILED,
-            )
-        else:
-            raise RemoteSqliteException(
-                message=f"SSH command failed: {e.stderr}",
-                status=ConnectionTestStates.FAILED,
-            )
-    except subprocess.TimeoutExpired:
+        return ssh_client.execute_command(command, timeout=30)
+    except SSHException as e:
         raise RemoteSqliteException(
-            message=f"SSH command timed out: {command}",
+            message=str(e),
             status=ConnectionTestStates.FAILED,
         )
 

ttnn_visualizer/sftp_operations.py

@@ -24,6 +24,7 @@ from ttnn_visualizer.exceptions import (
 )
 from ttnn_visualizer.models import RemoteConnection, RemoteReportFolder
 from ttnn_visualizer.sockets import FileProgress, FileStatus, emit_file_status
+from ttnn_visualizer.ssh_client import SSHClient
 from ttnn_visualizer.utils import update_last_synced
 
 logger = logging.getLogger(__name__)
@@ -33,53 +34,6 @@ TEST_PROFILER_FILE = "profile_log_device.csv"
 REPORT_DATA_DIRECTORY = Path(__file__).parent.absolute().joinpath("data")
 
 
-def handle_ssh_subprocess_error(
-    e: subprocess.CalledProcessError, remote_connection: RemoteConnection
-):
-    """
-    Convert subprocess SSH errors to appropriate SSH exceptions.
-
-    :param e: The subprocess.CalledProcessError
-    :param remote_connection: The RemoteConnection object for context
-    :raises: SSHException, AuthenticationException, or NoValidConnectionsError
-    """
-    stderr = e.stderr.lower() if e.stderr else ""
-
-    # Check for authentication failures
-    if any(
-        auth_err in stderr
-        for auth_err in [
-            "permission denied",
-            "authentication failed",
-            "publickey",
-            "password",
-            "host key verification failed",
-        ]
-    ):
-        raise AuthenticationException(f"SSH authentication failed: {e.stderr}")
-
-    # Check for connection failures
-    elif any(
-        conn_err in stderr
-        for conn_err in [
-            "connection refused",
-            "network is unreachable",
-            "no route to host",
-            "name or service not known",
-            "connection timed out",
-        ]
-    ):
-        raise NoValidConnectionsError(f"SSH connection failed: {e.stderr}")
-
-    # Check for general SSH protocol errors
-    elif "ssh:" in stderr or "protocol" in stderr:
-        raise SSHException(f"SSH protocol error: {e.stderr}")
-
-    # Default to generic SSH exception
-    else:
-        raise SSHException(f"SSH command failed: {e.stderr}")
-
-
 def start_background_task(task, *args):
     with current_app.app_context():
         if current_app.config["USE_WEBSOCKETS"]:
@@ -663,36 +617,8 @@ def read_remote_file(
     else:
         path = Path(remote_connection.profilerPath)
 
-    logger.info(f"Reading remote file {path}")
-
-    # Build SSH command to read the file
-    ssh_cmd = ["ssh", "-o", "PasswordAuthentication=no"]
-
-    # Handle non-standard SSH port
-    if remote_connection.port != 22:
-        ssh_cmd.extend(["-p", str(remote_connection.port)])
-
-    ssh_cmd.extend(
-        [f"{remote_connection.username}@{remote_connection.host}", f"cat '{path}'"]
-    )
-
-    try:
-        result = subprocess.run(ssh_cmd, capture_output=True, check=True, timeout=30)
-        return result.stdout
-    except subprocess.CalledProcessError as e:
-        if e.returncode == 255:  # SSH protocol errors
-            handle_ssh_subprocess_error(e, remote_connection)
-            return None
-        else:
-            # File not found or other command error
-            logger.error(f"File not found or cannot be read: {path}")
-            return None
-    except subprocess.TimeoutExpired:
-        logger.error(f"Timeout reading remote file: {path}")
-        return None
-    except Exception as e:
-        logger.error(f"Error reading remote file {path}: {e}")
-        return None
+    ssh_client = SSHClient(remote_connection)
+    return ssh_client.read_file(path, timeout=30)
 
 
 @remote_exception_handler
@@ -713,24 +639,11 @@ def check_remote_path_exists(remote_connection: RemoteConnection, path_key: str)
     """Check if a remote path exists using SSH test command."""
     path = getattr(remote_connection, path_key)
 
-    # Build SSH command to test if path exists
-    ssh_cmd = ["ssh", "-o", "PasswordAuthentication=no"]
-
-    # Handle non-standard SSH port
-    if remote_connection.port != 22:
-        ssh_cmd.extend(["-p", str(remote_connection.port)])
-
-    ssh_cmd.extend(
-        [f"{remote_connection.username}@{remote_connection.host}", f"test -d '{path}'"]
-    )
+    ssh_client = SSHClient(remote_connection)
 
     try:
-        result = subprocess.run(ssh_cmd, capture_output=True, check=True, timeout=10)
-        # If command succeeds, directory exists
-        return True
-    except subprocess.CalledProcessError as e:
-        if e.returncode == 255:  # SSH protocol errors
-            handle_ssh_subprocess_error(e, remote_connection)
+        if ssh_client.check_path_exists(path, timeout=10):
+            return True
         else:
             # Directory does not exist or is inaccessible
             if path_key == "performancePath":
@@ -742,10 +655,10 @@ def check_remote_path_exists(remote_connection: RemoteConnection, path_key: str)
             raise RemoteConnectionException(
                 message=message, status=ConnectionTestStates.FAILED
             )
-    except subprocess.TimeoutExpired:
-        logger.error(f"Timeout checking remote path: {path}")
+    except SSHException as e:
+        logger.error(f"Error checking remote path: {path}")
         raise RemoteConnectionException(
-            message=f"Timeout checking remote path: {path}",
+            message=f"Error checking remote path: {path}: {str(e)}",
             status=ConnectionTestStates.FAILED,
         )
 
@@ -902,11 +815,11 @@ def sync_remote_profiler_folders(
 def sync_remote_performance_folders(
     remote_connection: RemoteConnection,
     path_prefix: str,
-    profile: RemoteReportFolder,
+    performance: RemoteReportFolder,
     exclude_patterns: Optional[List[str]] = None,
     sid=None,
 ):
-    remote_folder_path = profile.remotePath
+    remote_folder_path = performance.remotePath
     profile_folder = Path(remote_folder_path).name
     destination_dir = Path(
         REPORT_DATA_DIRECTORY,