ttnn-visualizer 0.24.0__py3-none-any.whl

ttnn_visualizer/serializers.py

@@ -0,0 +1,249 @@
+import dataclasses
+from collections import defaultdict
+from typing import List
+# SPDX-License-Identifier: Apache-2.0
+#
+# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
+
+
+from ttnn_visualizer.models import BufferType, Operation, TensorComparisonRecord
+
+
+def serialize_operations(
+    inputs,
+    operation_arguments,
+    operations,
+    outputs,
+    stack_traces,
+    tensors,
+    devices,
+    producers_consumers,
+    device_operations,
+):
+    tensors_dict = {t.tensor_id: t for t in tensors}
+    device_operations_dict = {
+        do.operation_id: do.captured_graph
+        for do in device_operations
+        if hasattr(do, "operation_id")
+    }
+
+    stack_traces_dict = {st.operation_id: st.stack_trace for st in stack_traces}
+
+    arguments_dict = defaultdict(list)
+    for argument in operation_arguments:
+        arguments_dict[argument.operation_id].append(argument)
+
+    inputs_dict, outputs_dict = serialize_inputs_outputs(
+        inputs, outputs, producers_consumers, tensors_dict
+    )
+
+    results = []
+    for operation in operations:
+
+        inputs = inputs_dict[operation.operation_id]
+        outputs = outputs_dict[operation.operation_id]
+        arguments = [a.to_dict() for a in arguments_dict[operation.operation_id]]
+        operation_data = operation.to_dict()
+        operation_data["id"] = operation.operation_id
+        operation_device_operations = device_operations_dict.get(
+            operation.operation_id, []
+        )
+        id = operation_data.pop("operation_id", None)
+
+        results.append(
+            {
+                **operation_data,
+                "id": id,
+                "stack_trace": stack_traces_dict.get(operation.operation_id),
+                "device_operations": operation_device_operations,
+                "arguments": arguments,
+                "inputs": inputs,
+                "outputs": outputs,
+            }
+        )
+    return results
+
+
+def serialize_inputs_outputs(
+    inputs,
+    outputs,
+    producers_consumers,
+    tensors_dict,
+    comparisons=None,
+):
+    producers_consumers_dict = {pc.tensor_id: pc for pc in producers_consumers}
+
+    def attach_tensor_data(values):
+        values_dict = defaultdict(list)
+        for value in values:
+            tensor = tensors_dict.get(value.tensor_id)
+            tensor_dict = tensor.to_dict()
+            pc = producers_consumers_dict.get(value.tensor_id)
+            value_dict = dataclasses.asdict(value)
+            value_dict.pop("tensor_id", None)
+            value_dict.update(
+                {
+                    "id": tensor_dict.pop("tensor_id"),
+                    "consumers": pc.consumers if pc else [],
+                    "producers": pc.producers if pc else [],
+                }
+            )
+            if comparisons:
+                comparison = comparisons.get(value.tensor_id)
+                value_dict.update({"comparison": comparison})
+
+            values_dict[value.operation_id].append({**value_dict, **tensor_dict})
+        return values_dict
+
+    inputs_dict = attach_tensor_data(inputs)
+    outputs_dict = attach_tensor_data(outputs)
+    return inputs_dict, outputs_dict
+
+
+def serialize_buffer_pages(buffer_pages):
+    # Collect device-specific data if needed
+
+    # Serialize each buffer page to a dictionary using dataclasses.asdict
+    buffer_pages_list = [page.to_dict() for page in buffer_pages]
+
+    # Optionally, modify or adjust the serialized data as needed
+    for page_data in buffer_pages_list:
+        # Set a custom id field if needed
+        page_data["id"] = f"{page_data['operation_id']}_{page_data['page_index']}"
+
+        # If the buffer_type is handled by an enum, adjust it similarly to your BufferPage model
+        if "buffer_type" in page_data and isinstance(
+            page_data["buffer_type"], BufferType
+        ):
+            page_data["buffer_type"] = page_data["buffer_type"].value
+
+    return buffer_pages_list
+
+
+def comparisons_by_tensor_id(
+    local_comparisons: List[TensorComparisonRecord],
+    global_comparisons: List[TensorComparisonRecord],
+):
+    comparisons = defaultdict(dict)
+    for local_comparison in local_comparisons:
+        comparisons[local_comparison.tensor_id].update({"local": local_comparison})
+    for global_comparison in global_comparisons:
+        comparisons[global_comparison.tensor_id].update({"global": global_comparison})
+    return comparisons
+
+
+def serialize_operation(
+    buffers,
+    inputs,
+    operation,
+    operation_arguments,
+    outputs,
+    stack_trace,
+    tensors,
+    global_tensor_comparisons,
+    local_tensor_comparisons,
+    devices,
+    producers_consumers,
+    device_operations,
+):
+    tensors_dict = {t.tensor_id: t for t in tensors}
+    comparisons = comparisons_by_tensor_id(
+        local_tensor_comparisons, global_tensor_comparisons
+    )
+
+    inputs_dict, outputs_dict = serialize_inputs_outputs(
+        inputs,
+        outputs,
+        producers_consumers,
+        tensors_dict,
+        comparisons,
+    )
+
+    buffer_list = [buffer.to_dict() for buffer in buffers]
+
+    l1_sizes = [d.worker_l1_size for d in devices]
+    arguments_data = [argument.to_dict() for argument in operation_arguments]
+    operation_data = operation.to_dict()
+    operation_data["id"] = operation.operation_id
+
+    inputs_data = inputs_dict.get(operation.operation_id)
+    outputs_data = outputs_dict.get(operation.operation_id)
+
+    id = operation_data.pop("operation_id", None)
+
+    device_operations_data = []
+    for do in device_operations:
+        if do.operation_id == operation.operation_id:
+            device_operations_data = do.captured_graph
+            break
+
+    return {
+        **operation_data,
+        "id": id,
+        "l1_sizes": l1_sizes,
+        "device_operations": device_operations_data,
+        "stack_trace": stack_trace.stack_trace if stack_trace else "",
+        "buffers": buffer_list,
+        "arguments": arguments_data,
+        "inputs": inputs_data or [],
+        "outputs": outputs_data or [],
+    }
+
+
+def serialize_operation_buffers(operation: Operation, operation_buffers):
+    buffer_data = [b.to_dict() for b in operation_buffers]
+    for b in buffer_data:
+        b.pop("operation_id")
+        b.update({"size": b.pop("max_size_per_bank")})
+    return {
+        "id": operation.operation_id,
+        "name": operation.name,
+        "buffers": list(buffer_data),
+    }
+
+
+def serialize_devices(devices):
+    return [d.to_dict() for d in devices]
+
+
+def serialize_operations_buffers(operations, buffers):
+
+    buffer_dict = defaultdict(list)
+    for b in buffers:
+        buffer_dict[b.operation_id].append(b)
+
+    results = []
+    for operation in operations:
+        operation_buffers = buffer_dict.get(operation.operation_id, [])
+        results.append(serialize_operation_buffers(operation, operation_buffers))
+    return results
+
+
+def serialize_tensors(
+    tensors, producers_consumers, local_comparisons, global_comparisons
+):
+    producers_consumers_dict = {pc.tensor_id: pc for pc in producers_consumers}
+    results = []
+    comparisons = comparisons_by_tensor_id(local_comparisons, global_comparisons)
+    for tensor in tensors:
+        tensor_data = tensor.to_dict()
+        tensor_id = tensor_data.pop("tensor_id")
+        tensor_data.update(
+            {
+                "id": tensor_id,
+                "comparison": comparisons.get(tensor_id),
+                "consumers": (
+                    producers_consumers_dict[tensor_id].consumers
+                    if tensor_id in producers_consumers_dict
+                    else []
+                ),
+                "producers": (
+                    producers_consumers_dict[tensor_id].producers
+                    if tensor_id in producers_consumers_dict
+                    else []
+                ),
+            }
+        )
+        results.append(tensor_data)
+
+    return results

