From 8d795dd676567a5427f4a05b6114a4824b3e8d33 Mon Sep 17 00:00:00 2001
From: Tarry Singh <tarry.singh@deepkapha.com>
Date: Thu, 9 May 2024 11:09:29 +0000
Subject: [PATCH] pjrt: profiler support std writer API

expose a more low-level function to customize where to write profile
reports
---
 pjrt/profiler.zig | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/pjrt/profiler.zig b/pjrt/profiler.zig
index 1b0d1a6..77f1520 100644
--- a/pjrt/profiler.zig
+++ b/pjrt/profiler.zig
@@ -126,6 +126,19 @@ pub const Profiler = struct {
         allocator: std.mem.Allocator,
         dir: std.fs.Dir,
         file_name: []const u8,
+    ) !void {
+        log.info("Writing profiling data to {s}", .{file_name});
+        var output_file = try dir.createFile(file_name, .{});
+        defer output_file.close();
+        var buffered_writer = std.io.bufferedWriter(output_file.writer());
+        try self.dumpAsJsonToWriter(allocator, buffered_writer.writer());
+        try buffered_writer.flush();
+    }
+
+    pub fn dumpAsJsonToWriter(
+        self: *Profiler,
+        allocator: std.mem.Allocator,
+        writer: anytype,
     ) !void {
         const profile_data = try self.collectData(allocator);
         defer profile_data.free(allocator);
@@ -139,12 +152,7 @@ pub const Profiler = struct {
         defer converter.deinit();
         try converter.parseXSpaceBytes(profile_data.items(), 1_000_000);
 
-        var output_file = try dir.createFile(file_name, .{});
-        defer output_file.close();
-        var buffered_writer = std.io.bufferedWriter(output_file.writer());
-        log.info("Writing profiling data to {s}", .{file_name});
-        try converter.toJson(buffered_writer.writer());
-        try buffered_writer.flush();
+        try converter.toJson(writer);
     }
 
     fn check(self: *Profiler, c_error: ?*Error) !void {