Fix perfetto import formatting

Author: drisspg

transformer_nuggets/utils/__init__.py

@@ -19,7 +19,6 @@
 from transformer_nuggets.utils.memory_viz import generate_memory_comparison_html
 from transformer_nuggets.utils.perfetto import (
     TraceFormat,
-    chrome_trace_to_track_event_trace,
     default_trace_path,
     default_track_event_path,
     perfetto_trace_path,
@@ -29,4 +28,5 @@
     write_trace,
     write_track_event_trace,
 )
+from transformer_nuggets.utils.track_event import chrome_trace_to_track_event_trace
 # from transformer_nuggets.utils.model_extraction import extract_attention_data

transformer_nuggets/utils/perfetto.py

@@ -24,10 +24,15 @@
 from collections import defaultdict
 from contextlib import contextmanager
 from dataclasses import dataclass
-from pathlib import Path
-from typing import Any, Literal
 from collections.abc import Iterator
+from pathlib import Path
 from re import Pattern
+from typing import Any, Literal
+
+from transformer_nuggets.utils.track_event import (
+    default_track_event_path,
+    write_track_event_trace,
+)
 
 
 TraceFormat = Literal["chrome_json", "track_event"]
@@ -426,13 +431,6 @@ def _reassign_sort_indices(
         )
 
 
-from transformer_nuggets.utils.track_event import (
-    chrome_trace_to_track_event_trace,
-    default_track_event_path,
-    write_track_event_trace,
-)
-
-
 def perfetto_trace_path(
     path: str | Path,
     *,

transformer_nuggets/utils/track_event.py

@@ -527,10 +527,14 @@ def _paired_flow_ids_by_slice(
             flow_ids_by_slice[source_slice.index].add(flow_id)
 
         for destination in destinations:
-            destination_slice = min(
-                kernel_slices,
-                key=lambda slc: abs(slc.ts_us - destination.ts_us),
-            ) if kernel_slices else None
+            destination_slice = (
+                min(
+                    kernel_slices,
+                    key=lambda slc: abs(slc.ts_us - destination.ts_us),
+                )
+                if kernel_slices
+                else None
+            )
             if destination_slice is None:
                 destination_slice = _smallest_containing_slice(slices, destination)
             if destination_slice is None:
@@ -814,9 +818,15 @@ def _emit_duration_markers(
     add_packet = builder.add_packet
     slice_begin = protos.TrackEvent.TYPE_SLICE_BEGIN
     slice_end = protos.TrackEvent.TYPE_SLICE_END
-    for ts_ns, _begin_order, _duration_key, track_uuid, _slice_index, is_begin, slc in (
-        _duration_markers(trace, track_ids)
-    ):
+    for (
+        ts_ns,
+        _begin_order,
+        _duration_key,
+        track_uuid,
+        _slice_index,
+        is_begin,
+        slc,
+    ) in _duration_markers(trace, track_ids):
         packet = add_packet()
         packet.timestamp = ts_ns
         packet.trusted_packet_sequence_id = trusted_packet_sequence_id