ttnn-visualizer 0.25.1__py3-none-any.whl → 0.27.0__py3-none-any.whl

ttnn_visualizer/app.py

@@ -2,6 +2,7 @@
 #
 # SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
 
+import argparse
 import logging
 import os
 import subprocess
@@ -19,7 +20,8 @@ from flask_cors import CORS
 from werkzeug.debug import DebuggedApplication
 from werkzeug.middleware.proxy_fix import ProxyFix
 
-from ttnn_visualizer.exceptions import DatabaseFileNotFoundException
+from ttnn_visualizer.exceptions import DatabaseFileNotFoundException, InvalidProfilerPath, InvalidReportPath
+from ttnn_visualizer.sessions import create_instance_from_local_paths
 from ttnn_visualizer.settings import Config, DefaultConfig
 
 logger = logging.getLogger(__name__)
@@ -135,9 +137,10 @@ def middleware(app: flask.Flask):
     return None
 
 
-def open_browser(host, port, protocol="http"):
-
-    url = f"{protocol}://{host}:{port}"
+def open_browser(host, port, instance_id=None):
+    url = f"http://{host}:{port}"
+    if instance_id:
+        url = f"{url}?instanceId={instance_id}"
 
     print(f"Launching browser with url: {url}")
     try:
@@ -149,6 +152,13 @@ def open_browser(host, port, protocol="http"):
         print(f"Could not open the default browser: {e}")
 
 
+def parse_args():
+    parser = argparse.ArgumentParser(description="A tool for visualizing the Tenstorrent Neural Network model (TT-NN)")
+    parser.add_argument("--report-path", type=str, help="Specify a report path", default=None)
+    parser.add_argument("--profiler-path", help="Specify a profiler path", default=None)
+    return parser.parse_args()
+
+
 def main():
 
     run_command = sys.argv[0].split("/")
@@ -156,6 +166,23 @@ def main():
         os.environ.setdefault("FLASK_ENV", "production")
 
     config = cast(DefaultConfig, Config())
+    args = parse_args()
+    instance_id = None
+
+    if args.report_path or args.profiler_path:
+        app = create_app()
+        app.app_context().push()
+        try:
+            session = create_instance_from_local_paths(
+                report_path=args.report_path,
+                profiler_path=args.profiler_path,
+            )
+        except InvalidReportPath:
+            sys.exit("Invalid report path")
+        except InvalidProfilerPath:
+            sys.exit("Invalid profiler path")
+
+        instance_id = session.instance_id
 
     # Check if DEBUG environment variable is set
     debug_mode = os.environ.get("DEBUG", "false").lower() == "true"
@@ -182,7 +209,7 @@ def main():
         flask_env = os.getenv("FLASK_ENV", "development")
         port = config.PORT if flask_env == "production" else config.DEV_SERVER_PORT
         host = config.HOST if flask_env == "production" else config.DEV_SERVER_HOST
-        threading.Timer(2, open_browser, [host, port]).start()
+        threading.Timer(2, open_browser, [host, port, instance_id]).start()
     try:
         subprocess.run(gunicorn_args)
     except KeyboardInterrupt:

ttnn_visualizer/csv_queries.py

@@ -12,7 +12,7 @@ import pandas as pd
 from tt_perf_report import perf_report
 
 from ttnn_visualizer.exceptions import DataFormatError
-from ttnn_visualizer.models import TabSession
+from ttnn_visualizer.models import Instance
 from ttnn_visualizer.ssh_client import get_client
 
 
@@ -278,7 +278,7 @@ class DeviceLogProfilerQueries:
         "source file",
     ]
 
