ttnn-visualizer 0.36.0__py3-none-any.whl → 0.37.0__py3-none-any.whl

ttnn_visualizer/app.py

@@ -21,7 +21,7 @@ from werkzeug.debug import DebuggedApplication
 from werkzeug.middleware.proxy_fix import ProxyFix
 
 from ttnn_visualizer.exceptions import DatabaseFileNotFoundException, InvalidProfilerPath, InvalidReportPath
-from ttnn_visualizer.sessions import create_instance_from_local_paths
+from ttnn_visualizer.instances import create_instance_from_local_paths
 from ttnn_visualizer.settings import Config, DefaultConfig
 
 logger = logging.getLogger(__name__)
@@ -127,7 +127,7 @@ def middleware(app: flask.Flask):
         app.wsgi_app = ProxyFix(app.wsgi_app)
 
     # CORS configuration
-    origins = ["http://localhost:5173", "http://localhost:8000"]
+    origins = app.config["ALLOWED_ORIGINS"]
 
     CORS(
         app,

ttnn_visualizer/csv_queries.py

@@ -277,38 +277,38 @@ class DeviceLogProfilerQueries:
         "source file",
     ]
 
-    def __init__(self, session: Instance):
+    def __init__(self, instance: Instance):
         """
-        Initialize the profiler with a session object.
-        The session determines whether to use a local or remote runner.
+        Initialize the profiler with a instance object.
+        The instance determines whether to use a local or remote runner.
         """
-        self.session = session
+        self.instance = instance
         self.runner = None
 
     def __enter__(self):
         """
-        Determine the appropriate query runner based on the session's remote connection.
+        Determine the appropriate query runner based on the instance's remote connection.
         """
 
-        is_remote = self.session.remote_connection
+        is_remote = self.instance.remote_connection
         use_remote_querying = False
 
         # Disabled until we resolve the issue with sqlite versions
         # if is_remote:
-        #     use_remote_querying = self.session.remote_connection.useRemoteQuerying
+        #     use_remote_querying = self.instance.remote_connection.useRemoteQuerying
 
         # Determine if this is a local or remote operation
         if is_remote and use_remote_querying:
-            remote_profiler_folder = self.session.remote_profile_folder
+            remote_profiler_folder = self.instance.remote_profile_folder
             file_path = f"{remote_profiler_folder.remotePath}/{self.DEVICE_LOG_FILE}"
             self.runner = RemoteCSVQueryRunner(
                 file_path=file_path,
-                remote_connection=self.session.remote_connection,
+                remote_connection=self.instance.remote_connection,
                 offset=1,  # Skip the first line for device log files
             )
         else:
             self.runner = LocalCSVQueryRunner(
-                file_path=Path(self.session.performance_path).joinpath(
+                file_path=Path(self.instance.performance_path).joinpath(
                     self.DEVICE_LOG_FILE
                 ),
                 offset=1,  # Skip the first line for device log files
@@ -365,23 +365,23 @@ class DeviceLogProfilerQueries:
         )
 
     @staticmethod
-    def get_raw_csv(session: Instance):
+    def get_raw_csv(instance: Instance):
         from ttnn_visualizer.sftp_operations import read_remote_file
 
         if (
-            not session.remote_connection
-            or session.remote_connection
-            and not session.remote_connection.useRemoteQuerying
+            not instance.remote_connection
+            or instance.remote_connection
+            and not instance.remote_connection.useRemoteQuerying
         ):
             file_path = Path(
-                session.performance_path, DeviceLogProfilerQueries.DEVICE_LOG_FILE
+                instance.performance_path, DeviceLogProfilerQueries.DEVICE_LOG_FILE
             )
             with open(file_path, "r") as f:
                 return f.read()
         else:
-            profiler_folder = session.remote_profile_folder
+            profiler_folder = instance.remote_profile_folder
             return read_remote_file(
-                session.remote_connection,
+                instance.remote_connection,
                 f"{profiler_folder.remotePath}/{DeviceLogProfilerQueries.DEVICE_LOG_FILE}",
             )
 
@@ -454,11 +454,11 @@ class OpsPerformanceQueries:
         "HWCommandQueue_write_buffer_TT_HOST_FUNC [ns]",
     ]
 
-    def __init__(self, session: Instance):
+    def __init__(self, instance: Instance):
         """
-        Initialize the performance profiler with a session object.
+        Initialize the performance profiler with a instance object.
         """
-        self.session = session
+        self.instance = instance
         self.runner = None
 
     def __enter__(self):