ttnn_visualizer/sessions.py

@@ -0,0 +1,245 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
+
+from logging import getLogger
+
+from flask import request
+
+from ttnn_visualizer.utils import get_report_path, get_profiler_path
+from ttnn_visualizer.models import (
+    TabSessionTable,
+)
+from ttnn_visualizer.extensions import db
+
+logger = getLogger(__name__)
+
+from flask import jsonify, current_app
+from sqlalchemy.exc import SQLAlchemyError
+import json
+
+
+def update_existing_tab_session(
+    session_data,
+    report_name,
+    profile_name,
+    remote_connection,
+    remote_folder,
+    remote_profile_folder,
+    clear_remote,
+):
+    active_report = session_data.active_report or {}
+
+    if report_name:
+        active_report["report_name"] = report_name
+    if profile_name:
+        active_report["profile_name"] = profile_name
+
+    session_data.active_report = active_report
+
+    if remote_connection:
+        session_data.remote_connection = remote_connection.model_dump()
+    if remote_folder:
+        session_data.remote_folder = remote_folder.model_dump()
+    if remote_profile_folder:
+        session_data.remote_profile_folder = remote_profile_folder.model_dump()
+
+    if clear_remote:
+        clear_remote_data(session_data)
+
+    update_paths(
+        session_data, active_report, remote_connection, report_name, profile_name
+    )
+
+
+def clear_remote_data(session_data):
+    session_data.remote_connection = None
+    session_data.remote_folder = None
+    session_data.remote_profile_folder = None
+
+
+def handle_sqlalchemy_error(error):
+    current_app.logger.error(f"Failed to update tab session: {str(error)}")
+    db.session.rollback()
+
+
+def commit_and_log_session(session_data, tab_id):
+    db.session.commit()
+
+    session_data = TabSessionTable.query.filter_by(tab_id=tab_id).first()
+    current_app.logger.info(
+        f"Session data for tab {tab_id}: {json.dumps(session_data.to_dict(), indent=4)}"
+    )
+
+
+def update_paths(
+    session_data, active_report, remote_connection, report_name, profile_name
+):
+    if active_report.get("profile_name"):
+        session_data.profiler_path = get_profiler_path(
+            profile_name=active_report["profile_name"],
+            current_app=current_app,
+            remote_connection=remote_connection,
+        )
+
+    if active_report.get("report_name"):
+        session_data.report_path = get_report_path(
+            active_report=active_report,
+            current_app=current_app,
+            remote_connection=remote_connection,
+        )
+
+
+def create_new_tab_session(
+    tab_id,
+    report_name,
+    profile_name,
+    remote_connection,
+    remote_folder,
+    remote_profile_folder,
+    clear_remote,
+):
+    active_report = {}
+    if report_name:
+        active_report["report_name"] = report_name
+    if profile_name:
+        active_report["profile_name"] = profile_name
+
+    if clear_remote:
+        remote_connection = None
+        remote_folder = None
+        remote_profile_folder = None
+
+    session_data = TabSessionTable(
+        tab_id=tab_id,
+        active_report=active_report,
+        report_path=get_report_path(
+            active_report,
+            current_app=current_app,
+            remote_connection=remote_connection,
+        ),
+        remote_connection=(
+            remote_connection.model_dump() if remote_connection else None
+        ),
+        remote_folder=remote_folder.model_dump() if remote_folder else None,
+        remote_profile_folder=(
+            remote_profile_folder.model_dump() if remote_profile_folder else None
+        ),
+    )
+    db.session.add(session_data)
+    return session_data
+
+
+def update_tab_session(
+    tab_id,
+    report_name=None,
+    profile_name=None,
+    remote_connection=None,
+    remote_folder=None,
+    remote_profile_folder=None,
+    clear_remote=False,
+):
+    try:
+        session_data = get_or_create_tab_session(tab_id)
+
+        if session_data:
+            update_existing_tab_session(
+                session_data,
+                report_name,
+                profile_name,
+                remote_connection,
+                remote_folder,
+                remote_profile_folder,
+                clear_remote,
+            )
+        else:
+            session_data = create_new_tab_session(
+                tab_id,
+                report_name,
+                profile_name,
+                remote_connection,
+                remote_folder,
+                remote_profile_folder,
+                clear_remote,
+            )
+
+        commit_and_log_session(session_data, tab_id)
+        return jsonify({"message": "Tab session updated successfully"}), 200
+
+    except SQLAlchemyError as e:
+        handle_sqlalchemy_error(e)
+        return jsonify({"error": "Failed to update tab session"}), 500
+
+
+def get_or_create_tab_session(
+    tab_id,
+    report_name=None,
+    profile_name=None,
+    remote_connection=None,
+    remote_folder=None,
+):
+    """
+    Retrieve an existing tab session or create a new one if it doesn't exist.
+    Uses the TabSession model to manage session data and supports conditional updates.
+    """
+    try:
+        # Query the database for the tab session
+        session_data = TabSessionTable.query.filter_by(tab_id=tab_id).first()
+
+        # If session doesn't exist, initialize it
+        if not session_data:
+            session_data = TabSessionTable(
+                tab_id=tab_id,
+                active_report={},
+                remote_connection=None,
+                remote_folder=None,
+            )
+            db.session.add(session_data)
+            db.session.commit()
+
+        # Update the session if any new data is provided
+        if report_name or profile_name or remote_connection or remote_folder:
+            update_tab_session(
+                tab_id=tab_id,
+                report_name=report_name,
+                profile_name=profile_name,
+                remote_connection=remote_connection,
+                remote_folder=remote_folder,
+            )
+
+        # Query again to get the updated session data
+        session_data = TabSessionTable.query.filter_by(tab_id=tab_id).first()
+
+        return session_data
+
+    except SQLAlchemyError as e:
+        current_app.logger.error(f"Failed to get or create tab session: {str(e)}")
+        db.session.rollback()
+        return None
+
+
+def get_tab_session():
+    """
+    Middleware to retrieve or create a tab session based on the tab_id.
+    """
+    tab_id = request.args.get("tabId", None)
+
+    current_app.logger.info(f"get_tab_session: Received tab_id: {tab_id}")
+    if not tab_id:
+        current_app.logger.error("get_tab_session: No tab_id found")
+        return jsonify({"error": "tabId is required"}), 400
+
+    active_report = get_or_create_tab_session(tab_id)
+    current_app.logger.info(
+        f"get_tab_session: Session retrieved: {active_report.active_report}"
+    )
+
+    return jsonify({"active_report": active_report.active_report}), 200
+
+
+def init_sessions(app):
+    """
+    Initializes session middleware and hooks it into Flask.
+    """
+    app.before_request(get_tab_session)
+    app.logger.info("Sessions middleware initialized.")

