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

ttnn_visualizer/models.py

@@ -119,6 +119,7 @@ class Tensor(SerializeableDataclass):
     def __post_init__(self):
         self.memory_config = parse_memory_config(self.memory_config)
 
+
 @dataclasses.dataclass
 class InputTensor(SerializeableDataclass):
     operation_id: int
@@ -246,19 +247,21 @@ class InstanceTable(db.Model):
             "remote_performance_folder": self.remote_performance_folder,
             "profiler_path": self.profiler_path,
             "performance_path": self.performance_path,
-            "npe_path": self.npe_path
+            "npe_path": self.npe_path,
         }
 
     def to_pydantic(self) -> Instance:
         return Instance(
             instance_id=str(self.instance_id),
-            profiler_path=str(self.profiler_path) if self.profiler_path is not None else None,
-            performance_path=(
-                str(self.performance_path) if self.performance_path is not None else None
+            profiler_path=(
+                str(self.profiler_path) if self.profiler_path is not None else None
             ),
-            npe_path=(
-                str(self.npe_path) if self.npe_path is not None else None
+            performance_path=(
+                str(self.performance_path)
+                if self.performance_path is not None
+                else None
             ),
+            npe_path=(str(self.npe_path) if self.npe_path is not None else None),
             active_report=(
                 (ActiveReports(**self.active_report) if self.active_report else None)
                 if isinstance(self.active_report, dict)
@@ -270,7 +273,9 @@ class InstanceTable(db.Model):
                 else None
             ),
             remote_profiler_folder=(
-                RemoteReportFolder.model_validate(self.remote_profiler_folder, strict=False)
+                RemoteReportFolder.model_validate(
+                    self.remote_profiler_folder, strict=False
+                )
                 if self.remote_profiler_folder is not None
                 else None
             ),

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,116 @@
 # 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", "-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])
+
+    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 +120,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 +135,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 +162,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,9 @@ 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
+black==25.1.0

ttnn_visualizer/serializers.py

@@ -8,6 +8,7 @@ from typing import List
 
 from ttnn_visualizer.models import BufferType, Operation, TensorComparisonRecord
 
+
 def serialize_operations(
     inputs,
     operation_arguments,

ttnn_visualizer/settings.py

@@ -10,6 +10,7 @@ from ttnn_visualizer.utils import str_to_bool
 
 load_dotenv()
 
+
 class DefaultConfig(object):
     # General Settings
     SECRET_KEY = os.getenv("SECRET_KEY", "90909")
@@ -19,17 +20,20 @@ class DefaultConfig(object):
     SERVER_MODE = str_to_bool(os.getenv("SERVER_MODE", "false"))
     MALWARE_SCANNER = os.getenv("MALWARE_SCANNER")
     ALLOWED_ORIGINS = [
-        o for o in os.getenv("ALLOWED_ORIGINS", "http://localhost:5173,http://localhost:8000").split(",")
+        o
+        for o in os.getenv(
+            "ALLOWED_ORIGINS", "http://localhost:5173,http://localhost:8000"
+        ).split(",")
         if o
     ]
     BASE_PATH = os.getenv("BASE_PATH", "/")
-    MAX_CONTENT_LENGTH = (
-        None if not (v := os.getenv("MAX_CONTENT_LENGTH")) else int(v)
-    )
+    MAX_CONTENT_LENGTH = None if not (v := os.getenv("MAX_CONTENT_LENGTH")) else int(v)
 
     # Path Settings
     DB_VERSION = "0.29.0"  # App version when DB schema last changed
-    REPORT_DATA_DIRECTORY = os.getenv("REPORT_DATA_DIRECTORY", Path(__file__).parent.absolute().joinpath("data"))
+    REPORT_DATA_DIRECTORY = os.getenv(
+        "REPORT_DATA_DIRECTORY", Path(__file__).parent.absolute().joinpath("data")
+    )
     LOCAL_DATA_DIRECTORY = Path(REPORT_DATA_DIRECTORY).joinpath("local")
     REMOTE_DATA_DIRECTORY = Path(REPORT_DATA_DIRECTORY).joinpath("remote")
     PROFILER_DIRECTORY_NAME = "profiler-reports"