Skip to content

Commit 6c2d25d

Browse files
sannidhyachauhancopybara-github
sannidhyachauhan
authored and
copybara-github
committed
Add Consume and Serialize methods to ProfilerCollection and ProfilerSession, and introduce ProfilerSessionOrchestrator.
PiperOrigin-RevId: 891616851
1 parent 4fbda33 commit 6c2d25d

11 files changed

Lines changed: 408 additions & 9 deletions

‎tsl/profiler/lib/BUILD‎

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ load(
1212
"tf_profiler_copts",
1313
"tf_profiler_pybind_cc_library_wrapper",
1414
)
15+
load("//third_party/bazel_rules/rules_cc/cc:cc_library.bzl", "cc_library")
1516

1617
# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
1718

@@ -140,7 +141,7 @@ cc_library(
140141
]),
141142
deps = [
142143
"//tsl/profiler/protobuf:xplane_proto_cc",
143-
"@xla//xla/tsl/platform:status",
144+
"@com_google_absl//absl/status",
144145
],
145146
)
146147

@@ -378,7 +379,6 @@ cc_library(
378379
":profiler_interface",
379380
"//tsl/profiler/protobuf:xplane_proto_cc",
380381
"@com_google_absl//absl/status",
381-
"@xla//xla/tsl/platform:status",
382382
],
383383
)
384384

@@ -391,3 +391,36 @@ cc_library(
391391
"@com_google_absl//absl/strings:string_view",
392392
],
393393
)
394+
395+
cc_library(
396+
name = "profiler_orchestrator",
397+
srcs = ["profiler_orchestrator.cc"],
398+
hdrs = ["profiler_orchestrator.h"],
399+
visibility = internal_visibility([
400+
"@xla//xla/python:__pkg__",
401+
]),
402+
deps = [
403+
":profiler_session",
404+
"//tsl/profiler/protobuf:profiler_options_proto_cc",
405+
"//tsl/profiler/protobuf:xplane_proto_cc",
406+
"@com_google_absl//absl/status",
407+
"@com_google_absl//absl/status:statusor",
408+
"@xla//xla/tsl/platform:errors",
409+
"@xla//xla/tsl/platform:logging",
410+
],
411+
)
412+
413+
tsl_cc_test(
414+
name = "profiler_orchestrator_test",
415+
srcs = ["profiler_orchestrator_test.cc"],
416+
deps = [
417+
":profiler_factory",
418+
":profiler_interface",
419+
":profiler_orchestrator",
420+
":profiler_session",
421+
"//tsl/profiler/protobuf:profiler_options_proto_cc",
422+
"@com_google_absl//absl/memory",
423+
"@com_google_absl//absl/status",
424+
"@xla//xla/tsl/platform:test",
425+
],
426+
)

‎tsl/profiler/lib/profiler_collection.cc‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,5 +55,28 @@ absl::Status ProfilerCollection::CollectData(
5555
return status;
5656
}
5757

58+
absl::Status ProfilerCollection::Consume(void* ptr) {
59+
absl::Status status;
60+
for (auto& profiler : profilers_) {
61+
absl::Status s = profiler->Consume(ptr);
62+
if (!s.ok() && !absl::IsUnimplemented(s)) {
63+
status.Update(s);
64+
}
65+
}
66+
return status;
67+
}
68+
69+
absl::Status ProfilerCollection::Serialize(
70+
void* ptr, tensorflow::profiler::XSpace* output_space) {
71+
absl::Status status;
72+
for (auto& profiler : profilers_) {
73+
absl::Status s = profiler->Serialize(ptr, output_space);
74+
if (!s.ok() && !absl::IsUnimplemented(s)) {
75+
status.Update(s);
76+
}
77+
}
78+
return status;
79+
}
80+
5881
} // namespace profiler
5982
} // namespace tsl

‎tsl/profiler/lib/profiler_collection.h‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,11 @@ limitations under the License.
1515
#ifndef TENSORFLOW_TSL_PROFILER_LIB_PROFILER_COLLECTION_H_
1616
#define TENSORFLOW_TSL_PROFILER_LIB_PROFILER_COLLECTION_H_
1717

18+
#include <cstddef>
1819
#include <memory>
1920
#include <vector>
2021

