Make PEM heap dump scripts compatible with Kelvin (#2220)

Author: ddelnano

scripts/collect_heap_pprofs.sh

@@ -19,43 +19,68 @@
 usage() {
   echo "This script downloads all of the files listed in the mappings section of a heap profile."
   echo ""
-  echo "Usage: $0 <heap_profile_dir> <vizier_cluster_id> [<gcloud ssh opts>...]"
+  echo "Usage: $0 <heap_profile_dir> <vizier_cluster_id> [--gcloud-common-args] [--gcloud-ssh-args...]"
   echo "<heap_profile_dir> : the directory where the heap profile and memory mapped files will be stored. It will be created if it does not exist."
   echo "<vizier_cluster_id> : the ID of the Vizier cluster to connect to."
-  echo "Common gcloud ssh options include --project."
+  echo "[--gcloud-common-args] : common arguments to pass to gcloud commands, such as --project."
+  echo "[--gcloud-ssh-args] : additional arguments to pass to gcloud compute ssh commands, such as --internal-ip."
   exit 1
 }
 
-heap_profile_dir="$1"
-cluster_id="$2"
+parse_args() {
+  if [ $# -lt 2 ]; then
+    usage
+  fi
+
+  heap_profile_dir="$1"
+  shift
+
+  cluster_id="$1"
+  shift
+}
+
+check_args() {
+  if [ -z "$heap_profile_dir" ] || [ -z "$cluster_id" ]; then
+    usage
+  fi
+}
+
+parse_args "${@}"
+check_args
+
 script_dir=$(dirname "$(realpath "$0")")
 repo_root=$(git rev-parse --show-toplevel)
 
-if [ -z "$heap_profile_dir" ] || [ -z "$cluster_id" ]; then
-  usage
-fi
-
 mkdir -p "$heap_profile_dir"
 
-pxl_heap_output_file="${heap_profile_dir}/raw_output_from_hot_table_test.json"
+pxl_heap_output_file="${heap_profile_dir}/heap_dump_raw_output.json"
 
 px run -o json -c "$cluster_id" -f "${repo_root}/src/pxl_scripts/px/collect_heap_dumps.pxl"  > "$pxl_heap_output_file"
 
 while IFS= read -r line; do
+    asid=$(echo "$line" | jq -r '.asid')
     hostname=$(echo "$line" | jq -r '.hostname')
     heap_content=$(echo "$line" | jq -r '.heap')
-    echo "$heap_content" > "${heap_profile_dir}/${hostname}.txt"
-    echo "Wrote ${heap_profile_dir}/${hostname}.txt"
+    filename="${heap_profile_dir}/${hostname}-${asid}.txt"
+    echo "$heap_content" > "$filename"
+    echo "Wrote heap profile for ASID $asid on host $hostname to $filename"
 done < "$pxl_heap_output_file"
 
-nodes=()
+declare -a agents
 for file in "${heap_profile_dir}"/*.txt; do
-  hostname=$(basename "${file%.*}")
-  nodes+=("$hostname")
-  hostname_dir="${heap_profile_dir}/${hostname}"
-  mkdir -p "$hostname_dir"
+  filename=$(basename "$file")
+  filename_noext="${filename%.txt}"
+  hostname="${filename_noext%-*}"
+  asid="${filename_noext##*-}"
+
+  agents["$asid"]="$hostname"
+  agent_dir="${heap_profile_dir}/${hostname}-${asid}"
+  mkdir -p "$agent_dir"
 done
 
-for node in "${nodes[@]}"; do
-  "${script_dir}/download_heap_prof_mapped_files.sh" "${heap_profile_dir}/${node}.txt" "$node" "${@:3}"
+for agent in "${!agents[@]}"; do
+  asid="${agent}"
+  hostname="${agents[$agent]}"
+  filename="${heap_profile_dir}/${hostname}-${asid}.txt"
+  "${script_dir}/download_heap_prof_mapped_files.sh" "$filename" "$hostname" "${@:3}"
 done

scripts/download_heap_prof_mapped_files.sh

@@ -19,18 +19,61 @@
 usage() {
   echo "This script downloads all of the files listed in the mappings section of a heap profile."
   echo ""
-  echo "Usage: $0 <heap_profile> <node_name> [<gcloud ssh opts>...]"
-  echo "Common gcloud ssh options include --project."
+  echo "Usage: $0 <heap_profile> <node_name> [--gcloud-common-args] [--gcloud-ssh-args...]"
+  echo "<heap_profile> : the path to the heap profile file (in pprof format) to analyze."
+  echo "<node_name> : the name of the node to connect to (e.g., the name of a Vizier agent node)."
+  echo "[--gcloud-common-args] : common arguments to pass to gcloud commands, such as --project."
+  echo "[--gcloud-ssh-args] : additional arguments to pass to gcloud compute ssh commands, such as --internal-ip."
   exit 1
 }
 
-heap_profile="$1"
-node_name="$2"
-output_dir="${heap_profile%.txt}"
+parse_args() {
+  if [ $# -lt 2 ]; then
+    usage
+  fi
+
+  heap_profile="$1"
+  shift
+
+  node_name="$1"
+  shift
+
+  while test $# -gt 0; do
+    case "$1" in
+      --gcloud-common-args=*)
+        GCLOUD_COMMON_ARGS="${1#*=}"
+        shift
+        ;;
+      --gcloud-common-args)
+        GCLOUD_COMMON_ARGS="$2"
+        shift
+        shift
+        ;;
+      --gcloud-ssh-args=*)
+        GCLOUD_SSH_ARGS="${1#*=}"
+        shift
+        ;;
+      --gcloud-ssh-args)
+        GCLOUD_SSH_ARGS="$2"
+        shift
+        shift
+        ;;
+      *) usage ;;
+    esac
+  done;
+}
+
 
-if [ -z "$heap_profile" ] || [ -z "$node_name" ]; then
-  usage
-fi
+check_args() {
+  if [ -z "$heap_profile" ] || [ -z "$node_name" ]; then
+    usage
+  fi
+}
+
+parse_args "${@}"
+check_args
+
+output_dir="${heap_profile%.txt}"
 
 # Create the output directory at the beginning of the script.
 mkdir -p "$output_dir"
@@ -44,16 +87,16 @@ mkdir -p "$output_dir"
 mappings=$(awk 'BEGIN{m=0} /MAPPED_LIBRARIES/{m=1} { if(m) { print $6 }}' "$heap_profile" | grep "^/" | sort | uniq)
 
 err_file="$output_dir/gcloud_error.log"
-zone=$(gcloud compute instances list "${@:3}" --filter="$node_name" --format="table(name, zone)"| tail -n 1 | awk '{print $2}')
-procs=$(gcloud compute ssh --zone "$zone" --command='ps ax' "$node_name" "${@:3}" 2> "$err_file") || cat "$err_file" && rm "$err_file"
+zone=$(gcloud compute instances list "${GCLOUD_COMMON_ARGS}" --filter="$node_name" --format="table(name, zone)"| tail -n 1 | awk '{print $2}')
+procs=$(gcloud compute ssh --zone "$zone" --command='ps ax' "$node_name" "${GCLOUD_COMMON_ARGS}" "${GCLOUD_SSH_ARGS}" 2> "$err_file") || cat "$err_file" && rm "$err_file"
 
 # Find the mapping that corresponds to a process on the node.
 # We assume that the process was started by running one of the files in the mappings
 # (which is commonly the case and certainly the case for the binaries we care about, eg. pem, kelvin).
 for fname in $mappings
 do
   bname=$(basename "$fname")
-  matching_proc=$(echo "$procs" | grep "$bname" || true)
+  matching_proc=$(echo "$procs" | grep -E "$bname$" || true)
   if [[ -n "$matching_proc" ]]; then
     pid=$(echo "$matching_proc" | awk '{ print $1 }')
     matched_file="$fname"
@@ -80,13 +123,13 @@ output_on_err() {
 }
 
 # Create tar archive on node.
-output_on_err gcloud compute ssh  --zone "$zone" --command="$create_tar_cmd" "$node_name" "${@:3}"
+output_on_err gcloud compute ssh  --zone "$zone" --command="$create_tar_cmd" "$node_name" "${GCLOUD_COMMON_ARGS}" "${GCLOUD_SSH_ARGS}"
 
 # Copy archive to local machine.
-output_on_err gcloud compute scp --zone "$zone" "${@:3}" "$USER@$node_name:~/$tar_file" "${output_dir}/$tar_file"
+output_on_err gcloud compute scp --zone "$zone" "${GCLOUD_COMMON_ARGS}" "${GCLOUD_SSH_ARGS}" "$USER@$node_name:~/$tar_file" "${output_dir}/$tar_file"
 
 # Cleanup tar archive on node.
-output_on_err gcloud compute ssh --zone "$zone" --command="rm ~/$tar_file" "$node_name" "${@:3}"
+output_on_err gcloud compute ssh --zone "$zone" --command="rm ~/$tar_file" "$node_name" "${GCLOUD_COMMON_ARGS}" "${GCLOUD_SSH_ARGS}"
 
 tar --strip-components=1 -C "$output_dir" -xzf "${output_dir}/$tar_file"
 

src/pxl_scripts/px/collect_heap_dumps.pxl

@@ -16,9 +16,13 @@
 
 import px
 
-df = px.GetAgentStatus(False)
+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 = df[['hostname', 'heap']]
+df.asid = df.asid_x
+df = df[['asid', 'hostname', 'heap']]
 px.display(df)