ttnn-visualizer 0.24.0__py3-none-any.whl

ttnn_visualizer/decorators.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
+
+import re
+from ttnn_visualizer.enums import ConnectionTestStates
+
+
+from functools import wraps
+from flask import request, abort
+from paramiko.ssh_exception import (
+    AuthenticationException,
+    NoValidConnectionsError,
+    SSHException,
+)
+
+from ttnn_visualizer.exceptions import (
+    RemoteConnectionException,
+    NoProjectsException,
+    RemoteSqliteException,
+)
+from ttnn_visualizer.sessions import get_or_create_tab_session
+
+
+def with_session(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        from flask import current_app
+
+        tab_id = request.args.get("tabId")
+
+        if not tab_id:
+            current_app.logger.error("No tabId present on request, returning 404")
+            abort(404)
+
+        session_query_data = get_or_create_tab_session(tab_id=tab_id)
+        session = session_query_data.to_pydantic()
+
+        if not session.active_report:
+            current_app.logger.error(
+                f"No active report exists for tabId {tab_id}, returning 404"
+            )
+            # Raise 404 if report_path is missing or does not exist
+            abort(404)
+
+        kwargs["session"] = session
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+def remote_exception_handler(func):
+    from flask import current_app
+
+    def remote_handler(*args, **kwargs):
+        if kwargs.get("connection", None):
+            connection = kwargs["connection"]
+        elif kwargs.get("remote_connection", None):
+            connection = kwargs["remote_connection"]
+        else:
+            connection = args[0]
+        try:
+            return func(*args, **kwargs)
+        except AuthenticationException as err:
+            current_app.logger.error(f"Authentication failed {err}")
+            raise RemoteConnectionException(
+                status=ConnectionTestStates.FAILED,
+                message=f"Unable to authenticate: {str(err)}",
+            )
+        except FileNotFoundError as err:
+            current_app.logger.error(f"File not found: {str(err)}")
+            raise RemoteConnectionException(
+                status=ConnectionTestStates.FAILED,
+                message=f"Unable to open path {connection.path}: {str(err)}",
+            )
+        except NoProjectsException as err:
+            current_app.logger.error(f"No projects: {str(err)}")
+            raise RemoteConnectionException(
+                status=ConnectionTestStates.FAILED,
+                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()
+            raise RemoteConnectionException(
+                status=ConnectionTestStates.FAILED,
+                message=f"{message}",
+            )
+
+        except RemoteSqliteException as err:
+            current_app.logger.error(f"Remote Sqlite exception: {str(err)}")
+            message = err.message
+            if "No such file" in str(err):
+                message = "Unable to open SQLite binary, check path"
+            raise RemoteConnectionException(
+                status=ConnectionTestStates.FAILED, message=message
+            )
+        except IOError as err:
+            message = f"Error opening remote folder {connection.path}: {str(err)}"
+            if "Name or service not known" in str(err):
+                message = f"Unable to connect to {connection.host} - check hostname"
+            raise RemoteConnectionException(
+                status=ConnectionTestStates.FAILED,
+                message=message,
+            )
+        except SSHException as err:
+            if str(err) == "No existing session":
+                message = "Authentication failed - check credentials and ssh-agent"
+            else:
+                err_message = re.sub(r"\[.*?]", "", str(err)).strip()
+                message = f"Error connecting to host {connection.host}: {err_message}"
+
+            raise RemoteConnectionException(
+                status=ConnectionTestStates.FAILED, message=message
+            )
+
+    return remote_handler

ttnn_visualizer/enums.py

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
+
+import enum
+
+
+class ConnectionTestStates(enum.Enum):
+    IDLE = 0
+    PROGRESS = 1
+    FAILED = 2
+    OK = 3

ttnn_visualizer/exceptions.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
+
+from http import HTTPStatus
+
+from ttnn_visualizer.enums import ConnectionTestStates
+
+
+class RemoteConnectionException(Exception):
+    def __init__(self, message, status: ConnectionTestStates):
+        super().__init__(message)
+        self.message = message
+        self.status = status
+
+    @property
+    def http_status(self):
+        if self.status == ConnectionTestStates.FAILED:
+            return HTTPStatus.INTERNAL_SERVER_ERROR
+        if self.status == ConnectionTestStates.OK:
+            return HTTPStatus.OK
+
+
+class NoProjectsException(RemoteConnectionException):
+    pass
+
+
+class RemoteSqliteException(Exception):
+    def __init__(self, message, status):
+        super().__init__(message)
+        self.message = message
+        self.status = status
+
+
+class DatabaseFileNotFoundException(Exception):
+    pass
+
+
+class DataFormatError(Exception):
+    pass

ttnn_visualizer/extensions.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
+
+from flask_socketio import SocketIO
+from flask_static_digest import FlaskStaticDigest
+from flask_sqlalchemy import SQLAlchemy
+
+
+flask_static_digest = FlaskStaticDigest()
+# Initialize Flask SQLAlchemy
+db = SQLAlchemy()
+
+socketio = SocketIO(cors_allowed_origins="*", async_mode="gevent")

ttnn_visualizer/file_uploads.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# SPDX-FileCopyrightText: © 2024 Tenstorrent Inc.
+
+from pathlib import Path
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def validate_files(files, required_files, pattern=None):
+    """Validate uploaded files against required file names and an optional pattern."""
+    found_files = set()
+
+    for file in files:
+        file_path = Path(file.filename)
+
+        if file_path.name in required_files or (
+            pattern and file_path.name.startswith(pattern)
+        ):
+            found_files.add(file_path.name)
+            if len(file_path.parents) != 2:
+                logger.warning(
+                    f"File {file.filename} is not under a single parent folder."
+                )
+                return False
+
+    missing_files = required_files - found_files
+    if pattern and not any(name.startswith(pattern) for name in found_files):
+        missing_files.add(f"{pattern}*")
+
+    if missing_files:
+        logger.warning(f"Missing required files: {', '.join(missing_files)}")
+        return False
+
+    return True
+
+
+def extract_report_name(files):
+    """Extract the report name from the first file."""
+    if not files:
+        return None
+    unsplit_report_name = str(files[0].filename)
+    return unsplit_report_name.split("/")[0]
+
+
+def save_uploaded_files(
+    files,
+    target_directory,
+    report_name=None,
+):
+    """
+    Save uploaded files to the target directory.
+
+    :param files: List of files to be saved.
+    :param target_directory: The base directory for saving the files.
+    :param report_name: The report name to use for the directory.
+    :param modify_path: Optional function to modify the file path before saving.
+    :param flat_structure: If True, saves files directly under the report_name directory without subdirectories.
+    """
+    for file in files:
+        current_file_name = str(file.filename)
+        logger.info(f"Processing file: {current_file_name}")
+
+        file_path = Path(current_file_name)
+
+        destination_file = Path(target_directory).joinpath(str(file_path))
+
+        logger.info(f"Writing file to {destination_file}")
+
+        # Create directory if it doesn't exist
+        if not destination_file.parent.exists():
+            logger.info(
+                f"{destination_file.parent.name} does not exist. Creating directory"
+            )
+            destination_file.parent.mkdir(exist_ok=True, parents=True)
+
+        file.save(destination_file)

ttnn_visualizer/models.py

@@ -0,0 +1,275 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
+
+import dataclasses
+import enum
+import json
+from json import JSONDecodeError
+from typing import Optional, Any
+
+from pydantic import BaseModel, Field
+from sqlalchemy import Integer, Column, String, JSON
+from sqlalchemy.ext.mutable import MutableDict
+
+from ttnn_visualizer.utils import SerializeableDataclass
+from ttnn_visualizer.enums import ConnectionTestStates
+from ttnn_visualizer.extensions import db
+
+from ttnn_visualizer.utils import parse_memory_config
+
+
+class BufferType(enum.Enum):
+    DRAM = 0
+    L1 = 1
+    SYSTEM_MEMORY = 2
+    L1_SMALL = 3
+    TRACE = 4
+
+
+@dataclasses.dataclass
+class Operation(SerializeableDataclass):
+    operation_id: int
+    name: str
+    duration: float
+
+
+@dataclasses.dataclass
+class Device(SerializeableDataclass):
+    device_id: int
+    num_y_cores: int
+    num_x_cores: int
+    num_y_compute_cores: int
+    num_x_compute_cores: int
+    worker_l1_size: int
+    l1_num_banks: int
+    l1_bank_size: int
+    address_at_first_l1_bank: int
+    address_at_first_l1_cb_buffer: int
+    num_banks_per_storage_core: int
+    num_compute_cores: int
+    num_storage_cores: int
+    total_l1_memory: int
+    total_l1_for_tensors: int
+    total_l1_for_interleaved_buffers: int
+    total_l1_for_sharded_buffers: int
+    cb_limit: int
+
+
+@dataclasses.dataclass
+class DeviceOperation(SerializeableDataclass):
+    operation_id: int
+    captured_graph: str
+
+    def __post_init__(self):
+        try:
+            captured_graph = json.loads(self.captured_graph)
+            for graph in captured_graph:
+                id = graph.pop("counter")
+                graph.update({"id": id})
+
+            self.captured_graph = captured_graph
+
+        except JSONDecodeError:
+            self.captured_graph = json.dumps({})
+
+
+@dataclasses.dataclass
+class Buffer(SerializeableDataclass):
+    operation_id: int
+    device_id: int
+    address: int
+    max_size_per_bank: int
+    buffer_type: BufferType
+
+
+@dataclasses.dataclass
+class BufferPage(SerializeableDataclass):
+    operation_id: int
+    device_id: int
+    address: int
+    core_y: int
+    core_x: int
+    bank_id: int
+    page_index: int
+    page_address: int
+    page_size: int
+    buffer_type: BufferType
+
+
+@dataclasses.dataclass
+class ProducersConsumers(SerializeableDataclass):
+    tensor_id: int
+    producers: list[int]
+    consumers: list[int]
+
+
+@dataclasses.dataclass
+class Tensor(SerializeableDataclass):
+    tensor_id: int
+    shape: str
+    dtype: str
+    layout: str
+    memory_config: str | dict[str, Any] | None
+    device_id: int
+    address: int
+    buffer_type: BufferType
+    device_addresses: list[int]
+
+    def __post_init__(self):
+        self.memory_config = parse_memory_config(self.memory_config)
+
+@dataclasses.dataclass
+class InputTensor(SerializeableDataclass):
+    operation_id: int
+    input_index: int
+    tensor_id: int
+
+
+@dataclasses.dataclass
+class OutputTensor(SerializeableDataclass):
+    operation_id: int
+    output_index: int
+    tensor_id: int
+
+
+@dataclasses.dataclass
+class TensorComparisonRecord(SerializeableDataclass):
+    tensor_id: int
+    golden_tensor_id: int
+    matches: bool
+    desired_pcc: bool
+    actual_pcc: float
+
+
+@dataclasses.dataclass
+class OperationArgument(SerializeableDataclass):
+    operation_id: int
+    name: str
+    value: str
+
+
+@dataclasses.dataclass
+class StackTrace(SerializeableDataclass):
+    operation_id: int
+    stack_trace: str
+
+
+# Non Data Models
+
+
+class SerializeableModel(BaseModel):
+    class Config:
+        use_enum_values = True
+
+
+class RemoteConnection(SerializeableModel):
+    name: str
+    username: str
+    host: str
+    port: int = Field(ge=1, le=65535)
+    reportPath: str
+    performancePath: Optional[str] = None
+    sqliteBinaryPath: Optional[str] = None
+    useRemoteQuerying: bool = False
+
+
+class StatusMessage(SerializeableModel):
+    status: ConnectionTestStates
+    message: str
+
+
+class ActiveReport(SerializeableModel):
+    report_name: Optional[str] = None
+    profile_name: Optional[str] = None
+
+
+class RemoteReportFolder(SerializeableModel):
+    testName: str
+    remotePath: str
+    lastModified: int
+    lastSynced: Optional[int] = None
+
+
+class TabSession(BaseModel):
+    tab_id: str
+    report_path: Optional[str] = None
+    profiler_path: Optional[str] = None
+    active_report: Optional[ActiveReport] = None
+    remote_connection: Optional[RemoteConnection] = None
+    remote_folder: Optional[RemoteReportFolder] = None
+    remote_profile_folder: Optional[RemoteReportFolder] = None
+
+
+class TabSessionTable(db.Model):
+    __tablename__ = "tab_sessions"
+
+    id = Column(Integer, primary_key=True)
+    tab_id = Column(String, unique=True, nullable=False)
+    report_path = Column(String)
+    profiler_path = Column(String, nullable=True)
+    active_report = db.Column(MutableDict.as_mutable(JSON), nullable=False, default={})
+    remote_connection = Column(JSON, nullable=True)
+    remote_folder = Column(JSON, nullable=True)
+    remote_profile_folder = Column(JSON, nullable=True)
+
+    def __init__(
+        self,
+        tab_id,
+        active_report,
+        remote_connection=None,
+        remote_folder=None,
+        report_path=None,
+        profiler_path=None,
+        remote_profile_folder=None,
+    ):
+        self.tab_id = tab_id
+        self.active_report = active_report
+        self.report_path = report_path
+        self.remote_connection = remote_connection
+        self.remote_folder = remote_folder
+        self.profiler_path = profiler_path
+        self.remote_profile_folder = remote_profile_folder
+
+    def to_dict(self):
+        return {
+            "id": self.id,
+            "tab_id": self.tab_id,
+            "active_report": self.active_report,
+            "remote_connection": self.remote_connection,
+            "remote_folder": self.remote_folder,
+            "remote_profile_folder": self.remote_profile_folder,
+            "report_path": self.report_path,
+            "profiler_path": self.profiler_path,
+        }
+
+    def to_pydantic(self) -> TabSession:
+        return TabSession(
+            tab_id=str(self.tab_id),
+            report_path=str(self.report_path) if self.report_path is not None else None,
+            profiler_path=(
+                str(self.profiler_path) if self.profiler_path is not None else None
+            ),
+            active_report=(
+                (ActiveReport(**self.active_report) if self.active_report else None)
+                if isinstance(self.active_report, dict)
+                else None
+            ),
+            remote_connection=(
+                RemoteConnection.model_validate(self.remote_connection, strict=False)
+                if self.remote_connection is not None
+                else None
+            ),
+            remote_folder=(
+                RemoteReportFolder.model_validate(self.remote_folder, strict=False)
+                if self.remote_folder is not None
+                else None
+            ),
+            remote_profile_folder=(
+                RemoteReportFolder.model_validate(
+                    self.remote_profile_folder, strict=False
+                )
+                if self.remote_profile_folder is not None
+                else None
+            ),
+        )