-    def __init__(self, session: TabSession):
+    def __init__(self, session: Instance):
         """
         Initialize the profiler with a session object.
         The session determines whether to use a local or remote runner.
@@ -366,7 +366,7 @@ class DeviceLogProfilerQueries:
         )
 
     @staticmethod
-    def get_raw_csv(session: TabSession):
+    def get_raw_csv(session: Instance):
         from ttnn_visualizer.sftp_operations import read_remote_file
 
         if (
@@ -455,7 +455,7 @@ class OpsPerformanceQueries:
         "HWCommandQueue_write_buffer_TT_HOST_FUNC [ns]",
     ]
 
-    def __init__(self, session: TabSession):
+    def __init__(self, session: Instance):
         """
         Initialize the performance profiler with a session object.
         """
@@ -571,6 +571,10 @@ class OpsPerformanceReportQueries:
         "raw_op_code"
     ]
 
+    PASSTHROUGH_COLUMNS = {
+        "pm_ideal_ns": "PM IDEAL [ns]",
+    }
+
     DEFAULT_SIGNPOST = None
     DEFAULT_IGNORE_SIGNPOSTS = None
     DEFAULT_MIN_PERCENTAGE = 0.5
@@ -596,6 +600,12 @@ class OpsPerformanceReportQueries:
             True,
         )
 
+        ops_perf_results = []
+        ops_perf_results_reader = csv.DictReader(StringIO(raw_csv))
+
+        for row in ops_perf_results_reader:
+            ops_perf_results.append(row)
+
         report = []
 
         try:
@@ -610,6 +620,12 @@ class OpsPerformanceReportQueries:
                         processed_row["advice"] = processed_row["advice"].split(" • ")
                     else:
                         processed_row["advice"] = []
+
+                    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
+                        processed_row[key] = ops_perf_results[idx][value]
+
                     report.append(processed_row)
         except csv.Error as e:
             raise DataFormatError() from e

ttnn_visualizer/decorators.py

@@ -19,7 +19,7 @@ from ttnn_visualizer.exceptions import (
     NoProjectsException,
     RemoteSqliteException,
 )
-from ttnn_visualizer.sessions import get_or_create_tab_session
+from ttnn_visualizer.sessions import get_or_create_instance
 
 
 def with_session(func):
@@ -27,18 +27,18 @@ def with_session(func):
     def wrapper(*args, **kwargs):
         from flask import current_app
 
-        tab_id = request.args.get("tabId")
+        instance_id = request.args.get("instanceId")
 
-        if not tab_id:
-            current_app.logger.error("No tabId present on request, returning 404")
+        if not instance_id:
+            current_app.logger.error("No instanceId present on request, returning 404")
             abort(404)
 
-        session_query_data = get_or_create_tab_session(tab_id=tab_id)
+        session_query_data = get_or_create_instance(instance_id=instance_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"
+                f"No active report exists for instanceId {instance_id}, returning 404"
             )
             # Raise 404 if report_path is missing or does not exist
             abort(404)

ttnn_visualizer/exceptions.py

@@ -38,3 +38,10 @@ class DatabaseFileNotFoundException(Exception):
 
 class DataFormatError(Exception):
     pass
+
+
+class InvalidReportPath(Exception):
+    pass
+
+class InvalidProfilerPath(Exception):
+    pass

ttnn_visualizer/file_uploads.py

@@ -43,6 +43,13 @@ def extract_report_name(files):
     unsplit_report_name = str(files[0].filename)
     return unsplit_report_name.split("/")[0]
 
+def extract_npe_name(files):
+    if not files:
+        return None
+
+    file_path = Path(files[0].filename)
+    return file_path.stem
+
 
 def save_uploaded_files(
     files,

ttnn_visualizer/models.py

@@ -182,6 +182,7 @@ class StatusMessage(SerializeableModel):
 class ActiveReport(SerializeableModel):
     report_name: Optional[str] = None
     profile_name: Optional[str] = None
+    npe_name: Optional[str] = None
 
 
 class RemoteReportFolder(SerializeableModel):
@@ -191,23 +192,25 @@ class RemoteReportFolder(SerializeableModel):
     lastSynced: Optional[int] = None
 
 
-class TabSession(BaseModel):
-    tab_id: str
+class Instance(BaseModel):
+    instance_id: str
     report_path: Optional[str] = None
     profiler_path: Optional[str] = None
+    npe_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"
+class InstanceTable(db.Model):
+    __tablename__ = "instances"
 
     id = Column(Integer, primary_key=True)
-    tab_id = Column(String, unique=True, nullable=False)
+    instance_id = Column(String, unique=True, nullable=False)
     report_path = Column(String)
     profiler_path = Column(String, nullable=True)
+    npe_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)
@@ -215,17 +218,19 @@ class TabSessionTable(db.Model):
 
     def __init__(
         self,
-        tab_id,
+        instance_id,
         active_report,
         remote_connection=None,
         remote_folder=None,
+        remote_profile_folder=None,
         report_path=None,
         profiler_path=None,
-        remote_profile_folder=None,
+        npe_path=None,
     ):
-        self.tab_id = tab_id
+        self.instance_id = instance_id
         self.active_report = active_report
         self.report_path = report_path
+        self.npe_path = npe_path
         self.remote_connection = remote_connection
         self.remote_folder = remote_folder
         self.profiler_path = profiler_path
@@ -234,22 +239,26 @@ class TabSessionTable(db.Model):
     def to_dict(self):
         return {
             "id": self.id,
-            "tab_id": self.tab_id,
+            "instance_id": self.instance_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,
+            "npe_path": self.npe_path
         }
 
-    def to_pydantic(self) -> TabSession:
-        return TabSession(
-            tab_id=str(self.tab_id),
+    def to_pydantic(self) -> Instance:
+        return Instance(
+            instance_id=str(self.instance_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
             ),
+            npe_path=(
+                str(self.npe_path) if self.npe_path is not None else None
+            ),
             active_report=(
                 (ActiveReport(**self.active_report) if self.active_report else None)
                 if isinstance(self.active_report, dict)

ttnn_visualizer/queries.py

@@ -13,7 +13,7 @@ from ttnn_visualizer.models import (
     DeviceOperation,
     Buffer,
     BufferPage,
-    TabSession,
+    Instance,
     Tensor,
     OperationArgument,
     StackTrace,
@@ -31,7 +31,7 @@ import paramiko
 
 
 class LocalQueryRunner:
-    def __init__(self, session: Optional[TabSession] = None, connection=None):
+    def __init__(self, session: Optional[Instance] = None, connection=None):
 
         if connection:
             self.connection = connection
@@ -66,7 +66,7 @@ class LocalQueryRunner:
 class RemoteQueryRunner:
     column_delimiter = "|||"
 
-    def __init__(self, session: TabSession):
+    def __init__(self, session: Instance):
         self.session = session
         self._validate_session()
         self.ssh_client = self._get_ssh_client(self.session.remote_connection)
@@ -165,11 +165,11 @@ class RemoteQueryRunner:
 
 class DatabaseQueries:
 
-    session: Optional[TabSession] = None
+    session: Optional[Instance] = None
     ssh_client = None
     query_runner: LocalQueryRunner | RemoteQueryRunner
 
-    def __init__(self, session: Optional[TabSession] = None, connection=None):
+    def __init__(self, session: Optional[Instance] = None, connection=None):
         self.session = session
 
         if connection:

ttnn_visualizer/sessions.py

@@ -2,13 +2,18 @@
 #
 # SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
 
+import json
+import random
+import string
 from logging import getLogger
+from pathlib import Path
 
 from flask import request
 
-from ttnn_visualizer.utils import get_report_path, get_profiler_path
+from ttnn_visualizer.exceptions import InvalidReportPath, InvalidProfilerPath
+from ttnn_visualizer.utils import get_report_path, get_profiler_path, get_npe_path
 from ttnn_visualizer.models import (
-    TabSessionTable,
+    InstanceTable,
 )
 from ttnn_visualizer.extensions import db
 
@@ -16,13 +21,13 @@ logger = getLogger(__name__)
 
 from flask import jsonify, current_app
 from sqlalchemy.exc import SQLAlchemyError
-import json
 
 
-def update_existing_tab_session(
+def update_existing_instance(
     session_data,
     report_name,
     profile_name,
+    npe_name,
     remote_connection,
     remote_folder,
     remote_profile_folder,
@@ -34,6 +39,8 @@ def update_existing_tab_session(
         active_report["report_name"] = report_name
     if profile_name:
         active_report["profile_name"] = profile_name
+    if npe_name:
+        active_report["npe_name"] = npe_name
 
     session_data.active_report = active_report
 
@@ -48,7 +55,7 @@ def update_existing_tab_session(
         clear_remote_data(session_data)
 
     update_paths(
-        session_data, active_report, remote_connection, report_name, profile_name
+        session_data, active_report, remote_connection
     )
 
 
@@ -63,17 +70,17 @@ def handle_sqlalchemy_error(error):
     db.session.rollback()
 
 
-def commit_and_log_session(session_data, tab_id):
+def commit_and_log_session(session_data, instance_id):
     db.session.commit()
 
-    session_data = TabSessionTable.query.filter_by(tab_id=tab_id).first()
+    session_data = InstanceTable.query.filter_by(instance_id=instance_id).first()
     current_app.logger.info(
-        f"Session data for tab {tab_id}: {json.dumps(session_data.to_dict(), indent=4)}"
+        f"Session data for instance {instance_id}: {json.dumps(session_data.to_dict(), indent=4)}"
     )
 
 
 def update_paths(
-    session_data, active_report, remote_connection, report_name, profile_name
+    session_data, active_report, remote_connection
 ):
     if active_report.get("profile_name"):
         session_data.profiler_path = get_profiler_path(
@@ -89,11 +96,18 @@ def update_paths(
             remote_connection=remote_connection,
         )
 
+    if active_report.get("npe_name"):
+        session_data.npe_path = get_npe_path(
+            npe_name=active_report["npe_name"],
+            current_app=current_app
+        )
+
 
-def create_new_tab_session(
-    tab_id,
+def create_new_instance(
+    instance_id,
     report_name,
     profile_name,
+    npe_name,
     remote_connection,
     remote_folder,
     remote_profile_folder,
@@ -104,14 +118,16 @@ def create_new_tab_session(
         active_report["report_name"] = report_name
     if profile_name:
         active_report["profile_name"] = profile_name
+    if npe_name:
+        active_report["npe_name"] = npe_name
 
     if clear_remote:
         remote_connection = None
         remote_folder = None
         remote_profile_folder = None
 
-    session_data = TabSessionTable(
-        tab_id=tab_id,
+    session_data = InstanceTable(
+        instance_id=instance_id,
         active_report=active_report,
         report_path=get_report_path(
             active_report,
@@ -130,40 +146,43 @@ def create_new_tab_session(
     return session_data
 
 
-def update_tab_session(
-    tab_id,
+def update_instance(
+    instance_id,
     report_name=None,
     profile_name=None,
+    npe_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)
+        session_data = get_or_create_instance(instance_id)
 
         if session_data:
-            update_existing_tab_session(
+            update_existing_instance(
                 session_data,
                 report_name,
                 profile_name,
+                npe_name,
                 remote_connection,
                 remote_folder,
                 remote_profile_folder,
                 clear_remote,
             )
         else:
-            session_data = create_new_tab_session(
-                tab_id,
+            session_data = create_new_instance(
+                instance_id,
                 report_name,
                 profile_name,
+                npe_name,
                 remote_connection,
                 remote_folder,
                 remote_profile_folder,
                 clear_remote,
             )
 
-        commit_and_log_session(session_data, tab_id)
+        commit_and_log_session(session_data, instance_id)
         return jsonify({"message": "Tab session updated successfully"}), 200
 
     except SQLAlchemyError as e:
@@ -171,25 +190,26 @@ def update_tab_session(
         return jsonify({"error": "Failed to update tab session"}), 500
 
 
-def get_or_create_tab_session(
-    tab_id,
+def get_or_create_instance(
+    instance_id,
     report_name=None,
     profile_name=None,
+    npe_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.
+    Uses the Instance 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()
+        session_data = InstanceTable.query.filter_by(instance_id=instance_id).first()
 
         # If session doesn't exist, initialize it
         if not session_data:
-            session_data = TabSessionTable(
-                tab_id=tab_id,
+            session_data = InstanceTable(
+                instance_id=instance_id,
                 active_report={},
                 remote_connection=None,
                 remote_folder=None,
@@ -198,17 +218,18 @@ def get_or_create_tab_session(
             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,
+        if report_name or profile_name or npe_name or remote_connection or remote_folder:
+            update_instance(
+                instance_id=instance_id,
                 report_name=report_name,
                 profile_name=profile_name,
+                npe_name=npe_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()
+        session_data = InstanceTable.query.filter_by(instance_id=instance_id).first()
 
         return session_data
 
@@ -218,20 +239,20 @@ def get_or_create_tab_session(
         return None
 
 
-def get_tab_session():
+def get_instance():
     """
