more

Author: drisspg

test/test_memory_viz.py

@@ -48,6 +48,13 @@ def _extract_bootstrap(html):
     return json.loads(html[start:end])
 
 
+def _extract_named_bootstrap(html, script_id):
+    marker = f'<script id="{script_id}" type="application/json">'
+    start = html.index(marker) + len(marker)
+    end = html.index("</script>", start)
+    return json.loads(html[start:end])
+
+
 class TestExtractFrames:
     def test_filters_unwind_frames(self):
         frames = [
@@ -281,6 +288,65 @@ def test_uses_fixed_layout_instead_of_splitter(self, snapshot):
         assert "window.addEventListener('resize', applyPaneLayout);" in html
         assert 'id="splitter"' not in html
 
+    def test_uses_requested_devices_per_side(self):
+        snapshot_left = {
+            "device_traces": [
+                [_make_event("alloc", 0x1000, 128, time_us=1, filename="left.py", name="left")]
+            ],
+            "segments": [
+                {
+                    "device": 0,
+                    "stream": 0,
+                    "address": 0x1000,
+                    "blocks": [
+                        {
+                            "state": "active_allocated",
+                            "size": 64,
+                            "addr": 0x1000,
+                            "frames": [{"filename": "left.py", "name": "seed", "line": 1}],
+                        }
+                    ],
+                }
+            ],
+            "allocator_settings": {},
+        }
+        snapshot_right = {
+            "device_traces": [
+                [],
+                [_make_event("alloc", 0x2000, 256, time_us=2, filename="right.py", name="right")],
+            ],
+            "segments": [
+                {
+                    "device": 1,
+                    "stream": 0,
+                    "address": 0x2000,
+                    "blocks": [
+                        {
+                            "state": "active_allocated",
+                            "size": 128,
+                            "addr": 0x2000,
+                            "frames": [{"filename": "right.py", "name": "seed", "line": 1}],
+                        }
+                    ],
+                }
+            ],
+            "allocator_settings": {},
+        }
+        html = generate_memory_comparison_html(
+            snapshot_left,
+            snapshot_right,
+            device_left=0,
+            device_right=1,
+            title_left="Rank 0",
+            title_right="Rank 1",
+        )
+        bootstrap_left = _extract_named_bootstrap(html, "bootstrap-left")
+        bootstrap_right = _extract_named_bootstrap(html, "bootstrap-right")
+        assert bootstrap_left["meta"]["device"] == 0
+        assert bootstrap_right["meta"]["device"] == 1
+        assert bootstrap_left["meta"]["num_allocs"] == 2
+        assert bootstrap_right["meta"]["num_allocs"] == 2
+
 
 class TestNeverFreedAllocations:
     def test_never_freed_end_at_max_ts(self):

transformer_nuggets/utils/compare_memory.py

@@ -19,6 +19,8 @@ def main(
         "memory_comparison.html"
     ),
     device: Annotated[int, typer.Option("--device", help="CUDA device index.")] = 0,
+    device_left: Annotated[int | None, typer.Option("--device-left")] = None,
+    device_right: Annotated[int | None, typer.Option("--device-right")] = None,
     title_left: Annotated[str | None, typer.Option("--title-left")] = None,
     title_right: Annotated[str | None, typer.Option("--title-right")] = None,
 ):
@@ -39,6 +41,8 @@ def main(
         snapshot_left,
         snapshot_right,
         device=device,
+        device_left=device_left,
+        device_right=device_right,
         title_left=title_left or left.stem,
         title_right=title_right or right.stem,
     )

transformer_nuggets/utils/memory_viz.py

@@ -296,21 +296,27 @@ def generate_memory_comparison_html(
     snapshot_left: dict,
     snapshot_right: dict,
     device: int = 0,
+    device_left: int | None = None,
+    device_right: int | None = None,
     title_left: str = "Left",
     title_right: str = "Right",
 ) -> str:
+    if device_left is None:
+        device_left = device
+    if device_right is None:
+        device_right = device
     doc_title = f"{title_left} vs {title_right}"
     return (
         _MEMORY_COMPARISON_TEMPLATE.replace("__DOCUMENT_TITLE__", html.escape(doc_title))
         .replace("__TITLE_LEFT__", html.escape(title_left))
         .replace("__TITLE_RIGHT__", html.escape(title_right))
         .replace(
             "__BOOTSTRAP_LEFT__",
-            _json_for_html(_build_memory_viz_data(snapshot_left, device, title_left)),
+            _json_for_html(_build_memory_viz_data(snapshot_left, device_left, title_left)),
         )
         .replace(
             "__BOOTSTRAP_RIGHT__",
-            _json_for_html(_build_memory_viz_data(snapshot_right, device, title_right)),
+            _json_for_html(_build_memory_viz_data(snapshot_right, device_right, title_right)),
         )
     )