2122
#include "absl/status/status.h"
22-
#include "xla/tsl/platform/status.h"
2323
#include "tsl/profiler/lib/profiler_interface.h"
2424
#include "tsl/profiler/protobuf/xplane.pb.h"
2525

@@ -38,6 +38,9 @@ class ProfilerCollection : public ProfilerInterface {
3838
absl::Status Stop() override;
3939

4040
absl::Status CollectData(tensorflow::profiler::XSpace* space) override;
41+
absl::Status Consume(void* ptr) override;
42+
absl::Status Serialize(void* ptr,
43+
tensorflow::profiler::XSpace* output_space) override;
4144

4245
private:
4346
std::vector<std::unique_ptr<ProfilerInterface>> profilers_;

‎tsl/profiler/lib/profiler_controller.cc‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,5 +85,14 @@ absl::Status ProfilerController::CollectData(
8585
return status;
8686
}
8787

88+
absl::Status ProfilerController::Consume(void* ptr) {
89+
return profiler_->Consume(ptr);
90+
}
91+
92+
absl::Status ProfilerController::Serialize(
93+
void* ptr, tensorflow::profiler::XSpace* output_space) {
94+
return profiler_->Serialize(ptr, output_space);
95+
}
96+
8897
} // namespace profiler
8998
} // namespace tsl

‎tsl/profiler/lib/profiler_controller.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ limitations under the License.
1515
#ifndef TENSORFLOW_TSL_PROFILER_LIB_PROFILER_CONTROLLER_H_
1616
#define TENSORFLOW_TSL_PROFILER_LIB_PROFILER_CONTROLLER_H_
1717

18+
#include <cstddef>
1819
#include <memory>
1920

