ttnn-visualizer 0.41.0__py3-none-any.whl → 0.42.0__py3-none-any.whl

ttnn_visualizer/csv_queries.py

@@ -4,6 +4,7 @@
 import csv
 import json
 import os
+import subprocess
 import tempfile
 from io import StringIO
 from pathlib import Path
@@ -13,10 +14,50 @@ import pandas as pd
 from tt_perf_report import perf_report
 
 from ttnn_visualizer.exceptions import DataFormatError
+from ttnn_visualizer.models import Instance, RemoteConnection
+from ttnn_visualizer.exceptions import SSHException, AuthenticationException, NoValidConnectionsError
 from ttnn_visualizer.models import Instance
-from ttnn_visualizer.ssh_client import get_client
 from ttnn_visualizer.sftp_operations import read_remote_file
 
+
+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
@@ -110,7 +151,38 @@ class RemoteCSVQueryRunner:
         self.remote_connection = remote_connection
         self.sep = sep
         self.offset = offset
-        self.ssh_client = get_client(remote_connection)
+
+    def _execute_ssh_command(self, command: str) -> str:
+        """Execute an SSH command and return the output."""
+        ssh_cmd = ["ssh"]
+        
+        # Handle non-standard SSH port
+        if self.remote_connection.port != 22:
+            ssh_cmd.extend(["-p", str(self.remote_connection.port)])
+        
+        ssh_cmd.extend([
+            f"{self.remote_connection.username}@{self.remote_connection.host}",
+            command
+        ])
+        
+        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, self.remote_connection)
+                # This line should never be reached as handle_ssh_subprocess_error raises an exception
+                raise RuntimeError(f"SSH command failed: {e.stderr}")
+            else:
+                raise RuntimeError(f"SSH command failed: {e.stderr}")
+        except subprocess.TimeoutExpired:
+            raise RuntimeError(f"SSH command timed out: {command}")
 
     def execute_query(
         self,
@@ -128,12 +200,7 @@ class RemoteCSVQueryRunner:
         """
         # Fetch header row, accounting for the offset
         header_cmd = f"head -n {self.offset + 1} {self.file_path} | tail -n 1"
-        stdin, stdout, stderr = self.ssh_client.exec_command(header_cmd)
-        raw_header = stdout.read().decode("utf-8").strip()
-        error = stderr.read().decode("utf-8").strip()
-
-        if error:
-            raise RuntimeError(f"Error fetching header row: {error}")
+        raw_header = self._execute_ssh_command(header_cmd).strip()
 
         # Sanitize headers
         headers = [
@@ -160,12 +227,7 @@ class RemoteCSVQueryRunner:
         limit_clause = f"| head -n {limit}" if limit else ""
         awk_cmd = f"awk -F'{self.sep}' 'NR > {self.offset + 1} {f'&& {awk_filter}' if awk_filter else ''} {{print}}' {self.file_path} {limit_clause}"
 
-        stdin, stdout, stderr = self.ssh_client.exec_command(awk_cmd)
-        output = stdout.read().decode("utf-8").strip()
-        error = stderr.read().decode("utf-8").strip()
-
-        if error:
-            raise RuntimeError(f"Error executing AWK command: {error}")
+        output = self._execute_ssh_command(awk_cmd).strip()
 
         # Split rows into lists of strings
         rows = [
@@ -205,12 +267,7 @@ class RemoteCSVQueryRunner:
             if total_lines
             else f"cat {self.file_path}"
         )
-        stdin, stdout, stderr = self.ssh_client.exec_command(cmd)
-        output = stdout.read().decode("utf-8").strip()
-        error = stderr.read().decode("utf-8").strip()
-
-        if error:
-            raise RuntimeError(f"Error fetching raw rows: {error}")
+        output = self._execute_ssh_command(cmd).strip()
 
         return output.splitlines()[self.offset:]
 
@@ -220,12 +277,7 @@ class RemoteCSVQueryRunner:
         :return: Dictionary of headers.
         """
         header_cmd = f"head -n {self.offset + 1} {self.file_path} | tail -n 1"
-        stdin, stdout, stderr = self.ssh_client.exec_command(header_cmd)
-        header = stdout.read().decode("utf-8").strip()
-        error = stderr.read().decode("utf-8").strip()
-
-        if error:
-            raise RuntimeError(f"Error reading CSV header: {error}")
+        header = self._execute_ssh_command(header_cmd).strip()
 
         # Trim spaces in header names
         column_names = [name.strip() for name in header.split(self.sep)]
@@ -254,10 +306,9 @@ class RemoteCSVQueryRunner:
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         """
-        Clean up the SSH connection when exiting context.
+        Clean up resources when exiting context.
         """
-        if self.ssh_client:
-            self.ssh_client.close()
+        pass
 
 
 class NPEQueries:
@@ -267,7 +318,6 @@ class NPEQueries:
     @staticmethod
     def get_npe_manifest(instance: Instance):
 
-
         if (
             not instance.remote_connection
             or instance.remote_connection
@@ -285,6 +335,31 @@ class NPEQueries:
                 f"{profiler_folder.remotePath}/{NPEQueries.NPE_FOLDER}/{NPEQueries.MANIFEST_FILE}",
             )
 
+    @staticmethod
+    def get_npe_timeline(instance: Instance, filename: str):
+        if not filename:
+            raise ValueError("filename parameter is required and cannot be None or empty")
+
+        if (
+            not instance.remote_connection
+            or not instance.remote_connection.useRemoteQuerying
+        ):
+            if not instance.performance_path:
+                raise ValueError("instance.performance_path is None")
+
+            file_path = Path(
+                instance.performance_path, NPEQueries.NPE_FOLDER, filename
+            )
+            with open(file_path, "r") as f:
+                return json.load(f)
+        else:
+            profiler_folder = instance.remote_profile_folder
+            return read_remote_file(
+                instance.remote_connection,
+                f"{profiler_folder.remotePath}/{NPEQueries.NPE_FOLDER}/{filename}",
+            )
+
+
 
 class DeviceLogProfilerQueries:
     DEVICE_LOG_FILE = "profile_log_device.csv"
@@ -611,6 +686,7 @@ class OpsPerformanceReportQueries:
         "inner_dim_block_size",
         "output_subblock_h",
         "output_subblock_w",
+        "global_call_count",
         "advice",
         "raw_op_code"
     ]

