ttnn-visualizer 0.49.0__py3-none-any.whl → 0.64.0__py3-none-any.whl

ttnn_visualizer/views.py

@@ -25,12 +25,9 @@ from ttnn_visualizer.csv_queries import (
 from ttnn_visualizer.decorators import local_only, with_instance
 from ttnn_visualizer.enums import ConnectionTestStates
 from ttnn_visualizer.exceptions import (
-    AuthenticationException,
     AuthenticationFailedException,
     DataFormatError,
-    NoValidConnectionsError,
     RemoteConnectionException,
-    SSHException,
 )
 from ttnn_visualizer.file_uploads import (
     extract_folder_name_from_files,
@@ -46,7 +43,6 @@ from ttnn_visualizer.models import (
     StatusMessage,
 )
 from ttnn_visualizer.queries import DatabaseQueries
-from ttnn_visualizer.remote_sqlite_setup import check_sqlite_path, get_sqlite_path
 from ttnn_visualizer.serializers import (
     serialize_buffer,
     serialize_buffer_pages,
@@ -72,6 +68,7 @@ from ttnn_visualizer.utils import (
     create_path_resolver,
     get_cluster_descriptor_path,
     read_last_synced_file,
+    str_to_bool,
     timer,
 )
 
@@ -103,7 +100,11 @@ def operation_list(instance: Instance):
         devices = list(db.query_devices())
         producers_consumers = list(db.query_producers_consumers())
 
-        return serialize_operations(
+        error_records = None
+        if db._check_table_exists("errors"):
+            error_records = list(db.query_error_records())
+
+        serialized_operations = serialize_operations(
             inputs,
             operation_arguments,
             operations,
@@ -113,6 +114,11 @@ def operation_list(instance: Instance):
             devices,
             producers_consumers,
             device_operations,
+            error_records,
+        )
+        return Response(
+            orjson.dumps(serialized_operations),
+            mimetype="application/json",
         )
 
 
@@ -173,7 +179,15 @@ def operation_detail(operation_id, instance: Instance):
 
         devices = list(db.query_devices())
 
-        return serialize_operation(
+        error_record = None
+        if db._check_table_exists("errors"):
+            error_records = list(
+                db.query_error_records(filters={"operation_id": operation_id})
+            )
+            if error_records:
+                error_record = error_records[0]
+
+        serialized_operation = serialize_operation(
             buffers,
             inputs,
             operation,
@@ -186,6 +200,12 @@ def operation_detail(operation_id, instance: Instance):
             devices,
             producers_consumers,
             device_operations,
+            error_record,
+        )
+
+        return Response(
+            orjson.dumps(serialized_operation),
+            mimetype="application/json",
         )
 
 
@@ -198,9 +218,36 @@ def operation_history(instance: Instance):
         Path(str(instance.profiler_path)).parent / operation_history_filename
     )
     if not operation_history_file.exists():
-        return []
+        return jsonify([])
     with open(operation_history_file, "r") as file:
-        return json.load(file)
+        return Response(
+            orjson.dumps(json.load(file)),
+            mimetype="application/json",
+        )
+
+
+@api.route("/errors", methods=["GET"])
+@with_instance
+@timer
+def errors_list(instance: Instance):
+    with DatabaseQueries(instance) as db:
+        if not db._check_table_exists("errors"):
+            return (
+                jsonify(
+                    {
+                        "error": "Error records table does not exist in this report database."
+                    }
+                ),
+                HTTPStatus.UNPROCESSABLE_ENTITY,
+            )
+
+        error_records = list(db.query_error_records())
+        serialized_errors = [dataclasses.asdict(error) for error in error_records]
+
+        return Response(
+            orjson.dumps(serialized_errors),
+            mimetype="application/json",
+        )
 
 
 @api.route("/config")
@@ -211,7 +258,10 @@ def get_config(instance: Instance):
     if not config_file.exists():
         return {}
     with open(config_file, "r") as file:
-        return json.load(file)
+        return Response(
+            orjson.dumps(json.load(file)),
+            mimetype="application/json",
+        )
 
 
 @api.route("/tensors", methods=["GET"])
@@ -224,9 +274,13 @@ def tensors_list(instance: Instance):
         local_comparisons = list(db.query_tensor_comparisons())
         global_comparisons = list(db.query_tensor_comparisons(local=False))
         producers_consumers = list(db.query_producers_consumers())
-        return serialize_tensors(
+        serialized_tensors = serialize_tensors(
             tensors, producers_consumers, local_comparisons, global_comparisons
         )
+        return Response(
+            orjson.dumps(serialized_tensors),
+            mimetype="application/json",
+        )
 
 
 @api.route("/buffer", methods=["GET"])
@@ -248,7 +302,10 @@ def buffer_detail(instance: Instance):
         buffer = db.query_next_buffer(operation_id, address)
         if not buffer:
             return Response(status=HTTPStatus.NOT_FOUND)
-        return dataclasses.asdict(buffer)
+        return Response(
+            orjson.dumps(dataclasses.asdict(buffer)),
+            mimetype="application/json",
+        )
 
 
 @api.route("/buffer-pages", methods=["GET"])
@@ -283,7 +340,10 @@ def buffer_pages(instance: Instance):
                 )
             )
         )
-        return serialize_buffer_pages(buffers)
+        return Response(
+            orjson.dumps(serialize_buffer_pages(buffers)),
+            mimetype="application/json",
+        )
 
 
 @api.route("/tensors/<tensor_id>", methods=["GET"])
@@ -295,7 +355,10 @@ def tensor_detail(tensor_id, instance: Instance):
         if not tensors:
             return Response(status=HTTPStatus.NOT_FOUND)
 
-        return dataclasses.asdict(tensors[0])
+        return Response(
+            orjson.dumps(dataclasses.asdict(tensors[0])),
+            mimetype="application/json",
+        )
 
 
 @api.route("/buffers", methods=["GET"])
@@ -315,7 +378,7 @@ def get_all_buffers(instance: Instance):
             )
         )
         serialized = [serialize_buffer(b) for b in buffers]