2021
#include "absl/status/status.h"
@@ -45,6 +46,11 @@ class ProfilerController : public ProfilerInterface {
4546

4647
absl::Status CollectData(tensorflow::profiler::XSpace* space) override;
4748

49+
absl::Status Consume(void* ptr) override;
50+
51+
absl::Status Serialize(void* ptr,
52+
tensorflow::profiler::XSpace* output_space) override;
53+
4854
private:
4955
enum class ProfilerState {
5056
kInit = 0,

‎tsl/profiler/lib/profiler_interface.h‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ limitations under the License.
1515
#ifndef TENSORFLOW_TSL_PROFILER_LIB_PROFILER_INTERFACE_H_
1616
#define TENSORFLOW_TSL_PROFILER_LIB_PROFILER_INTERFACE_H_
1717

18-
#include "xla/tsl/platform/status.h"
18+
#include "absl/status/status.h"
1919
#include "tsl/profiler/protobuf/xplane.pb.h"
2020

2121
namespace tsl {
@@ -41,6 +41,18 @@ class ProfilerInterface {
4141

4242
// Saves collected profile data into XSpace.
4343
virtual absl::Status CollectData(tensorflow::profiler::XSpace* space) = 0;
44+
45+
// Pulls collected profile data into arbitrary raw memory.
46+
// Size refers to the amount of data memory `ptr` points to.
47+
virtual absl::Status Consume(void* ptr) {
48+
return absl::UnimplementedError("Consume not implemented");
49+
}
50+
51+
// Serializes collected profile data into XSpace.
52+
virtual absl::Status Serialize(void* ptr,
53+
tensorflow::profiler::XSpace* output_space) {
54+
return absl::UnimplementedError("Serialize not implemented");
55+
}
4456
};
4557

4658
} // namespace profiler

‎tsl/profiler/lib/profiler_orchestrator.cc‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
#include "tsl/profiler/lib/profiler_orchestrator.h"
16+
17+
#include <cstdint>
18+
#include <memory>
19+
#include <vector>
20+
21+
#include "absl/status/status.h"
22+
#include "xla/tsl/platform/errors.h"
23+
#include "xla/tsl/platform/logging.h"
24+
#include "tsl/profiler/lib/profiler_session.h"
25+
#include "tsl/profiler/protobuf/profiler_options.pb.h"
26+
27+
namespace tsl {
28+
namespace profiler {
29+
30+
ProfilerSessionOrchestrator::ProfilerSessionOrchestrator(
31+
const tensorflow::ProfileOptions& options)
32+
: options_(options) {}
33+
34+
ProfilerSessionOrchestrator::~ProfilerSessionOrchestrator() {
35+
Stop().IgnoreError();
36+
}
37+
38+
absl::Status ProfilerSessionOrchestrator::Start() {
39+
if (session_ != nullptr) {
40+
return absl::FailedPreconditionError("Session already started.");
41+
}
42+
session_ = tsl::ProfilerSession::Create(options_);
43+
if (session_ == nullptr) {
44+
return absl::InternalError("Failed to create ProfilerSession.");
45+
}
46+
return session_->Status();
47+
}
48+
49+
absl::Status ProfilerSessionOrchestrator::Stop() {
50+
if (session_ == nullptr) {
51+
return absl::OkStatus(); // Already stopped or not started.
52+
}
53+
session_.reset();
54+
return absl::OkStatus();
55+
}
56+
57+
absl::StatusOr<int> ProfilerSessionOrchestrator::Consume() {
58+
if (session_ == nullptr) {
59+
return absl::FailedPreconditionError("Session not started.");
60+
}
61+
62+
consume_buffers_.emplace_back(kDefaultConsumeSize);
63+
auto& buffer = consume_buffers_.back();
64+
TF_RETURN_IF_ERROR(session_->Consume(buffer.data()));
65+
return consume_buffers_.size() - 1;
66+
}
67+
68+
absl::Status ProfilerSessionOrchestrator::Serialize(int buffer_index) {
69+
if (session_ == nullptr) {
70+
return absl::FailedPreconditionError("Session not started.");
71+
}
72+
73+
if (buffer_index < 0 || buffer_index >= consume_buffers_.size()) {
74+
return absl::InvalidArgumentError("Invalid buffer index.");
75+
}
76+
serialize_space_.Clear();
77+
auto& buffer = consume_buffers_[buffer_index];
78+
return session_->Serialize(buffer.data(), &serialize_space_);
79+
}
80+
81+
void ProfilerSessionOrchestrator::ClearConsumeBuffers() {
82+
std::vector<std::vector<uint8_t>>().swap(consume_buffers_);
83+
}
84+
85+
} // namespace profiler
86+
} // namespace tsl

‎tsl/profiler/lib/profiler_orchestrator.h‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
#ifndef TENSORFLOW_TSL_PROFILER_LIB_PROFILER_ORCHESTRATOR_H_
16+
#define TENSORFLOW_TSL_PROFILER_LIB_PROFILER_ORCHESTRATOR_H_
17+
18+
#include <cstdint>
19+
#include <memory>
20+
#include <vector>
21+
22+
#include "absl/status/status.h"
23+
#include "absl/status/statusor.h"
24+
#include "tsl/profiler/lib/profiler_session.h"
25+
#include "tsl/profiler/protobuf/profiler_options.pb.h"
26+
#include "tsl/profiler/protobuf/xplane.pb.h"
27+
28+
namespace tsl {
29+
namespace profiler {
30+
31+
class ProfilerSessionOrchestrator {
32+
public:
33+
explicit ProfilerSessionOrchestrator(
34+
const tensorflow::ProfileOptions& options);
35+
~ProfilerSessionOrchestrator();
36+
37+
absl::Status Start();
38+
39+
absl::StatusOr<int> Consume();
40+
41+
absl::Status Serialize(int buffer_index);
42+
43+
absl::Status Stop();
44+
45+
void ClearConsumeBuffers();
46+
47+
const std::vector<uint8_t>& GetConsumeBuffer(int index) const {
48+
return consume_buffers_[index];
49+
}
50+
const tensorflow::profiler::XSpace& GetSerializeSpace() const {
51+
return serialize_space_;
52+
}
53+
54+
private:
55+
tensorflow::ProfileOptions options_;
56+
std::unique_ptr<tsl::ProfilerSession> session_;
57+
std::vector<std::vector<uint8_t>> consume_buffers_;
58+
tensorflow::profiler::XSpace serialize_space_;
59+
};
60+
61+
} // namespace profiler
62+
} // namespace tsl
63+
64+
#endif // TENSORFLOW_TSL_PROFILER_LIB_PROFILER_ORCHESTRATOR_H_

0 commit comments

Comments
 (0)