ttnn_visualizer/settings.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
+
+import os
+from pathlib import Path
+from dotenv import load_dotenv
+
+from ttnn_visualizer.utils import str_to_bool
+
+load_dotenv()
+
+class DefaultConfig(object):
+    # General Settings
+    SECRET_KEY = os.getenv("SECRET_KEY", "90909")
+    DEBUG = bool(str_to_bool(os.getenv("FLASK_DEBUG", "false")))
+    TESTING = False
+    PRINT_ENV = True
+
+    # Path Settings
+    REPORT_DATA_DIRECTORY = Path(__file__).parent.absolute().joinpath("data")
+    VERSION = "0.14.2"
+    LOCAL_DATA_DIRECTORY = Path(REPORT_DATA_DIRECTORY).joinpath("local")
+    REMOTE_DATA_DIRECTORY = Path(REPORT_DATA_DIRECTORY).joinpath("remote")
+    APPLICATION_DIR = os.path.abspath(os.path.join(__file__, "..", os.pardir))
+    STATIC_ASSETS_DIR = Path(APPLICATION_DIR).joinpath("ttnn_visualizer", "static")
+    SEND_FILE_MAX_AGE_DEFAULT = 0
+
+    LAUNCH_BROWSER_ON_START = str_to_bool(os.getenv("LAUNCH_BROWSER_ON_START", "true"))
+
+    # File Name Configs
+    TEST_CONFIG_FILE = "config.json"
+    SQLITE_DB_PATH = "db.sqlite"
+
+    # For development you may want to disable sockets
+    USE_WEBSOCKETS = str_to_bool(os.getenv("USE_WEBSOCKETS", "true"))
+
+    # SQL Alchemy Settings
+    SQLALCHEMY_DATABASE_URI = (
+        f"sqlite:///{os.path.join(APPLICATION_DIR, f"ttnn_{VERSION}.db")}"
+    )
+    SQLALCHEMY_ENGINE_OPTIONS = {
+        "pool_size": 10,  # Adjust pool size as needed (default is 5)
+        "max_overflow": 20,  # Allow overflow of the pool size if necessary
+        "pool_timeout": 30,  # Timeout in seconds before giving up on getting a connection
+    }
+    SQLALCHEMY_TRACK_MODIFICATIONS = False
+
+    # Gunicorn settings
+    GUNICORN_WORKER_CLASS = os.getenv("GUNICORN_WORKER_CLASS", "gevent")
+    GUNICORN_WORKERS = os.getenv("GUNICORN_WORKERS", "1")
+    PORT = os.getenv("PORT", "8000")
+    HOST = "localhost"
+    DEV_SERVER_PORT = "5173"
+    DEV_SERVER_HOST = "localhost"
+
+    GUNICORN_BIND = f"{HOST}:{PORT}"
+    GUNICORN_APP_MODULE = os.getenv(
+        "GUNICORN_APP_MODULE", "ttnn_visualizer.app:create_app()"
+    )
+
+    # Session Settings
+    SESSION_COOKIE_SAMESITE = "Lax"
+    SESSION_COOKIE_SECURE = False  # For development on HTTP
+
+    def override_with_env_variables(self):
+        """Override config values with environment variables."""
+        for key, value in self.__class__.__dict__.items():
+            if not key.startswith("_"):  # Skip private/protected attributes
+                env_value = os.getenv(key)
+                if env_value is not None:
+                    setattr(self, key, env_value)
+
+    def to_dict(self):
+        """Return all config values as a dictionary, including inherited attributes."""
+        return {
+            key: getattr(self, key)
+            for key in dir(self)
+            if not key.startswith("_") and not callable(getattr(self, key))
+        }
+
+
+class DevelopmentConfig(DefaultConfig):
+    pass
+
+
+class TestingConfig(DefaultConfig):
+    DEBUG = bool(str_to_bool(os.getenv("FLASK_DEBUG", "True")))
+    TESTING = True
+
+
+class ProductionConfig(DefaultConfig):
+    DEBUG = False
+    TESTING = False
+
+
+class Config:
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super(Config, cls).__new__(cls)
+            cls._instance = cls._determine_config()
+            cls._instance.override_with_env_variables()
+        return cls._instance
+
+    @staticmethod
+    def _determine_config():
+        # Determine the environment
+        flask_env = os.getenv("FLASK_ENV", "development").lower()
+
+        # Choose the correct configuration class based on FLASK_ENV
+        if flask_env == "production":
+            return ProductionConfig()
+        elif flask_env == "testing":
+            return TestingConfig()
+        else:
+            return DevelopmentConfig()