-    Middleware to retrieve or create a tab session based on the tab_id.
+    Middleware to retrieve or create a tab session based on the instance_id.
     """
-    tab_id = request.args.get("tabId", None)
+    instance_id = request.args.get("instanceId", 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
+    current_app.logger.info(f"get_instance: Received instance_id: {instance_id}")
+    if not instance_id:
+        current_app.logger.error("get_instance: No instance_id found")
+        return jsonify({"error": "instanceId is required"}), 400
 
-    active_report = get_or_create_tab_session(tab_id)
+    active_report = get_or_create_instance(instance_id)
     current_app.logger.info(
-        f"get_tab_session: Session retrieved: {active_report.active_report}"
+        f"get_instance: Session retrieved: {active_report.active_report}"
     )
 
     return jsonify({"active_report": active_report.active_report}), 200
@@ -241,5 +262,39 @@ def init_sessions(app):
     """
     Initializes session middleware and hooks it into Flask.
     """
-    app.before_request(get_tab_session)
+    app.before_request(get_instance)
     app.logger.info("Sessions middleware initialized.")
+
+
+def create_random_instance_id():
+    return ''.join(random.choices(string.ascii_lowercase + string.digits, k=10))
+
+
+def create_instance_from_local_paths(report_path, profiler_path):
+    _report_path = Path(report_path)
+    _profiler_path = Path(profiler_path)
+
+    if not _report_path.exists():
+        raise InvalidReportPath()
+
+    if not _profiler_path.exists():
+        raise InvalidProfilerPath()
+
+    report_name = _report_path.parts[-2] if len(_report_path.parts) > 2 else ""
+    profile_name = _profiler_path.parts[-1] if len(_profiler_path.parts) > 2 else ""
+    session_data = InstanceTable(
+        instance_id=create_random_instance_id(),
+        active_report={
+            "report_name": report_name,
+            "profile_name": profile_name,
+            "npe_name": None,
+        },
+        report_path=report_path,
+        profiler_path=profiler_path,
+        remote_connection=None,
+        remote_folder=None,
+        remote_profile_folder=None,
+    )
+    db.session.add(session_data)
+    db.session.commit()
+    return session_data

ttnn_visualizer/settings.py

@@ -19,7 +19,7 @@ class DefaultConfig(object):
 
     # Path Settings
     REPORT_DATA_DIRECTORY = Path(__file__).parent.absolute().joinpath("data")
-    VERSION = "0.14.2"
+    VERSION = "0.26.0"
     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))