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

ttnn_visualizer/__init__.py

@@ -1,4 +1,3 @@
 # SPDX-License-Identifier: Apache-2.0
 #
 # SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
-

ttnn_visualizer/app.py

@@ -21,7 +21,11 @@ from flask_cors import CORS
 from werkzeug.debug import DebuggedApplication
 from werkzeug.middleware.proxy_fix import ProxyFix
 
-from ttnn_visualizer.exceptions import DatabaseFileNotFoundException, InvalidProfilerPath, InvalidReportPath
+from ttnn_visualizer.exceptions import (
+    DatabaseFileNotFoundException,
+    InvalidProfilerPath,
+    InvalidReportPath,
+)
 from ttnn_visualizer.instances import create_instance_from_local_paths
 from ttnn_visualizer.settings import Config, DefaultConfig
 
@@ -65,6 +69,7 @@ def create_app(settings_override=None):
     extensions(app)
 
     if flask_env == "production":
+
         @app.route(f"{app.config['BASE_PATH']}", defaults={"path": ""})
         @app.route(f"{app.config['BASE_PATH']}<path:path>")
         def catch_all(path):
@@ -179,9 +184,15 @@ def open_browser(host, port, instance_id=None):
 
 
 def parse_args():
-    parser = argparse.ArgumentParser(description="A tool for visualizing the Tenstorrent Neural Network model (TT-NN)")
-    parser.add_argument("--profiler-path", type=str, help="Specify a profiler path", default=None)
-    parser.add_argument("--performance-path", help="Specify a performance path", default=None)
+    parser = argparse.ArgumentParser(
+        description="A tool for visualizing the Tenstorrent Neural Network model (TT-NN)"
+    )
+    parser.add_argument(
+        "--profiler-path", type=str, help="Specify a profiler path", default=None
+    )
+    parser.add_argument(
+        "--performance-path", help="Specify a performance path", default=None
+    )
     return parser.parse_args()
 
 

ttnn_visualizer/csv_queries.py

@@ -4,19 +4,73 @@
 import csv
 import json
 import os
+import subprocess
 import tempfile
 from io import StringIO
 from pathlib import Path
 from typing import List, Dict, Union, Optional
 
 import pandas as pd
+import zstd
 from tt_perf_report import perf_report
 
 from ttnn_visualizer.exceptions import DataFormatError
-from ttnn_visualizer.models import Instance
-from ttnn_visualizer.ssh_client import get_client
+from ttnn_visualizer.exceptions import (
+    SSHException,
+    AuthenticationException,
+    NoValidConnectionsError,
+)
+from ttnn_visualizer.models import Instance, RemoteConnection
 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 +164,36 @@ 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", "-o", "PasswordAuthentication=no"]