-        return jsonify(serialized)
+        return Response(orjson.dumps(serialized), mimetype="application/json")
 
 
 @api.route("/operation-buffers", methods=["GET"])
@@ -323,6 +386,7 @@ def get_all_buffers(instance: Instance):
 def get_operations_buffers(instance: Instance):
     buffer_type = request.args.get("buffer_type", "")
     device_id = request.args.get("device_id", None)
+
     if buffer_type and str.isdigit(buffer_type):
         buffer_type = int(buffer_type)
     else:
@@ -335,7 +399,10 @@ def get_operations_buffers(instance: Instance):
             )
         )
         operations = list(db.query_operations())
-        return serialize_operations_buffers(operations, buffers)
+        return Response(
+            orjson.dumps(serialize_operations_buffers(operations, buffers)),
+            mimetype="application/json",
+        )
 
 
 @api.route("/operation-buffers/<operation_id>", methods=["GET"])
@@ -364,7 +431,11 @@ def get_operation_buffers(operation_id, instance: Instance):
         )
         if not operation:
             return Response(status=HTTPStatus.NOT_FOUND)
-        return serialize_operation_buffers(operation, buffers)
+
+        return Response(
+            orjson.dumps(serialize_operation_buffers(operation, buffers)),
+            mimetype="application/json",
+        )
 
 
 @api.route("/profiler", methods=["GET"])
@@ -374,12 +445,12 @@ def get_profiler_data_list(instance: Instance):
     resolver = create_path_resolver(current_app)
 
     # Note: "profiler" in app terminology maps to tt-metal's ttnn/reports