@@ -466,7 +466,7 @@ class OpsPerformanceQueries:
 
         :return:
         """
-        file_path = OpsPerformanceQueries.get_local_ops_perf_file_path(self.session)
+        file_path = OpsPerformanceQueries.get_local_ops_perf_file_path(self.instance)
         self.runner = LocalCSVQueryRunner(file_path=file_path, offset=1)
         self.runner.__enter__()
 
@@ -477,8 +477,8 @@ class OpsPerformanceQueries:
         return self
 
     @staticmethod
-    def get_local_ops_perf_file_path(session):
-        performance_path = Path(session.performance_path)
+    def get_local_ops_perf_file_path(instance):
+        performance_path = Path(instance.performance_path)
 
         # Find the latest file with the correct prefix
         perf_files = list(
@@ -494,29 +494,29 @@ class OpsPerformanceQueries:
         return str(latest_file)
 
     @staticmethod
-    def get_remote_ops_perf_file_path(session):
+    def get_remote_ops_perf_file_path(instance):
         from ttnn_visualizer.sftp_operations import resolve_file_path
 
-        remote_profile_folder = session.remote_profile_folder.remotePath
+        remote_profile_folder = instance.remote_profile_folder.remotePath
         return resolve_file_path(
-            session.remote_connection,
+            instance.remote_connection,
             f"{remote_profile_folder}/{OpsPerformanceQueries.PERF_RESULTS_PREFIX}*",
         )
 
     @staticmethod
-    def get_raw_csv(session):
+    def get_raw_csv(instance):
         from ttnn_visualizer.sftp_operations import read_remote_file
 
         if (
-            not session.remote_connection
-            or session.remote_connection
-            and not session.remote_connection.useRemoteQuerying
+            not instance.remote_connection
+            or instance.remote_connection
+            and not instance.remote_connection.useRemoteQuerying
         ):
-            with open(OpsPerformanceQueries.get_local_ops_perf_file_path(session)) 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(session)
-            return read_remote_file(session.remote_connection, path)
+            path = OpsPerformanceQueries.get_remote_ops_perf_file_path(instance)
+            return read_remote_file(instance.remote_connection, path)
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         """
@@ -600,8 +600,8 @@ class OpsPerformanceReportQueries:
     DEFAULT_TRACING_MODE = False
 
     @classmethod