+
+        # 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 +211,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 +238,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,14 +278,9 @@ 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:]
+        return output.splitlines()[self.offset :]
 
     def get_csv_header(self) -> Dict[str, int]:
         """
@@ -220,12 +288,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 +317,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,14 +329,15 @@ class NPEQueries:
     @staticmethod
     def get_npe_manifest(instance: Instance):
 
-
         if (
             not instance.remote_connection
             or instance.remote_connection
             and not instance.remote_connection.useRemoteQuerying
         ):
             file_path = Path(
-                instance.performance_path, NPEQueries.NPE_FOLDER, NPEQueries.MANIFEST_FILE
+                instance.performance_path,
+                NPEQueries.NPE_FOLDER,
+                NPEQueries.MANIFEST_FILE,
             )
             with open(file_path, "r") as f:
                 return json.load(f)
@@ -285,6 +348,48 @@ 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)
+
+            if filename.endswith(".zst"):
+                with open(file_path, "rb") as file:
+                    compressed_data = file.read()
+                    uncompressed_data = zstd.uncompress(compressed_data)
+                    return json.loads(uncompressed_data)
+            else:
+                with open(file_path, "r") as f:
+                    return json.load(f)
+
+        else:
+            profiler_folder = instance.remote_profile_folder
+            remote_path = (
+                f"{profiler_folder.remotePath}/{NPEQueries.NPE_FOLDER}/{filename}"
+            )
+            remote_data = read_remote_file(instance.remote_connection, remote_path)
+
+            if filename.endswith(".zst"):
+                if isinstance(remote_data, str):
+                    remote_data = remote_data.encode("utf-8")
+                uncompressed_data = zstd.decompress(remote_data)
+                return json.loads(uncompressed_data)
+            else:
+                if isinstance(remote_data, bytes):
+                    remote_data = remote_data.decode("utf-8")
+                return json.loads(remote_data)
+
 
 class DeviceLogProfilerQueries:
     DEVICE_LOG_FILE = "profile_log_device.csv"
@@ -539,7 +644,9 @@ class OpsPerformanceQueries:
             or instance.remote_connection
             and not instance.remote_connection.useRemoteQuerying
         ):
-            with open(OpsPerformanceQueries.get_local_ops_perf_file_path(instance)) as f:
+            with open(
+                OpsPerformanceQueries.get_local_ops_perf_file_path(instance)
+            ) as f:
                 return f.read()
         else:
             path = OpsPerformanceQueries.get_remote_ops_perf_file_path(instance)
@@ -581,9 +688,7 @@ class OpsPerformanceQueries:
         """
         try:
             return [
-                folder.name
-                for folder in Path(directory).iterdir()
-                if folder.is_dir()
+                folder.name for folder in Path(directory).iterdir() if folder.is_dir()
             ]
         except Exception as e:
             raise RuntimeError(f"Error accessing directory: {e}")
@@ -611,8 +716,9 @@ class OpsPerformanceReportQueries:
         "inner_dim_block_size",
         "output_subblock_h",
         "output_subblock_w",
+        "global_call_count",
         "advice",
-        "raw_op_code"
+        "raw_op_code",
     ]
 
     PASSTHROUGH_COLUMNS = {
@@ -658,7 +764,9 @@ class OpsPerformanceReportQueries:
                 next(reader, None)
                 for row in reader:
                     processed_row = {
-                        column: row[index] for index, column in enumerate(cls.REPORT_COLUMNS) if index < len(row)
+                        column: row[index]
+                        for index, column in enumerate(cls.REPORT_COLUMNS)
+                        if index < len(row)
                     }
                     if "advice" in processed_row and processed_row["advice"]:
                         processed_row["advice"] = processed_row["advice"].split(" • ")
@@ -667,7 +775,9 @@ class OpsPerformanceReportQueries:
 
                     for key, value in cls.PASSTHROUGH_COLUMNS.items():
                         op_id = int(row[0])
-                        idx = op_id - 2  # IDs in result column one correspond to row numbers in ops perf results csv
+                        idx = (
+                            op_id - 2
+                        )  # IDs in result column one correspond to row numbers in ops perf results csv
                         processed_row[key] = ops_perf_results[idx][value]
 
                     report.append(processed_row)

ttnn_visualizer/decorators.py

@@ -9,14 +9,12 @@ 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,
+    AuthenticationFailedException,
     NoProjectsException,
     RemoteSqliteException,
 )
@@ -37,15 +35,23 @@ 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
 
-        if 'instances' not in session:
-            session['instances'] = []
+        if "instances" not in session:
+            session["instances"] = []
 
-        if instance.instance_id not in session['instances']:
-            session['instances'] = session.get('instances', []) + [instance.instance_id]
+        if instance.instance_id not in session["instances"]:
+            session["instances"] = session.get("instances", []) + [instance.instance_id]
 
         return func(*args, **kwargs)
 
@@ -65,10 +71,21 @@ def remote_exception_handler(func):
         try:
             return func(*args, **kwargs)
         except AuthenticationException as err:
-            current_app.logger.error(f"Authentication failed {err}")
-            raise RemoteConnectionException(
+            # Log the detailed error for debugging, but don't show full traceback
+            current_app.logger.warning(
+                f"SSH authentication failed for {connection.username}@{connection.host}: SSH key authentication required"
+            )
+
+            # Return user-friendly error message about SSH keys
+            user_message = (
+                "SSH authentication failed. This application requires SSH key-based authentication. "
+                "Please ensure your SSH public key is added to the authorized_keys file on the remote server. "
+                "Password authentication is not supported."
+            )
+
+            raise AuthenticationFailedException(
+                message=user_message,
                 status=ConnectionTestStates.FAILED,
-                message=f"Unable to authenticate: {str(err)}",
             )
         except FileNotFoundError as err:
             current_app.logger.error(f"File not found: {str(err)}")
@@ -83,11 +100,20 @@ def remote_exception_handler(func):
                 message=f"No projects found at remote location: {connection.path}",
             )
         except NoValidConnectionsError as err:
-            current_app.logger.error(f"No valid connections: {str(err)}")
-            message = re.sub(r"\[.*?]", "", str(err)).strip()
+            current_app.logger.warning(
+                f"SSH connection failed for {connection.username}@{connection.host}: {str(err)}"
+            )
+
+            # Provide user-friendly message for connection issues
+            user_message = (
+                f"Unable to establish SSH connection to {connection.host}. "
+                "Please check the hostname, port, and network connectivity. "
+                "Ensure SSH key-based authentication is properly configured."
+            )
+
             raise RemoteConnectionException(
                 status=ConnectionTestStates.FAILED,
-                message=f"{message}",
+                message=user_message,
             )
 
         except RemoteSqliteException as err:
@@ -108,10 +134,10 @@ def remote_exception_handler(func):
             )
         except SSHException as err:
             if str(err) == "No existing session":
-                message = "Authentication failed - check credentials and ssh-agent"
+                message = "SSH authentication failed. Please ensure SSH keys are configured and ssh-agent is running."
             else:
                 err_message = re.sub(r"\[.*?]", "", str(err)).strip()