-    path = resolver.get_base_report_path("profiler", instance.remote_connection)
+    path = resolver.get_base_report_path("profiler")
 
     if not path.exists():
         if resolver.is_direct_report_mode:
             logger.warning(f"TT-Metal profiler reports not found: {path}")
-            return jsonify([])
+            return []
         else:
             path.mkdir(parents=True, exist_ok=True)
 
@@ -439,7 +510,7 @@ def get_profiler_data_list(instance: Instance):
             continue
         valid_dirs.append({"path": dir_path.name, "reportName": report_name})
 
-    return jsonify(valid_dirs)
+    return Response(orjson.dumps(valid_dirs), mimetype="application/json")
 
 
 @api.route("/profiler/<profiler_name>", methods=["DELETE"])
@@ -494,16 +565,12 @@ def get_performance_data_list(instance: Instance):
     resolver = create_path_resolver(current_app)
 
     # Note: "performance" in app terminology maps to tt-metal's profiler/reports
-    path = resolver.get_base_report_path("performance", instance.remote_connection)
-
-    is_remote = True if instance.remote_connection else False
+    path = resolver.get_base_report_path("performance")
 
     if not path.exists():
         if resolver.is_direct_report_mode:
             logger.warning(f"TT-Metal performance reports not found: {path}")
             return jsonify([])
-        elif not is_remote:
-            path.mkdir(parents=True, exist_ok=True)
 
     if current_app.config["SERVER_MODE"]:
         session_instances = session.get("instances", [])
@@ -564,7 +631,7 @@ def get_performance_data_list(instance: Instance):
             }
         )
 
-    return jsonify(valid_dirs)
+    return Response(orjson.dumps(valid_dirs), mimetype="application/json")
 
 
 @api.route("/performance/device-log", methods=["GET"])
@@ -572,9 +639,10 @@ def get_performance_data_list(instance: Instance):
 def get_performance_data(instance: Instance):
     if not instance.performance_path:
         return Response(status=HTTPStatus.NOT_FOUND)
+
     with DeviceLogProfilerQueries(instance) as csv:
         result = csv.get_all_entries(as_dict=True, limit=100)
-        return jsonify(result)
+        return Response(orjson.dumps(result), mimetype="application/json")
 
 
 @api.route("/performance/perf-results", methods=["GET"])
@@ -585,7 +653,7 @@ def get_profiler_performance_data(instance: Instance):
     with OpsPerformanceQueries(instance) as csv:
         # result = csv.query_by_op_code(op_code="(torch) contiguous", as_dict=True)
         result = csv.get_all_entries(as_dict=True, limit=100)
-        return jsonify(result)
+        return Response(orjson.dumps(result), mimetype="application/json")
 
 
 @api.route("/performance/<performance_name>", methods=["DELETE"])
@@ -659,6 +727,10 @@ def get_performance_results_report(instance: Instance):
         )
 
     name = request.args.get("name", None)
+    start_signpost = request.args.get("start_signpost", None)
+    end_signpost = request.args.get("end_signpost", None)
+    stack_by_in0 = str_to_bool(request.args.get("stack_by_in0", "true"))
+    hide_host_ops = str_to_bool(request.args.get("hide_host_ops", "true"))
 
     if name and not current_app.config["SERVER_MODE"]:
         performance_path = Path(instance.performance_path).parent / name
@@ -666,11 +738,17 @@ def get_performance_results_report(instance: Instance):
         logger.info(f"************ Performance path set to {instance.performance_path}")
 
     try:
-        report = OpsPerformanceReportQueries.generate_report(instance)
+        report = OpsPerformanceReportQueries.generate_report(
+            instance,
+            stack_by_in0=stack_by_in0,
+            start_signpost=start_signpost,
+            end_signpost=end_signpost,
+            hide_host_ops=hide_host_ops,
+        )
     except DataFormatError:
         return Response(status=HTTPStatus.UNPROCESSABLE_ENTITY)
 
