diff --git a/tsl/profiler/lib/BUILD b/tsl/profiler/lib/BUILD
index bcca03f6f..ac34b6b77 100644
--- a/tsl/profiler/lib/BUILD
+++ b/tsl/profiler/lib/BUILD
@@ -213,6 +213,7 @@ cc_library(
         "//tsl/profiler/protobuf:profiler_options_proto_cc",
         "//tsl/profiler/protobuf:xplane_proto_cc",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/synchronization",
         "@xla//xla/tsl/platform:errors",
         "@xla//xla/tsl/platform:status",
@@ -242,6 +243,8 @@ cc_library(
         "//tsl/profiler/protobuf:xplane_proto_cc",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/synchronization",
         "@xla//xla/tsl/platform:env",
         "@xla//xla/tsl/platform:errors",
diff --git a/tsl/profiler/lib/profiler_session.cc b/tsl/profiler/lib/profiler_session.cc
index 2f5352f47..460f0ece6 100644
--- a/tsl/profiler/lib/profiler_session.cc
+++ b/tsl/profiler/lib/profiler_session.cc
@@ -15,12 +15,14 @@ limitations under the License.
 
 #include "tsl/profiler/lib/profiler_session.h"
 
+#include <any>
 #include <cstdint>
 #include <memory>
 #include <utility>
 
 #include "absl/memory/memory.h"
 #include "absl/status/status.h"
+#include "absl/status/status_macros.h"
 #include "absl/synchronization/mutex.h"
 #include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/logging.h"
@@ -80,7 +82,7 @@ absl::Status ProfilerSession::Status() {
 #if !defined(IS_MOBILE_PLATFORM)
 absl::Status ProfilerSession::CollectDataInternal(XSpace* space) {
   absl::MutexLock l(mutex_);
-  TF_RETURN_IF_ERROR(status_);
+  RETURN_IF_ERROR(status_);
   LOG(INFO) << "Profiler session collecting data.";
   if (profilers_ != nullptr) {
     profilers_->Stop().IgnoreError();
@@ -97,7 +99,7 @@ absl::Status ProfilerSession::CollectDataInternal(XSpace* space) {
 absl::Status ProfilerSession::CollectData(XSpace* space) {
 #if !defined(IS_MOBILE_PLATFORM)
   space->add_hostnames(port::Hostname());
-  TF_RETURN_IF_ERROR(CollectDataInternal(space));
+  RETURN_IF_ERROR(CollectDataInternal(space));
   profiler::SetXSpacePidIfNotSet(*space, tsl::Env::Default()->GetProcessId());
   profiler::PostProcessSingleHostXSpace(space, start_time_ns_, stop_time_ns_);
 #endif
@@ -105,6 +107,40 @@ absl::Status ProfilerSession::CollectData(XSpace* space) {
   return absl::OkStatus();
 }
 
+absl::StatusOr<profiler::ConsumeResult> ProfilerSession::Consume() {
+#if !defined(IS_MOBILE_PLATFORM)
+  absl::MutexLock l(mutex_);
+  RETURN_IF_ERROR(status_);
+  LOG(INFO) << "Profiler session consuming data.";
+  if (profilers_ == nullptr) {
+    return absl::FailedPreconditionError("Profiler session is not active.");
+  }
+  return profilers_->Consume();
+#else
+  return absl::UnimplementedError("Consume not implemented on mobile.");
+#endif
+}
+
+absl::Status ProfilerSession::Serialize(std::any data,
+                                        tensorflow::profiler::XSpace* space) {
+#if !defined(IS_MOBILE_PLATFORM)
+  absl::MutexLock l(mutex_);
+  RETURN_IF_ERROR(status_);
+  if (profilers_ == nullptr) {
+    return absl::FailedPreconditionError("Profiler session is not active.");
+  }
+  absl::Status status = profilers_->Serialize(std::move(data), space);
+  if (!status.ok()) return status;
+
+  space->add_hostnames(port::Hostname());
+  profiler::SetXSpacePidIfNotSet(*space, tsl::Env::Default()->GetProcessId());
+  profiler::PostProcessSingleHostXSpace(space, start_time_ns_,
+                                        profiler::GetCurrentTimeNanos());
+#endif
+  SetProfileOptionsIntoSpace(options_, space);
+  return absl::OkStatus();
+}
+
 ProfilerSession::ProfilerSession(const ProfileOptions& options)
 #if defined(IS_MOBILE_PLATFORM)
     : status_(errors::Unimplemented(
diff --git a/tsl/profiler/lib/profiler_session.h b/tsl/profiler/lib/profiler_session.h
index 6644017f5..79a14f566 100644
--- a/tsl/profiler/lib/profiler_session.h
+++ b/tsl/profiler/lib/profiler_session.h
@@ -15,11 +15,13 @@ limitations under the License.
 #ifndef TENSORFLOW_TSL_PROFILER_LIB_PROFILER_SESSION_H_
 #define TENSORFLOW_TSL_PROFILER_LIB_PROFILER_SESSION_H_
 
+#include <any>
 #include <functional>
 #include <memory>
 #include <vector>
 
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/synchronization/mutex.h"
 #include "xla/tsl/platform/status.h"
 #include "xla/tsl/platform/types.h"
@@ -67,6 +69,13 @@ class ProfilerSession {
   absl::Status CollectData(tensorflow::profiler::XSpace* space)
       TF_LOCKS_EXCLUDED(mutex_);
 
+  // Consumes collected profile data and returns a container of consumed data.
+  absl::StatusOr<profiler::ConsumeResult> Consume() TF_LOCKS_EXCLUDED(mutex_);
+
+  // Serializes consumed data into XSpace.
+  absl::Status Serialize(std::any data, tensorflow::profiler::XSpace* space)
+      TF_LOCKS_EXCLUDED(mutex_);
+
  private:
   // Constructs an instance of the class and starts profiling
   explicit ProfilerSession(const tensorflow::ProfileOptions& options);