From eb35471caafcd32bacbc93508f58d57c50ef74fe Mon Sep 17 00:00:00 2001
From: Dom Del Nano <ddelnano@gmail.com>
Date: Thu, 24 Jul 2025 21:21:05 +0000
Subject: [PATCH] Fix script bundle creation and improve PEM/Kelvin heap dump
 collection for large clusters

Signed-off-by: Dom Del Nano <ddelnano@gmail.com>
(cherry picked from commit f1d8034ea6772bcaef4d5acff70c5034bfe7ce11)
---
 .../collect_agent_heaps.pxl}                  | 22 ++++++-----
 .../px/collect_agent_heaps/manifest.yaml      |  3 ++
 .../px/collect_agent_heaps/vis.json           | 39 +++++++++++++++++++
 3 files changed, 54 insertions(+), 10 deletions(-)
 rename src/pxl_scripts/px/{collect_heap_dumps.pxl => collect_agent_heaps/collect_agent_heaps.pxl} (52%)
 create mode 100644 src/pxl_scripts/px/collect_agent_heaps/manifest.yaml
 create mode 100644 src/pxl_scripts/px/collect_agent_heaps/vis.json

diff --git a/src/pxl_scripts/px/collect_heap_dumps.pxl b/src/pxl_scripts/px/collect_agent_heaps/collect_agent_heaps.pxl
similarity index 52%
rename from src/pxl_scripts/px/collect_heap_dumps.pxl
rename to src/pxl_scripts/px/collect_agent_heaps/collect_agent_heaps.pxl
index 88f6b2d8f22..f49bd5f2741 100644
--- a/src/pxl_scripts/px/collect_heap_dumps.pxl
+++ b/src/pxl_scripts/px/collect_agent_heaps/collect_agent_heaps.pxl
@@ -16,13 +16,15 @@
 
 import px
 
-df = px.GetAgentStatus()
-df.ip_address = px.pluck_array(px.split(df.ip_address, ":"), 0)
-df.hostname_by_ip = px.pod_id_to_node_name(px.ip_to_pod_id(df.ip_address))
-df.hostname = px.select(df.hostname_by_ip == "", df.hostname, df.hostname_by_ip)
-df = df[['asid', 'hostname']]
-heap_stats = px._HeapGrowthStacks()
-df = df.merge(heap_stats, how='inner', left_on='asid', right_on='asid')
-df.asid = df.asid_x
-df = df[['asid', 'hostname', 'heap']]
-px.display(df)
+
+# TODO(ddelnano): asid is unused until gh#2245 is addressed.
+def collect_pprofs(asid: int):
+    df = px.GetAgentStatus()
+    df.ip_address = px.pluck_array(px.split(df.ip_address, ":"), 0)
+    df.hostname_by_ip = px.pod_id_to_node_name(px.ip_to_pod_id(df.ip_address))
+    df.hostname = px.select(df.hostname_by_ip == "", df.hostname, df.hostname_by_ip)
+    df = df[['asid', 'hostname']]
+    heap_stats = px._HeapGrowthStacks()
+    df = df.merge(heap_stats, how='inner', left_on='asid', right_on='asid')
+    df.asid = df.asid_x
+    return df[['asid', 'hostname', 'heap']]
diff --git a/src/pxl_scripts/px/collect_agent_heaps/manifest.yaml b/src/pxl_scripts/px/collect_agent_heaps/manifest.yaml
new file mode 100644
index 00000000000..f31c74f81eb
--- /dev/null
+++ b/src/pxl_scripts/px/collect_agent_heaps/manifest.yaml
@@ -0,0 +1,3 @@
+---
+short: Collect Agent Heap Dumps
+long: Script useful for debugging kelvin and PEM memory footprint.
diff --git a/src/pxl_scripts/px/collect_agent_heaps/vis.json b/src/pxl_scripts/px/collect_agent_heaps/vis.json
new file mode 100644
index 00000000000..596db8b8053
--- /dev/null
+++ b/src/pxl_scripts/px/collect_agent_heaps/vis.json
@@ -0,0 +1,39 @@
+{
+    "variables": [
+        {
+            "name": "asid",
+            "type": "PX_INT64",
+            "description": "Whether to filter the results to a particular ASID",
+            "defaultValue": "-1"
+        }
+    ],
+    "globalFuncs": [
+        {
+            "outputName": "collect_pprofs",
+            "func": {
+                "name": "collect_pprofs",
+                "args": [
+                    {
+                        "name": "asid",
+                        "variable": "asid"
+                    }
+                ]
+            }
+        }
+    ],
+    "widgets": [
+        {
+            "name": "Table",
+            "position": {
+                "x": 0,
+                "y": 0,
+                "w": 12,
+                "h": 4
+            },
+            "globalFuncOutputName": "collect_pprofs",
+            "displaySpec": {
+                "@type": "types.px.dev/px.vispb.Table"
+            }
+        }
+    ]
+}