-    return jsonify(report), 200
+    return Response(orjson.dumps(report), mimetype="application/json")
 
 
 @api.route("/performance/device-log/raw", methods=["GET"])
@@ -678,7 +756,16 @@ def get_performance_results_report(instance: Instance):
 def get_performance_data_raw(instance: Instance):
     if not instance.performance_path:
         return Response(status=HTTPStatus.NOT_FOUND)
+
+    name = request.args.get("name", None)
+
+    if name and not current_app.config["SERVER_MODE"]:
+        performance_path = Path(instance.performance_path).parent / name
+        instance.performance_path = str(performance_path)
+        logger.info(f"************ Performance path set to {instance.performance_path}")
+
     content = DeviceLogProfilerQueries.get_raw_csv(instance)
+
     return Response(
         content,
         mimetype="text/csv",
@@ -696,7 +783,7 @@ def get_npe_manifest(instance: Instance):
     except FileNotFoundError:
         return jsonify([])
 
-    return jsonify(content)
+    return Response(orjson.dumps(content), mimetype="application/json")
 
 
 @api.route("/performance/npe/timeline", methods=["GET"])
@@ -708,16 +795,16 @@ def get_npe_timeline(instance: Instance):
     filename = request.args.get("filename", default=None)
 
     if not filename:
-        return jsonify({})
+        return Response(orjson.dumps({}), mimetype="application/json")
 
     filename = Path(filename).name
 
     try:
         content = NPEQueries.get_npe_timeline(instance, filename=filename)
     except FileNotFoundError:
-        return jsonify({})
+        return Response(orjson.dumps({}), mimetype="application/json")
 
-    return jsonify(content)
+    return Response(orjson.dumps(content), mimetype="application/json")
 
 
 @api.route("/performance/device-log/zone/<zone>", methods=["GET"])
@@ -727,7 +814,7 @@ def get_zone_statistics(zone, instance: Instance):
         return Response(status=HTTPStatus.NOT_FOUND)
     with DeviceLogProfilerQueries(instance) as csv:
         result = csv.query_zone_statistics(zone_name=zone, as_dict=True)
-        return jsonify(result)
+        return Response(orjson.dumps(result), mimetype="application/json")
 
 
 @api.route("/devices", methods=["GET"])
@@ -735,7 +822,10 @@ def get_zone_statistics(zone, instance: Instance):
 def get_devices(instance: Instance):
     with DatabaseQueries(instance) as db:
         devices = list(db.query_devices())
-        return serialize_devices(devices)
+        return Response(
+            orjson.dumps(serialize_devices(devices)),
+            mimetype="application/json",
+        )
 
 
 @api.route("/local/upload/profiler", methods=["POST"])
@@ -866,7 +956,11 @@ def create_npe_files():
     data_directory = current_app.config["LOCAL_DATA_DIRECTORY"]
 
     for file in files:
-        if not file.filename.endswith(".json") and not file.filename.endswith(".zst"):
+        if (
+            not file.filename.endswith(".json")
+            and not file.filename.endswith(".zst")
+            and not file.filename.endswith(".npeviz")
+        ):
             return StatusMessage(
                 status=ConnectionTestStates.FAILED,
                 message="NPE requires a valid .json or .zst file",
@@ -915,7 +1009,10 @@ def get_remote_folders_profiler():
             if not rf.lastSynced:
                 logger.info(f"{directory_name} not yet synced")
 
-        return [r.model_dump() for r in remote_folders]
+        return Response(
+            orjson.dumps([r.model_dump() for r in remote_folders]),
+            mimetype="application/json",
+        )
     except RemoteConnectionException as e:
         return Response(status=e.http_status, response=e.message)
 
@@ -948,7 +1045,10 @@ def get_remote_folders_performance():
             if not rf.lastSynced:
                 logger.info(f"{performance_name} not yet synced")
 
-        return [r.model_dump() for r in remote_performance_folders]
+        return Response(
+            orjson.dumps([r.model_dump() for r in remote_performance_folders]),
+            mimetype="application/json",
+        )
     except RemoteConnectionException as e:
         return Response(status=e.http_status, response=e.message)
 
@@ -981,7 +1081,7 @@ def get_cluster_descriptor(instance: Instance):
         try:
             with open(local_path) as cluster_desc_file:
                 yaml_data = yaml.safe_load(cluster_desc_file)
-                return jsonify(yaml_data), 200
+                return jsonify(yaml_data)  # yaml_data is not compatible with orjson
         except yaml.YAMLError as e:
             return jsonify({"error": f"Failed to parse YAML: {str(e)}"}), 400
 
@@ -1011,7 +1111,7 @@ def test_remote_folder():
         add_status(
             ConnectionTestStates.FAILED.value, e.message, getattr(e, "detail", None)
         )
-        return [status.model_dump() for status in statuses], e.http_status
+        return jsonify([status.model_dump() for status in statuses]), e.http_status
     except RemoteConnectionException as e:
         add_status(
             ConnectionTestStates.FAILED.value, e.message, getattr(e, "detail", None)
@@ -1026,7 +1126,7 @@ def test_remote_folder():
             add_status(
                 ConnectionTestStates.FAILED.value, e.message, getattr(e, "detail", None)
             )
-            return [status.model_dump() for status in statuses], e.http_status
+            return jsonify([status.model_dump() for status in statuses]), e.http_status
         except RemoteConnectionException as e:
             add_status(
                 ConnectionTestStates.FAILED.value, e.message, getattr(e, "detail", None)
@@ -1041,7 +1141,7 @@ def test_remote_folder():
             add_status(
                 ConnectionTestStates.FAILED.value, e.message, getattr(e, "detail", None)
             )
-            return [status.model_dump() for status in statuses], e.http_status
+            return jsonify([status.model_dump() for status in statuses]), e.http_status
         except RemoteConnectionException as e:
             add_status(
                 ConnectionTestStates.FAILED.value, e.message, getattr(e, "detail", None)
@@ -1055,21 +1155,32 @@ def test_remote_folder():
             add_status(
                 ConnectionTestStates.FAILED.value, e.message, getattr(e, "detail", None)
             )
-            return [status.model_dump() for status in statuses], e.http_status
+            return jsonify([status.model_dump() for status in statuses]), e.http_status
         except RemoteConnectionException as e:
             add_status(
                 ConnectionTestStates.FAILED.value, e.message, getattr(e, "detail", None)
             )
 
-    return [status.model_dump() for status in statuses]
+    return Response(
+        orjson.dumps([status.model_dump() for status in statuses]),
+        mimetype="application/json",
+    )
 
 
 @api.route("/remote/read", methods=["POST"])
-def read_remote_folder():
-    connection = RemoteConnection.model_validate(request.json, strict=False)
+@with_instance
+def read_remote_folder(instance: Instance):
+    file_path = request.json.get("filePath")
+
+    remote_connection = instance.remote_connection
+    if not remote_connection:
+        return Response(
+            status=HTTPStatus.BAD_REQUEST,
+            response="No remote connection found in instance.",
+        )
+
     try:
-        # Only profilerPath is relevant here as we're reading the stack trace file
-        content = read_remote_file(connection, remote_path=connection.profilerPath)
+        content = read_remote_file(remote_connection, remote_path=file_path)
     except RemoteConnectionException as e:
         return Response(status=e.http_status, response=e.message)
     return Response(status=200, response=content)
@@ -1126,37 +1237,15 @@ def sync_remote_folder():
 
         remote_profiler_folder.lastSynced = int(time.time())
 
-        return remote_profiler_folder.model_dump()
+        return Response(
+            orjson.dumps(remote_profiler_folder.model_dump()),
+            mimetype="application/json",
+        )
 
     except RemoteConnectionException as e:
         return Response(status=e.http_status, response=e.message)
 
 
-@api.route("/remote/sqlite/detect-path", methods=["POST"])
-def detect_sqlite_path():
-    connection = request.json
-    connection = RemoteConnection.model_validate(connection, strict=False)
-    status_message = StatusMessage(
-        status=ConnectionTestStates.OK, message="Unable to Detect Path"
-    )
-    try:
-        path = get_sqlite_path(connection=connection)
-        if path:
-            status_message = StatusMessage(status=ConnectionTestStates.OK, message=path)
-        else:
-            status_message = StatusMessage(
-                status=ConnectionTestStates.OK, message="Unable to Detect Path"
-            )
-    except RemoteConnectionException as e:
-        current_app.logger.error(f"Unable to detect SQLite3 path {str(e)}")
-        status_message = StatusMessage(
-            status=ConnectionTestStates.FAILED,
-            message="Unable to detect SQLite3 path. See logs",
-        )
-    finally:
-        return status_message.model_dump()
-
-
 @api.route("/remote/use", methods=["POST"])
 def use_remote_folder():
     data = request.get_json(force=True)
@@ -1204,7 +1293,10 @@ def health_check():
 @with_instance
 def get_instance(instance: Instance):
     # Used to gate UI functions if no report is active
-    return instance.model_dump()
+    return Response(
+        orjson.dumps(instance.model_dump()),
+        mimetype="application/json",
+    )
 
 
 @api.route("/instance", methods=["PUT"])
@@ -1247,7 +1339,7 @@ def get_npe_data(instance: Instance):
     if instance.npe_path.endswith(".zst"):
         compressed_path = Path(instance.npe_path)
         uncompressed_path = None
-    elif instance.npe_path.endswith(".json"):
+    elif instance.npe_path.endswith(".json") or instance.npe_path.endswith(".npeviz"):
         compressed_path = None
         uncompressed_path = Path(instance.npe_path)
     else:
@@ -1266,17 +1358,15 @@ def get_npe_data(instance: Instance):
         if compressed_path and compressed_path.exists():
             with open(compressed_path, "rb") as file:
                 compressed_data = file.read()
-                uncompressed_data = zstd.uncompress(compressed_data)
-                npe_data = json.loads(uncompressed_data)
+                npe_data = zstd.uncompress(compressed_data)
         else:
             with open(uncompressed_path, "r") as file:
-                npe_data = json.load(file)
-    except json.JSONDecodeError as e:
-        logger.error(f"Invalid JSON in NPE file: {e}")
+                npe_data = file.read()
+    except Exception as e:
+        logger.error(f"Error reading NPE file: {e}")
         return Response(status=HTTPStatus.UNPROCESSABLE_ENTITY)
 
-    # Use orjson for much faster JSON serialization of large files
-    return Response(orjson.dumps(npe_data), mimetype="application/json")
+    return Response(npe_data, mimetype="application/json")
 
 
 @api.route("/notify", methods=["POST"])
@@ -1323,8 +1413,8 @@ def notify_report_update():
 
         logger.info(f"Report generated notification processed: {report_name}")
 
-        return (
-            jsonify(
+        return Response(
+            orjson.dumps(
                 {
                     "report_name": report_name,
                     "profiler_path": report_generated.profiler_path,
@@ -1333,9 +1423,13 @@ def notify_report_update():
                     "timestamp": report_generated.timestamp,
                 }
             ),
-            200,
+            mimetype="application/json",
         )
 
     except Exception as e:
         logger.error(f"Error processing report update notification: {str(e)}")
-        return jsonify({"error": "Internal server error"}), 500
+        return Response(
+            orjson.dumps({"error": "Internal server error"}),
+            mimetype="application/json",
+            status=HTTPStatus.INTERNAL_SERVER_ERROR,
+        )