combine download_profile and xplane_to_metrics (GoogleCloudPlatform#744)

Author: shuningjin

xlml/apis/task.py

@@ -413,15 +413,13 @@ def post_process(self, result_location: Optional[str] = None) -> DAGNode:
       )
 
       if self.task_metric_config and self.task_metric_config.profile:
-        profile_tmp_dir = metric.download_profile(
-            process_id, self.task_metric_config.profile.file_location
-        )
-        self.task_metric_config.profile.metrics = metric.xplane_to_metrics(
-            profile_tmp_dir
+        self.task_metric_config.profile.metrics = (
+            metric.xplane_to_metrics.override(retries=0)(
+                self.task_metric_config.profile.file_location
+            )
         )
         (
             process_id
-            >> profile_tmp_dir
             >> self.task_metric_config.profile.metrics
             >> post_process_metrics
         )

xlml/utils/metric.py

@@ -347,96 +347,80 @@ def get_gcs_file_location_with_regex(
   return f"gs://{bucket_name}/{selected_blob.name}"
 
 
-def get_gcs_profile_locations(dir_location: str) -> str:
-  """
-  Find first gcs location matching regex `{dir_location}/.*/*xplane.pb`
-
-  Returns:
-    The matched gcs location, or "" if no match
-  """
-  storage_client = storage.Client()
-  url = urlparse(dir_location)
-  bucket_name = url.netloc
-  file_path = url.path.lstrip("/")
-  file_path_regex = re.compile(file_path + "/.*/*xplane.pb")
-  matched_blobs = [
-      b
-      for b in storage_client.list_blobs(bucket_name, prefix=file_path)
-      if file_path_regex.match(b.name)
-  ]
-  if not matched_blobs:
-    logging.warning(
-        f"No objects matched supplied regex: {dir_location}/.*/*xplane.pb"
-    )
-    return ""
-
-  selected_blob = matched_blobs[0]
-  return f"gs://{bucket_name}/{selected_blob.name}"
-
-
-@task
-def download_profile(
-    process_id: str, dir_location: str | airflow.XComArg
-) -> str:
-  """
-  Find the first file_location matching regex `{dir_location}/.*/*xplane.pb`
-  Create tmp_dir = f"/tmp/profile-{process_id}"
-  Download profile to tmp_location = f"{tmp_dir}/gke.xplane.pb"
-
-  Returns:
-    tmp_dir, or "" if no match
-  """
-  if isinstance(dir_location, airflow.XComArg):
-    dir_location = dir_location.resolve(get_current_context())
-  file_location = get_gcs_profile_locations(dir_location)
-  # no match profile
-  if not file_location:
-    return ""
-  print(f"For local download, run: gcloud storage cp {file_location} .")
-  # unique directory to accomodate intermediate files from profile extraction
-  tmp_dir = f"/tmp/profile-{process_id}"
-  os.makedirs(tmp_dir, exist_ok=False)
-  download_object_from_gcs(file_location, f"{tmp_dir}/gke.xplane.pb")
-  return tmp_dir
-
-
 @task.virtualenv(
     task_id="process_profile_metrics",
     requirements=["tensorboard_plugin_profile==2.19.4"],
-    system_site_packages=False,
+    system_site_packages=True,
 )
-def xplane_to_metrics(input_dir: str) -> dict:
+def xplane_to_metrics(dir_location: str | airflow.XComArg) -> dict:
   """
-  Extract metrics from xplane file `{input_dir}/gke.xplane.pb`.
-  Remove temporory input_dir generated by `download_profile`.
+  Find the first profile matching regex `{dir_location}/.*/*xplane.pb`.
+    Download profile to temporary directory.
+    Extract metrics from xplane file.
 
   Returns:
-    A dictionary of profile metrics, or None if input_dir is ""
+    A dictionary of profile metrics, or None if no match
   """
   # pylint: disable=redefined-outer-name, import-outside-toplevel, reimported
-  from tensorboard_plugin_profile.convert import raw_to_tool_data
   import json
   import logging
-  import shutil
+  import re
+  import tempfile
+  from urllib.parse import urlparse
+  import airflow
+  from google.cloud import storage
+  from tensorboard_plugin_profile.convert import raw_to_tool_data
   # pylint: enable=redefined-outer-name, import-outside-toplevel, reimported
 
-  if not input_dir:
-    return None
-  input_path = f"{input_dir}/gke.xplane.pb"
-  print(f"processing xplane file: {input_path}")
+  # --- Find and Download Profile ---
+  # pylint: disable-next=redefined-outer-name
+  def download_object_from_gcs(
+      source_location: str, destination_location: str
+  ) -> None:
+    """Download object from GCS bucket.
+
+    Args:
+      source_location: The full path of a file in GCS.
+      destination_location: The local path of the file.
+    """
+    storage_client = storage.Client()
+    bucket_name = source_location.split("/")[2]
+    object_name = "/".join(source_location.split("/")[3:])
+
+    bucket = storage_client.bucket(bucket_name)
+    blob = bucket.blob(object_name)
+    blob.download_to_filename(destination_location)
+    logging.info(
+        ("Download file from" f" {source_location} to {destination_location}")
+    )
 
-  def cleanup(input_dir):
-    # remove temporary directory generated by `download_profile`
-    if input_dir.startswith("/tmp/profile-"):
-      shutil.rmtree(input_dir)
-      print(f"remove temporary directory: {input_dir}")
+  def get_gcs_profile_location(dir_location: str) -> str:
+    """
+    Find first gcs location matching regex `{dir_location}/.*/*xplane.pb`.
+
+    Returns:
+      The matched gcs location, or "" if no match
+    """
+    storage_client = storage.Client()
+    url = urlparse(dir_location)
+    bucket_name = url.netloc
+    file_path = url.path.lstrip("/")
+    file_path_regex = re.compile(file_path + "/.*/*xplane.pb")
+    for b in storage_client.list_blobs(bucket_name, prefix=file_path):
+      if file_path_regex.match(b.name):
+        return f"gs://{bucket_name}/{b.name}"
+    logging.warning(
+        f"No objects matched supplied regex: {dir_location}/.*/*xplane.pb"
+    )
+    return ""
 
+  # --- Extract Metrics from Profile ---
   def round_number(number: float | str, decimal: int) -> float:
     if isinstance(number, str):
       number = float(number)
     return float(f"{number:.{decimal}f}")
 
-  def xplane_tool(input_path: str, tool: str, params: dict) -> dict:
+  def get_tool_data(input_path: str, tool: str, params: dict) -> dict:
     data, content_type = raw_to_tool_data.xspace_to_tool_data(
         xspace_paths=[input_path],
         tool=tool,
@@ -447,96 +431,111 @@ def xplane_tool(input_path: str, tool: str, params: dict) -> dict:
     parsed_data = json.loads(data)
     return parsed_data
 
-  out = {}
-  # check available tools
-  available_tools = set(
-      raw_to_tool_data.xspace_to_tool_names(xspace_paths=[input_path])
-  )
-  # 1 overview_page
-  # generate ALL_HOSTS.op_stats.pb
-  if "overview_page" not in available_tools:
-    logging.warning("overview_page: unavailable tool")
-  else:
+  def get_tool_metrics(input_path: str) -> dict:
+    out = {}
+    # check available tools
+    available_tools = set(
+        raw_to_tool_data.xspace_to_tool_names(xspace_paths=[input_path])
+    )
+    # 1 overview_page
+    # generate ALL_HOSTS.op_stats.pb
+    if "overview_page" not in available_tools:
+      logging.warning("overview_page: unavailable tool")
+    else:
+      try:
+        overview_page = get_tool_data(
+            input_path, tool="overview_page", params={}
+        )
+        out.update({
+            "Device Type": overview_page[2]["p"]["device_type"],
+            "Device Core Count": int(
+                overview_page[2]["p"]["device_core_count"]
+            ),
+            "Average Tensor Core Step Time (ms)": float(
+                overview_page[1]["p"]["steptime_ms_average"]
+            ),
+        })
+      except ValueError as e:
+        logging.warning(e)
+
+    # 2 op_profile
+    if "op_profile" not in available_tools:
+      logging.warning("op_profile: unavailable tool")
+      return out
     try:
-      overview_page = xplane_tool(input_path, tool="overview_page", params={})
+      op_profile = get_tool_data(input_path, tool="op_profile", params={})
       out.update({
-          "device_type": overview_page[2]["p"]["device_type"],
-          "device_core_count": int(overview_page[2]["p"]["device_core_count"]),
-          "Average Tensor Core Step Time (ms)": float(
-              overview_page[1]["p"]["steptime_ms_average"]
+          "TPU FLOPS Utilization (%): exclude_idle": round_number(
+              op_profile["byProgramExcludeIdle"]["metrics"]["flops"] * 100, 2
+          ),
+          "HBM Bandwidth Utilization (%): exclude_idle": round_number(
+              op_profile["byProgramExcludeIdle"]["metrics"]["bandwidthUtils"][0]
+              * 100,
+              2,
           ),
       })
     except ValueError as e:
       logging.warning(e)
-
-  # 2 op_profile
-  if "op_profile" not in available_tools:
-    logging.warning("op_profile: unavailable tool")
-    cleanup(input_dir)
-    return out
-  try:
-    op_profile = xplane_tool(input_path, tool="op_profile", params={})
-    out.update({
-        "TPU FLOPS Utilization (%): exclude_idle": round_number(
-            op_profile["byProgramExcludeIdle"]["metrics"]["flops"] * 100, 2
-        ),
-        "HBM Bandwidth Utilization (%): exclude_idle": round_number(
-            op_profile["byProgramExcludeIdle"]["metrics"]["bandwidthUtils"][0]
-            * 100,
-            2,
-        ),
-    })
-  except ValueError as e:
-    logging.warning(e)
-    cleanup(input_dir)
-    return out
-
-  # 3 memory_viewer: jit_train_step
-  # generate jit*.hlo_proto.pb
-  if "memory_viewer" not in available_tools:
-    logging.warning("memory_viewer: unavailable tool")
-    cleanup(input_dir)
-    return out
-  # find jit_train_step from 2
-  # example: "jit_train_step(4869159985936022652)"
-  jit_train_step = None
-  for program in op_profile["byProgramExcludeIdle"]["children"]:
-    if program["name"].startswith("jit_train_step"):
-      jit_train_step = program["name"]
-      break
-  if jit_train_step is None:
-    logging.warning("memory_viewer: cannot find jit_train_step in op_profile")
-    cleanup(input_dir)
+      return out
+
+    # 3 memory_viewer: jit_train_step
+    # generate jit*.hlo_proto.pb
+    if "memory_viewer" not in available_tools:
+      logging.warning("memory_viewer: unavailable tool")
+      return out
+    # find jit_train_step from 2
+    # example: "jit_train_step(4869159985936022652)"
+    jit_train_step = None
+    for program in op_profile["byProgramExcludeIdle"]["children"]:
+      if program["name"].startswith("jit_train_step"):
+        jit_train_step = program["name"]
+        break
+    if jit_train_step is None:
+      logging.warning("memory_viewer: cannot find jit_train_step in op_profile")
+      return out
+    # 3.1 hbm
+    MEMORY_SPACE_HBM = "0"
+    try:
+      memory_viewer_hbm = get_tool_data(
+          input_path,
+          "memory_viewer",
+          params={"host": jit_train_step, "memory_space": MEMORY_SPACE_HBM},
+      )
+      out["Peak memory allocation (MiB): jit_train_step, HBM"] = round_number(
+          memory_viewer_hbm["totalBufferAllocationMib"], 2
+      )
+    except ValueError as e:
+      logging.warning(e)
+    # 3.2 host
+    MEMORY_SPACE_HOST = "5"
+    try:
+      memory_viewer_host = get_tool_data(
+          input_path,
+          "memory_viewer",
+          params={"host": jit_train_step, "memory_space": MEMORY_SPACE_HOST},
+      )
+      out["Peak memory allocation (MiB): jit_train_step, host"] = round_number(
+          memory_viewer_host["totalBufferAllocationMib"], 2
+      )
+    except ValueError as e:
+      logging.warning(e)
     return out
-  # 3.1 hbm
-  MEMORY_SPACE_HBM = "0"
-  try:
-    memory_viewer_hbm = xplane_tool(
-        input_path,
-        "memory_viewer",
-        params={"host": jit_train_step, "memory_space": MEMORY_SPACE_HBM},
-    )
-    out["Peak memory allocation (MiB): jit_train_step, HBM"] = round_number(
-        memory_viewer_hbm["totalBufferAllocationMib"], 2
-    )
-  except ValueError as e:
-    logging.warning(e)
-  # 3.2 host
-  MEMORY_SPACE_HOST = "5"
-  try:
-    memory_viewer_host = xplane_tool(
-        input_path,
-        "memory_viewer",
-        params={"host": jit_train_step, "memory_space": MEMORY_SPACE_HOST},
-    )
-    out["Peak memory allocation (MiB): jit_train_step, host"] = round_number(
-        memory_viewer_host["totalBufferAllocationMib"], 2
-    )
-  except ValueError as e:
-    logging.warning(e)
 
-  cleanup(input_dir)
-  return out
+  # --- Main Logic ---
+  # Find the first file_location matching regex `{dir_location}/.*/*xplane.pb`
+  if isinstance(dir_location, airflow.XComArg):
+    dir_location = dir_location.resolve(get_current_context())
+  file_location = get_gcs_profile_location(dir_location)
+  if file_location:
+    print(f"For local download, run: gcloud storage cp {file_location} .")
+    # Temporary directory for file download and extraction cache
+    with tempfile.TemporaryDirectory() as tmp_dir:
+      input_path = f"{tmp_dir}/gke.xplane.pb"
+      download_object_from_gcs(file_location, input_path)
+      # Extract profile
+      return get_tool_metrics(input_path)
+  # No match profile
+  return None
 
 
 def process_profile(