ttnn_visualizer/decorators.py

@@ -9,13 +9,10 @@ from ttnn_visualizer.enums import ConnectionTestStates
 
 from functools import wraps
 from flask import abort, request, session
-from paramiko.ssh_exception import (
+from ttnn_visualizer.exceptions import (
     AuthenticationException,
     NoValidConnectionsError,
     SSHException,
-)
-
-from ttnn_visualizer.exceptions import (
     RemoteConnectionException,
     NoProjectsException,
     RemoteSqliteException,
@@ -37,6 +34,12 @@ def with_instance(func):
             abort(404)
 
         instance_query_data = get_or_create_instance(instance_id=instance_id)
+        
+        # Handle case where get_or_create_instance returns None due to database error
+        if instance_query_data is None:
+            current_app.logger.error(f"Failed to get or create instance with ID: {instance_id}")
+            abort(500)
+            
         instance = instance_query_data.to_pydantic()
 
         kwargs["instance"] = instance

ttnn_visualizer/exceptions.py

@@ -43,5 +43,21 @@ class DataFormatError(Exception):
 class InvalidReportPath(Exception):
     pass
 
+
 class InvalidProfilerPath(Exception):
     pass
+
+
+class SSHException(Exception):
+    """Base SSH exception for subprocess SSH operations"""
+    pass
+
+
+class AuthenticationException(SSHException):
+    """Raised when SSH authentication fails"""
+    pass
+
+
+class NoValidConnectionsError(SSHException):
+    """Raised when SSH connection cannot be established"""
+    pass

ttnn_visualizer/queries.py

@@ -2,7 +2,6 @@
 #
 # SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
 
-import json
 from typing import Generator, Dict, Any, Union
 
 from ttnn_visualizer.exceptions import (
@@ -23,11 +22,9 @@ from ttnn_visualizer.models import (
     ProducersConsumers,
     TensorComparisonRecord,
 )
-from ttnn_visualizer.ssh_client import get_client
 import sqlite3
 from typing import List, Optional
 from pathlib import Path
-import paramiko
 
 
 class LocalQueryRunner:
@@ -63,111 +60,10 @@ class LocalQueryRunner:
             self.connection.close()
 
 
-class RemoteQueryRunner:
-    column_delimiter = "|||"
-
-    def __init__(self, instance: Instance):
-        self.instance = instance
-        self._validate_instance()
-        self.ssh_client = self._get_ssh_client(self.instance.remote_connection)
-        self.sqlite_binary = self.instance.remote_connection.sqliteBinaryPath
-        self.remote_db_path = str(
-            Path(self.instance.remote_profiler_folder.remotePath, "db.sqlite")
-        )
-
-    def _validate_instance(self):
-        """
-        Validate that the instance has all required remote connection attributes.
-        """
-        if (
-            not self.instance.remote_connection
-            or not self.instance.remote_connection.sqliteBinaryPath
-            or not self.instance.remote_profiler_folder
-            or not self.instance.remote_profiler_folder.remotePath
-        ):
-            raise ValueError(
-                "Remote connections require remote path and sqliteBinaryPath"
-            )
-
-    def _get_ssh_client(self, remote_connection) -> paramiko.SSHClient:
-        """
-        Retrieve the SSH client for the given remote connection.
-        """
-        return get_client(remote_connection=remote_connection)
-
-    def _format_query(self, query: str, params: Optional[List] = None) -> str:
-        """
-        Format the query by replacing placeholders with properly quoted parameters.
-        """
-        if not params:
-            return query
-
-        formatted_params = [
-            f"'{param}'" if isinstance(param, str) else str(param) for param in params
-        ]
-        return query.replace("?", "{}").format(*formatted_params)
-
-    def _build_command(self, formatted_query: str) -> str:
-        """
-        Build the remote SQLite command.
-        """
-        return f'{self.sqlite_binary} {self.remote_db_path} "{formatted_query}" -json'
-
-    def _execute_ssh_command(self, command: str) -> tuple:
-        """
-        Execute the SSH command and return the standard output and error.
-        """
-        stdin, stdout, stderr = self.ssh_client.exec_command(command)
-        output = stdout.read().decode("utf-8").strip()
-        error_output = stderr.read().decode("utf-8").strip()
-        return output, error_output
-
-    def _parse_output(self, output: str, command: str) -> List:
-        """
-        Parse the output from the SQLite command. Attempt JSON parsing first,
-        then fall back to line-based parsing.
-        """
-        if not output.strip():
-            return []
-
-        try:
-            rows = json.loads(output)
-            return [tuple(row.values()) for row in rows]
-        except json.JSONDecodeError:
-            print(
-                f"Output is not valid JSON, attempting manual parsing.\nCommand: {command}"
-            )
-            return [tuple(line.split("|")) for line in output.splitlines()]
-
-    def execute_query(self, query: str, params: Optional[List] = None) -> List:
-        """
-        Execute a remote SQLite query using the instance's SSH client.
-        """
-        self._validate_instance()
-        formatted_query = self._format_query(query, params)
-        command = self._build_command(formatted_query)
-        output, error_output = self._execute_ssh_command(command)
-
-        if error_output:
-            raise RuntimeError(
-                f"Error executing query remotely: {error_output}\nCommand: {command}"
-            )
-
-        return self._parse_output(output, command)
-
-    def close(self):
-        """
-        Close the SSH connection.
-        """
-        if self.ssh_client:
-            self.ssh_client.close()
-
-
 class DatabaseQueries:
 
     instance: Optional[Instance] = None
-    ssh_client = None
-    query_runner: LocalQueryRunner | RemoteQueryRunner
+    query_runner: LocalQueryRunner
 
     def __init__(self, instance: Optional[Instance] = None, connection=None):
         self.instance = instance
@@ -181,7 +77,7 @@ class DatabaseQueries:
                 )
             remote_connection = instance.remote_connection if instance else None
             if remote_connection and remote_connection.useRemoteQuerying:
-                self.query_runner = RemoteQueryRunner(instance=instance)
+                raise NotImplementedError("Remote querying is not implemented yet")
             else:
                 self.query_runner = LocalQueryRunner(instance=instance)
 
@@ -382,7 +278,5 @@ class DatabaseQueries:
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        if isinstance(self.query_runner, RemoteQueryRunner):
-            self.query_runner.close()
-        elif isinstance(self.query_runner, LocalQueryRunner):
+        if isinstance(self.query_runner, LocalQueryRunner):
             self.query_runner.close()

ttnn_visualizer/remote_sqlite_setup.py

@@ -3,29 +3,109 @@
 # 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 RemoteSqliteException
+from ttnn_visualizer.exceptions import RemoteSqliteException, SSHException, AuthenticationException, NoValidConnectionsError
 from ttnn_visualizer.models import RemoteConnection
-from ttnn_visualizer.ssh_client import get_client
+
+
+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}")
 
 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"]
+    
+    # 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
+    ])
+    
+    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:
+        raise RemoteSqliteException(
+            message=f"SSH command timed out: {command}",
+            status=ConnectionTestStates.FAILED,
+        )
+
+
 def find_sqlite_binary(connection):
     """Check if SQLite is installed on the remote machine and return its path."""
