ttnn-visualizer 0.42.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

@@ -11,16 +11,22 @@ 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.exceptions import (
+    SSHException,
+    AuthenticationException,
+    NoValidConnectionsError,
+)
 from ttnn_visualizer.models import Instance, RemoteConnection
-from ttnn_visualizer.exceptions import SSHException, AuthenticationException, NoValidConnectionsError
-from ttnn_visualizer.models import Instance
 from ttnn_visualizer.sftp_operations import read_remote_file
 
 
-def handle_ssh_subprocess_error(e: subprocess.CalledProcessError, remote_connection: RemoteConnection):
+def handle_ssh_subprocess_error(
+    e: subprocess.CalledProcessError, remote_connection: RemoteConnection
+):
     """
     Convert subprocess SSH errors to appropriate SSH exceptions.
 
@@ -31,23 +37,29 @@ def handle_ssh_subprocess_error(e: subprocess.CalledProcessError, remote_connect
     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"
-    ]):
+    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"
-    ]):
+    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
@@ -58,6 +70,7 @@ def handle_ssh_subprocess_error(e: subprocess.CalledProcessError, remote_connect
     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
@@ -154,24 +167,22 @@ class RemoteCSVQueryRunner:
 
     def _execute_ssh_command(self, command: str) -> str:
         """Execute an SSH command and return the output."""
-        ssh_cmd = ["ssh"]
-        
+        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
-        ])
-        
+
+        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
+                ssh_cmd, capture_output=True, text=True, check=True, timeout=30
             )
             return result.stdout
         except subprocess.CalledProcessError as e:
@@ -269,7 +280,7 @@ class RemoteCSVQueryRunner:
         )
         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]:
         """
@@ -324,7 +335,9 @@ class NPEQueries:
             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)
@@ -338,7 +351,9 @@ class NPEQueries:
     @staticmethod
     def get_npe_timeline(instance: Instance, filename: str):
         if not filename:
-            raise ValueError("filename parameter is required and cannot be None or empty")
+            raise ValueError(
+                "filename parameter is required and cannot be None or empty"
+            )
 
         if (
             not instance.remote_connection
@@ -347,18 +362,33 @@ class NPEQueries:
             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)
+            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
-            return read_remote_file(
-                instance.remote_connection,
-                f"{profiler_folder.remotePath}/{NPEQueries.NPE_FOLDER}/{filename}",
+            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:
@@ -614,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)
@@ -656,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}")
@@ -688,7 +718,7 @@ class OpsPerformanceReportQueries:
         "output_subblock_w",
         "global_call_count",
         "advice",
-        "raw_op_code"
+        "raw_op_code",
     ]
 
     PASSTHROUGH_COLUMNS = {
@@ -734,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(" • ")
@@ -743,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

@@ -14,6 +14,7 @@ from ttnn_visualizer.exceptions import (
     NoValidConnectionsError,
     SSHException,
     RemoteConnectionException,
+    AuthenticationFailedException,
     NoProjectsException,
     RemoteSqliteException,
 )
@@ -34,21 +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}")
+            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)
 
@@ -68,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)}")
@@ -86,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:
@@ -111,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
 
@@ -50,14 +75,17 @@ class InvalidProfilerPath(Exception):
 
 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={

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
             ),