Create parse_xla_logs.py

Author: juliagmt-google

.github/workflows/parse_xla_logs.py

@@ -0,0 +1,120 @@
+import re
+import sys
+
+class BenchmarkResult:  # Define the class here
+    def __init__(self):
+        self.name = ""
+        self.flavor = ""
+        self.description = ""
+        self.xla_flags = ""
+        self.build_flags = []
+        self.source_uri = ""
+        self.device_name = ""
+        self.time_us = 0.0
+        self.last_average_loss = 0.0
+        self.memory_read_written = 0
+        self.buffer_allocation = []
+        self.nvptx_compilation_time_us = 0
+        self.hlo_passes_time_us = 0
+        self.run_backend_time_us = 0
+        self.device_time_us = 0.0
+        self.device_memcpy_time_us = 0.0
+        self.xprof_session_id = ""
+        self.error = ""
+        self.custom_metrics = []
+
+class LogParser:
+    def __init__(self):
+        self.compilation_time = {}  # You'll need to initialize this appropriately
+    # Units defined in google3/third_party/tensorflow/tsl/platform/numbers.cc
+    _TIME_UNITS = {"us": 1e-6, "ms": 1e-3, "s": 1, "min": 60, "h": 3600}
+    _TIME_REGEXP = re.compile(r"time: (\d+\.?\d*) (%s)" % "|".join(_TIME_UNITS))
+    _SIZE_UNITS = "KMGTPE"
+    _SIZE_REGEXP = re.compile(r"(\d+\.?\d*)([%s]i)?B" % _SIZE_UNITS)
+
+    def _ParseTimeFromLog(self, time_str: bytes, position: int = -1) -> int:
+      """Returns the time in microseconds parsed from XLA logs."""
+      match = self._TIME_REGEXP.search(time_str.decode())
+      assert match, "Unable to parse the time on log line %d" % position
+      exp_ = self._TIME_UNITS[match.group(2)] * 1e6
+      return int(float(match.group(1)) * exp_)
+
+    def _ParseBytesFromLog(self, size_str: bytes, position: int = -1) -> int:
+      """Returns the size in bytes parsed from XLA logs."""
+      match = self._SIZE_REGEXP.search(size_str.decode())
+      assert match, "Unable to parse the size on log line %d" % position
+      prefix = match.group(2) or "-"
+      exp_ = pow(2, 10 * (self._SIZE_UNITS.find(prefix[0]) + 1))
+      return int(float(match.group(1)) * exp_)
+
+    def _ParseLogLine(self, result, line, position) -> None:
+    """Parses a single XLA log line and updates the BenchmarkResult proto.
+
+    It looks for certain text patterns and updates the result proto with
+    compilation stats and Xprof session id.
+
+    Args:
+      result: a benchmark result proto to be updated
+      line: a single line of benchmark log output
+      position: line number (for debugging)
+    """
+      # Log output generated by --vmodule=nvptx_compiler=1
+      if b"NVPTXCompiler::CompileTargetBinary - CompileToPtx" in line:
+        thread_id = re.split(rb"\s+", line)[2]
+        self.compilation_time[thread_id] += self._ParseTimeFromLog(line, position)
+        return
+  
+      # Log output generated by --vmodule=gpu_compiler=1
+      if b"HLO memory read+written:" in line:
+        result.memory_read_written += self._ParseBytesFromLog(line, position)
+        return
+      if b"GpuCompiler::RunHloPasses for" in line:
+        result.hlo_passes_time_us += self._ParseTimeFromLog(line, position)
+        return
+      if b"GpuCompiler::RunBackend for" in line:
+        result.run_backend_time_us += self._ParseTimeFromLog(line, position)
+        return
+  
+      # Log output generated by --vmodule=bfc_allocator=2
+      if b"New Peak memory usage" in line and b"for GPU" in line:
+        match = re.search(rb"(\d+) bytes", line)
+        assert match, "Unable to parse the size on log line %d" % position
+        alloc_size = int(match.group(1))
+        if alloc_size > max(result.buffer_allocation, default=0):
+          result.buffer_allocation[:] = [alloc_size]
+        return
+  
+      # Log output generated by --xprof_end_2_end_upload
+      if b"XprofResponse uploaded to http://xprof/" in line:
+        match = re.search(rb"session_id=([\w\-]+)", line)
+        assert match, "Unable to parse the XProf link"
+        result.xprof_session_id = match.group(1)
+      return
+
+    def parse_log_file(self, log_file_path):
+        result = BenchmarkResult()
+        with open(log_file_path, "rb") as f:
+            for i, line in enumerate(f):
+                self._ParseLogLine(result, line, i + 1)
+
+        # Update compilation time (total across threads) in the result object
+        for time_taken in self.compilation_time.values():
+            result.nvptx_compilation_time_us += time_taken
+        return result  # Return the result object
+
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        print("Usage: python parse_log.py <log_file_path>")
+        sys.exit(1)
+
+    log_file_path = sys.argv[1]
+
+    parser = LogParser()
+    result = parser.parse_log_file(log_file_path)
+
+    # Access results
+    print(f"Memory Read+Written: {result.memory_read_written} bytes")
+    print(f"NVPTX Compilation Time: {result.nvptx_compilation_time_us} us")
+    print(f"HLO Passes Time: {result.hlo_passes_time_us} us")
+    print(f"Run Backend Time: {result.run_backend_time_us} us")
+    print(f"Buffer Allocation: {result.buffer_allocation} bytes")