-    def generate_report(cls, session):
-        raw_csv = OpsPerformanceQueries.get_raw_csv(session)
+    def generate_report(cls, instance):
+        raw_csv = OpsPerformanceQueries.get_raw_csv(instance)
         csv_file = StringIO(raw_csv)
         csv_output_file = tempfile.mktemp(suffix=".csv")
         perf_report.generate_perf_report(

ttnn_visualizer/decorators.py

@@ -7,7 +7,7 @@ from ttnn_visualizer.enums import ConnectionTestStates
 
 
 from functools import wraps
-from flask import request, abort
+from flask import abort, request, session
 from paramiko.ssh_exception import (
     AuthenticationException,
     NoValidConnectionsError,
@@ -19,10 +19,10 @@ from ttnn_visualizer.exceptions import (
     NoProjectsException,
     RemoteSqliteException,
 )
-from ttnn_visualizer.sessions import get_or_create_instance
+from ttnn_visualizer.instances import get_or_create_instance
 
 
-def with_session(func):
+def with_instance(func):
     @wraps(func)
     def wrapper(*args, **kwargs):
         from flask import current_app
@@ -33,10 +33,17 @@ def with_session(func):
             current_app.logger.error("No instanceId present on request, returning 404")
             abort(404)
 
-        session_query_data = get_or_create_instance(instance_id=instance_id)
-        session = session_query_data.to_pydantic()
+        instance_query_data = get_or_create_instance(instance_id=instance_id)
+        instance = instance_query_data.to_pydantic()
+
+        kwargs["instance"] = instance
+
+        if 'instances' not in session:
+            session['instances'] = []
+
+        if instance.instance_id not in session['instances']:
+            session['instances'].append(instance.instance_id)
 
-        kwargs["session"] = session
         return func(*args, **kwargs)
 
     return wrapper

ttnn_visualizer/file_uploads.py

@@ -2,9 +2,18 @@
 #
 # SPDX-FileCopyrightText: © 2024 Tenstorrent Inc.
 
-from pathlib import Path
 import logging
+import os
 import re
+import shlex
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+
+from flask import current_app
+
+from ttnn_visualizer.exceptions import DataFormatError
 
 logger = logging.getLogger(__name__)
 
@@ -37,13 +46,12 @@ def validate_files(files, required_files, pattern=None, folder_name=None):
     return True
 
 
-def extract_profiler_name(files):
+def extract_folder_name_from_files(files):
     """Extract the report name from the first file."""
     if not files:
         return None
-    unsplit_profiler_name = str(files[0].filename)
-    return unsplit_profiler_name.split("/")[0]
-
+    unsplit_name = str(files[0].filename)
+    return unsplit_name.split("/")[0]
 
 def extract_npe_name(files):
     if not files:
@@ -54,8 +62,8 @@ def extract_npe_name(files):
 
 def save_uploaded_files(
     files,
-    target_directory,
-    folder_name=None,
+    base_directory,
+    parent_folder_name=None,
 ):
     """
     Save uploaded files to the target directory.
@@ -64,24 +72,82 @@ def save_uploaded_files(
     :param target_directory: The base directory for saving the files.
     :param folder_name: The name to use for the directory.
     """
-    for file in files:
-        current_file_name = str(file.filename)
-        logger.info(f"Processing file: {current_file_name}")
+    if current_app.config["MALWARE_SCANNER"]:
+        scanned_files = scan_uploaded_files(files, base_directory, parent_folder_name)
+
+        for temp_path, dest_path in scanned_files:
+            if dest_path.parent.exists():
+                dest_path.parent.mkdir(exist_ok=True, parents=True)
+
+            logger.info(f"Saving uploaded file (clean): {dest_path}")
+            shutil.move(temp_path, dest_path)
+    else:
+        for file in files:
+            dest_path = construct_dest_path(file, base_directory, parent_folder_name)
+            logger.info(f"Writing file to {dest_path}")
+
+            # Create directory if it doesn't exist
+            if not dest_path.parent.exists():
+                logger.info(
+                    f"{dest_path.parent.name} does not exist. Creating directory"
+                )
+                dest_path.parent.mkdir(exist_ok=True, parents=True)
 
-        file_path = Path(current_file_name)
+            logger.info(f"Saving uploaded file: {dest_path}")
+            file.save(dest_path)
 
-        if folder_name:
-            destination_file = Path(target_directory) / folder_name / str(file_path)
-        else:
-            destination_file = Path(target_directory) / str(file_path)
+            # Update the modified time of the parent directory (for sorting purposes)
+            os.utime(dest_path.parent, None)
 
-        logger.info(f"Writing file to {destination_file}")
+    # Update the modified time of the uploaded directory
+    if parent_folder_name:
+        uploaded_dir = Path(base_directory) / parent_folder_name
+    else:
+        uploaded_dir = Path(base_directory)
+    if uploaded_dir.exists():
+        os.utime(uploaded_dir, None)
 
-        # 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)
+def scan_uploaded_files(
+    files,
+    target_directory,
+    folder_name=None,
+):
+    scanned_files = []
+
+    for file in files:
+        (_, temp_path) = tempfile.mkstemp()
+        file.save(temp_path)
+        dest_path = construct_dest_path(file, target_directory, folder_name)
+
+        cmd_list = shlex.split(current_app.config["MALWARE_SCANNER"])
+        cmd_list.append(temp_path)
+
+        try:
+            result = subprocess.run(
+                cmd_list,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+            )
+            if result.returncode == 0:
+                scanned_files.append((temp_path, dest_path))
+            else:
+                os.unlink(temp_path)
+                logger.warning(f"Malware scanner flagged file: {file.filename}")
+                raise DataFormatError()
+        except Exception as e:
+            if os.path.exists(temp_path):
+                os.unlink(temp_path)
+            raise
+
+    return scanned_files
+
+
+def construct_dest_path(file, target_directory, folder_name):
+    if folder_name:
+        dest_path = Path(target_directory) / folder_name / str(file.filename)
+    else:
+        dest_path = Path(target_directory) / str(file.filename)
+
+    return dest_path

ttnn_visualizer/{sessions.py → instances.py}

@@ -20,11 +20,11 @@ from ttnn_visualizer.extensions import db
 logger = getLogger(__name__)
 
 from flask import jsonify, current_app
-from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.exc import IntegrityError, SQLAlchemyError
 
 
 def update_existing_instance(
-    session_data,
+    instance_data,
     profiler_name,
     performance_name,
     npe_name,
@@ -33,7 +33,7 @@ def update_existing_instance(
     remote_performance_folder,
     clear_remote,
 ):
-    active_report = session_data.active_report or {}
+    active_report = instance_data.active_report or {}
 
     # First ifs are explicit deletes and elifs are updates
     if profiler_name == "":
@@ -51,62 +51,62 @@ def update_existing_instance(
     elif npe_name is not None:
         active_report["npe_name"] = npe_name
 
-    session_data.active_report = active_report
+    instance_data.active_report = active_report
 
     if remote_connection:
-        session_data.remote_connection = remote_connection.model_dump()
+        instance_data.remote_connection = remote_connection.model_dump()
     if remote_profiler_folder:
-        session_data.remote_profiler_folder = remote_profiler_folder.model_dump()
+        instance_data.remote_profiler_folder = remote_profiler_folder.model_dump()
     if remote_performance_folder:
-        session_data.remote_performance_folder = remote_performance_folder.model_dump()
+        instance_data.remote_performance_folder = remote_performance_folder.model_dump()
 
     if clear_remote:
-        clear_remote_data(session_data)
+        clear_remote_data(instance_data)
 
     update_paths(
-        session_data, active_report, remote_connection
+        instance_data, active_report, remote_connection
     )
 
 
-def clear_remote_data(session_data):
-    session_data.remote_connection = None
-    session_data.remote_profiler_folder = None
-    session_data.remote_performance_folder = None
+def clear_remote_data(instance_data):
+    instance_data.remote_connection = None
+    instance_data.remote_profiler_folder = None
+    instance_data.remote_performance_folder = None
 
 
 def handle_sqlalchemy_error(error):
-    current_app.logger.error(f"Failed to update tab session: {str(error)}")
+    current_app.logger.error(f"Failed to update tab instance: {str(error)}")
     db.session.rollback()
 
 
-def commit_and_log_session(session_data, instance_id):
+def commit_and_log_session(instance_data, instance_id):
     db.session.commit()
 
-    session_data = InstanceTable.query.filter_by(instance_id=instance_id).first()
+    instance_data = InstanceTable.query.filter_by(instance_id=instance_id).first()
     current_app.logger.info(
-        f"Session data for instance {instance_id}: {json.dumps(session_data.to_dict(), indent=4)}"
+        f"Data for instance {instance_id}: {json.dumps(instance_data.to_dict(), indent=4)}"
     )
 
 
 def update_paths(
-    session_data, active_report, remote_connection
+    instance_data, active_report, remote_connection
 ):
     if active_report.get("performance_name"):
-        session_data.performance_path = get_performance_path(
+        instance_data.performance_path = get_performance_path(
             performance_name=active_report["performance_name"],
             current_app=current_app,
             remote_connection=remote_connection,
         )
 
     if active_report.get("profiler_name"):
-        session_data.profiler_path = get_profiler_path(
+        instance_data.profiler_path = get_profiler_path(
             profiler_name=active_report["profiler_name"],
             current_app=current_app,
             remote_connection=remote_connection,
         )
 
     if active_report.get("npe_name"):
-        session_data.npe_path = get_npe_path(
+        instance_data.npe_path = get_npe_path(
             npe_name=active_report["npe_name"],
             current_app=current_app
         )
@@ -135,7 +135,7 @@ def create_new_instance(
         remote_profiler_folder = None
         remote_performance_folder = None
 
-    session_data = InstanceTable(
+    instance_data = InstanceTable(
         instance_id=instance_id,
         active_report=active_report,
         profiler_path=get_profiler_path(
@@ -151,8 +151,8 @@ def create_new_instance(
             remote_performance_folder.model_dump() if remote_performance_folder else None
         ),
     )
-    db.session.add(session_data)
-    return session_data
+    db.session.add(instance_data)
+    return instance_data
 
 
 def update_instance(
@@ -166,11 +166,11 @@ def update_instance(
     clear_remote=False,
 ):
     try:
-        session_data = get_or_create_instance(instance_id)
+        instance_data = get_or_create_instance(instance_id)
 
-        if session_data:
+        if instance_data:
             update_existing_instance(
-                session_data,
+                instance_data,
                 profiler_name,
                 performance_name,
                 npe_name,
@@ -180,7 +180,7 @@ def update_instance(
                 clear_remote,
             )
         else:
-            session_data = create_new_instance(
+            instance_data = create_new_instance(
                 instance_id,
                 profiler_name,
                 performance_name,
@@ -191,12 +191,12 @@ def update_instance(
                 clear_remote,
             )
 
-        commit_and_log_session(session_data, instance_id)
-        return jsonify({"message": "Tab session updated successfully"}), 200
+        commit_and_log_session(instance_data, instance_id)
+        return jsonify({"message": "Tab instance updated successfully"}), 200
 
     except SQLAlchemyError as e:
         handle_sqlalchemy_error(e)
-        return jsonify({"error": "Failed to update tab session"}), 500
+        return jsonify({"error": "Failed to update tab instance"}), 500
 
 
 def get_or_create_instance(
@@ -208,25 +208,30 @@ def get_or_create_instance(
     remote_profiler_folder=None,
 ):
     """
-    Retrieve an existing tab session or create a new one if it doesn't exist.
-    Uses the Instance model to manage session data and supports conditional updates.
+    Retrieve an existing tab instance or create a new one if it doesn't exist.
+    Uses the Instance model to manage instance data and supports conditional updates.
     """
     try:
-        # Query the database for the tab session
-        session_data = InstanceTable.query.filter_by(instance_id=instance_id).first()
+        # Query the database for the tab instance
+        instance_data = InstanceTable.query.filter_by(instance_id=instance_id).first()
 
-        # If session doesn't exist, initialize it
-        if not session_data:
-            session_data = InstanceTable(
+        # If instance doesn't exist, initialize it
+        if not instance_data:
+            instance_data = InstanceTable(
                 instance_id=instance_id,
                 active_report={},
                 remote_connection=None,
                 remote_profiler_folder=None,
             )
-            db.session.add(session_data)
-            db.session.commit()
+            db.session.add(instance_data)
 
-        # Update the session if any new data is provided
+            try:
+                db.session.commit()
+            except IntegrityError:
+                db.session.rollback()
+                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:
             update_instance(
                 instance_id=instance_id,
@@ -237,20 +242,20 @@ def get_or_create_instance(
                 remote_profiler_folder=remote_profiler_folder,
             )
 
-        # Query again to get the updated session data
-        session_data = InstanceTable.query.filter_by(instance_id=instance_id).first()
+            # Query again to get the updated instance data
+            instance_data = InstanceTable.query.filter_by(instance_id=instance_id).first()
 
-        return session_data
+        return instance_data
 
     except SQLAlchemyError as e:
-        current_app.logger.error(f"Failed to get or create tab session: {str(e)}")
+        current_app.logger.error(f"Failed to get or create tab instance: {str(e)}")
         db.session.rollback()
         return None
 
 
 def get_instance():
     """
-    Middleware to retrieve or create a tab session based on the instance_id.
+    Middleware to retrieve or create a tab instance based on the instance_id.
     """
     instance_id = request.args.get("instanceId", None)
 
@@ -261,22 +266,33 @@ def get_instance():
 
     active_report = get_or_create_instance(instance_id)
     current_app.logger.info(
-        f"get_instance: Session retrieved: {active_report.active_report}"
+        f"get_instance: active report retrieved: {active_report.active_report}"
     )
 
     return jsonify({"active_report": active_report.active_report}), 200
 
 
-def init_sessions(app):
+def get_instances(instance_ids):
+    instances = []
+
+    for instance_id in instance_ids:
+        instance = InstanceTable.query.filter_by(instance_id=instance_id).first()
+        if instance:
+            instances.append(instance)
+
+    return instances
+
+
+def init_instances(app):
     """
-    Initializes session middleware and hooks it into Flask.
+    Initializes instance middleware and hooks it into Flask.
     """
     app.before_request(get_instance)
-    app.logger.info("Sessions middleware initialized.")
+    app.logger.info("Instances middleware initialized.")
 
 
 def create_random_instance_id():
-    return ''.join(random.choices(string.ascii_lowercase + string.digits, k=10))
+    return ''.join(random.choices(string.ascii_lowercase + string.digits, k=45))
 
 
 def create_instance_from_local_paths(profiler_path, performance_path):
@@ -291,7 +307,7 @@ def create_instance_from_local_paths(profiler_path, performance_path):
 
     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 ""
-    session_data = InstanceTable(
+    instance_data = InstanceTable(
         instance_id=create_random_instance_id(),
         active_report={
             "profiler_name": profiler_name,
@@ -304,6 +320,6 @@ def create_instance_from_local_paths(profiler_path, performance_path):
         remote_profiler_folder=None,
         remote_performance_folder=None,
     )
-    db.session.add(session_data)
+    db.session.add(instance_data)
     db.session.commit()
-    return session_data
+    return instance_data

ttnn_visualizer/models.py

@@ -186,7 +186,7 @@ class ActiveReports(SerializeableModel):
 
 
 class RemoteReportFolder(SerializeableModel):
-    testName: str
+    reportName: str
     remotePath: str
     lastModified: int
     lastSynced: Optional[int] = None