-    ssh_client = get_client(connection)
     try:
-        stdin, stdout, stderr = ssh_client.exec_command("which sqlite3")
-        binary_path = stdout.read().decode().strip()
-        error = stderr.read().decode().strip()
+        output = _execute_ssh_command(connection, "which sqlite3")
+        binary_path = output.strip()
         if binary_path:
             print(f"SQLite binary found at: {binary_path}")
             return binary_path
-        elif error:
-            print(f"Error checking SQLite binary: {error}")
         return None
+    except RemoteSqliteException:
+        # Re-raise RemoteSqliteException as-is
+        raise
     except Exception as e:
         raise RemoteSqliteException(
             message=f"Error finding SQLite binary: {str(e)}",
@@ -33,17 +113,13 @@ def find_sqlite_binary(connection):
         )
 
 
-def is_sqlite_executable(ssh_client, binary_path):
+def is_sqlite_executable(remote_connection: RemoteConnection, binary_path):
     """Check if the SQLite binary is executable by trying to run it."""
     try:
-        stdin, stdout, stderr = ssh_client.exec_command(f"{binary_path} --version")
-        output = stdout.read().decode().strip()
-        error = stderr.read().decode().strip()
-        stdout.channel.recv_exit_status()
-        if error:
-            raise Exception(f"Error while trying to run SQLite binary: {error}")
-
-        version = get_sqlite_version(output)
+        output = _execute_ssh_command(remote_connection, f"{binary_path} --version")
+        version_output = output.strip()
+        
+        version = get_sqlite_version(version_output)
         if not is_version_at_least(version, MINIMUM_SQLITE_VERSION):
             raise Exception(
                 f"SQLite version {version} is below the required minimum of {MINIMUM_SQLITE_VERSION}."
@@ -52,6 +128,9 @@ def is_sqlite_executable(ssh_client, binary_path):
         print(f"SQLite binary at {binary_path} is executable. Version: {version}")
         return True
 
+    except RemoteSqliteException:
+        # Re-raise RemoteSqliteException as-is
+        raise
     except Exception as e:
         raise Exception(f"Error checking SQLite executability: {str(e)}")
 
@@ -76,8 +155,7 @@ def is_version_at_least(version, minimum_version):
 @remote_exception_handler
 def check_sqlite_path(remote_connection: RemoteConnection):
     try:
-        client = get_client(remote_connection)
-        is_sqlite_executable(client, remote_connection.sqliteBinaryPath)
+        is_sqlite_executable(remote_connection, remote_connection.sqliteBinaryPath)
     except Exception as e:
         raise RemoteSqliteException(message=str(e), status=ConnectionTestStates.FAILED)
 

ttnn_visualizer/requirements.txt

@@ -1,7 +1,6 @@
 Flask==3.1.1
 gunicorn~=22.0.0
 uvicorn==0.30.1
-paramiko~=3.4.0
 flask_cors==4.0.1
 pydantic==2.7.3
 pydantic_core==2.18.4
@@ -16,9 +15,8 @@ wheel
 build
 PyYAML==6.0.2
 python-dotenv==1.0.1
-tt-perf-report==1.0.6
+tt-perf-report==1.0.7
 zstd==1.5.7.0
 
 # Dev dependencies
 mypy
-types-paramiko