-                message = f"Error connecting to host {connection.host}: {err_message}"
+                message = f"SSH connection error to {connection.host}: {err_message}. Ensure SSH key-based authentication is properly configured."
 
             raise RemoteConnectionException(
                 status=ConnectionTestStates.FAILED, message=message

ttnn_visualizer/exceptions.py

@@ -3,24 +3,49 @@
 # SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
 
 from http import HTTPStatus
+from typing import Optional
 
 from ttnn_visualizer.enums import ConnectionTestStates
 
 
 class RemoteConnectionException(Exception):
-    def __init__(self, message, status: ConnectionTestStates):
+    def __init__(
+        self,
+        message,
+        status: ConnectionTestStates,
+        http_status_code: Optional[HTTPStatus] = None,
+    ):
         super().__init__(message)
         self.message = message
         self.status = status
+        self._http_status_code = http_status_code
 
     @property
     def http_status(self):
+        # Use custom HTTP status code if provided
+        if self._http_status_code is not None:
+            return self._http_status_code
+
+        # Default behavior
         if self.status == ConnectionTestStates.FAILED:
             return HTTPStatus.INTERNAL_SERVER_ERROR
         if self.status == ConnectionTestStates.OK:
             return HTTPStatus.OK
 
 
+class AuthenticationFailedException(RemoteConnectionException):
+    """Exception for SSH authentication failures that should return HTTP 422"""
+
+    def __init__(
+        self, message, status: ConnectionTestStates = ConnectionTestStates.FAILED
+    ):
+        super().__init__(
+            message=message,
+            status=status,
+            http_status_code=HTTPStatus.UNPROCESSABLE_ENTITY,  # 422
+        )
+
+
 class NoProjectsException(RemoteConnectionException):
     pass
 
@@ -43,5 +68,24 @@ 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/file_uploads.py

@@ -54,6 +54,7 @@ def extract_folder_name_from_files(files):
     unsplit_name = str(files[0].filename)
     return unsplit_name.split("/")[0]
 
+
 def extract_npe_name(files):
     if not files:
         return None

ttnn_visualizer/instances.py

@@ -94,8 +94,7 @@ def update_existing_instance(
     else:
         if active_report.get("npe_name"):
             instance_data.npe_path = get_npe_path(
-                npe_name=active_report["npe_name"],
-                current_app=current_app
+                npe_name=active_report["npe_name"], current_app=current_app
             )
 
 
@@ -148,17 +147,25 @@ def create_new_instance(
     instance_data = InstanceTable(
         instance_id=instance_id,
         active_report=active_report,
-        profiler_path=profiler_path if profiler_path is not _sentinel else get_profiler_path(
-            active_report["profiler_name"],
-            current_app=current_app,
-            remote_connection=remote_connection,
+        profiler_path=(
+            profiler_path
+            if profiler_path is not _sentinel
+            else get_profiler_path(
+                active_report["profiler_name"],
+                current_app=current_app,
+                remote_connection=remote_connection,
+            )
         ),
         remote_connection=(
             remote_connection.model_dump() if remote_connection else None
         ),
-        remote_profiler_folder=remote_profiler_folder.model_dump() if remote_profiler_folder else None,
+        remote_profiler_folder=(
+            remote_profiler_folder.model_dump() if remote_profiler_folder else None
+        ),
         remote_performance_folder=(
-            remote_performance_folder.model_dump() if remote_performance_folder else None
+            remote_performance_folder.model_dump()
+            if remote_performance_folder
+            else None
         ),
     )
 
@@ -255,10 +262,18 @@ def get_or_create_instance(
                 db.session.commit()
             except IntegrityError:
                 db.session.rollback()
-                instance_data = InstanceTable.query.filter_by(instance_id=instance_id).first()
+                instance_data = InstanceTable.query.filter_by(
+                    instance_id=instance_id
+                ).first()
 
         # Update the instance if any new data is provided
-        if profiler_name or performance_name or npe_name or remote_connection or remote_profiler_folder:
+        if (
+            profiler_name
+            or performance_name
+            or npe_name
+            or remote_connection
+            or remote_profiler_folder
+        ):
             update_instance(
                 instance_id=instance_id,
                 profiler_name=profiler_name,
@@ -269,7 +284,9 @@ def get_or_create_instance(
             )
 
             # Query again to get the updated instance data
-            instance_data = InstanceTable.query.filter_by(instance_id=instance_id).first()
+            instance_data = InstanceTable.query.filter_by(
+                instance_id=instance_id
+            ).first()
 
         return instance_data
 
@@ -318,7 +335,7 @@ def init_instances(app):
 
 
 def create_random_instance_id():
-    return ''.join(random.choices(string.ascii_lowercase + string.digits, k=45))
+    return "".join(random.choices(string.ascii_lowercase + string.digits, k=45))
 
 
 def create_instance_from_local_paths(profiler_path, performance_path):
@@ -328,11 +345,21 @@ def create_instance_from_local_paths(profiler_path, performance_path):
     if _profiler_path and (not _profiler_path.exists() or not _profiler_path.is_dir()):
         raise InvalidReportPath()
 
-    if _performance_path and (not _performance_path.exists() or not _performance_path.is_dir()):
+    if _performance_path and (
+        not _performance_path.exists() or not _performance_path.is_dir()
+    ):
         raise InvalidProfilerPath()
 
-    profiler_name = _profiler_path.parts[-1] if _profiler_path and len(_profiler_path.parts) > 2 else ""
-    performance_name = _performance_path.parts[-1] if _performance_path and len(_performance_path.parts) > 2 else ""
+    profiler_name = (
+        _profiler_path.parts[-1]
+        if _profiler_path and len(_profiler_path.parts) > 2
+        else ""
+    )
+    performance_name = (
+        _performance_path.parts[-1]
+        if _performance_path and len(_performance_path.parts) > 2
+        else ""
+    )
     instance_data = InstanceTable(
         instance_id=create_random_instance_id(),
         active_report={