diff --git a/tensorflow/core/util/reporter.cc b/tensorflow/core/util/reporter.cc index 02687095c9c..eb69e292116 100644 --- a/tensorflow/core/util/reporter.cc +++ b/tensorflow/core/util/reporter.cc @@ -21,25 +21,57 @@ limitations under the License. namespace tensorflow { -TestReporter::TestReporter(const string& fname, const string& test_name) +TestReportFile::TestReportFile(const string& fname, const string& test_name) : closed_(true), fname_(fname), test_name_(test_name) {} -Status TestReporter::Close() { +Status TestReportFile::Append(const string& content) { if (closed_) return Status::OK(); + return log_file_->Append(content); +} + +Status TestReportFile::Close() { + if (closed_) return Status::OK(); + closed_ = true; + return log_file_->Close(); +} + +Status TestReportFile::Initialize() { + if (fname_.empty()) { + return Status::OK(); + } + string mangled_fname = strings::StrCat( + fname_, absl::StrJoin(str_util::Split(test_name_, '/'), "__")); + Env* env = Env::Default(); + if (env->FileExists(mangled_fname).ok()) { + return errors::InvalidArgument( + "Cannot create TestReportFile, file exists: ", mangled_fname); + } + TF_RETURN_IF_ERROR(env->NewWritableFile(mangled_fname, &log_file_)); + TF_RETURN_IF_ERROR(log_file_->Flush()); + + closed_ = false; + return Status::OK(); +} + +TestReporter::TestReporter(const string& fname, const string& test_name) + : report_file_(fname, test_name) { + benchmark_entry_.set_name(test_name); +} + +Status TestReporter::Close() { + if (report_file_.IsClosed()) return Status::OK(); BenchmarkEntries entries; *entries.add_entry() = benchmark_entry_; - TF_RETURN_IF_ERROR(log_file_->Append(entries.SerializeAsString())); - + TF_RETURN_IF_ERROR(report_file_.Append(entries.SerializeAsString())); benchmark_entry_.Clear(); - closed_ = true; - return log_file_->Close(); + return report_file_.Close(); } Status TestReporter::Benchmark(int64 iters, double cpu_time, double wall_time, double throughput) { - if (closed_) return Status::OK(); + if (report_file_.IsClosed()) return Status::OK(); benchmark_entry_.set_iters(iters); benchmark_entry_.set_cpu_time(cpu_time / iters); benchmark_entry_.set_wall_time(wall_time / iters); @@ -48,34 +80,17 @@ Status TestReporter::Benchmark(int64 iters, double cpu_time, double wall_time, } Status TestReporter::SetProperty(const string& name, const string& value) { - if (closed_) return Status::OK(); + if (report_file_.IsClosed()) return Status::OK(); (*benchmark_entry_.mutable_extras())[name].set_string_value(value); return Status::OK(); } Status TestReporter::SetProperty(const string& name, double value) { - if (closed_) return Status::OK(); + if (report_file_.IsClosed()) return Status::OK(); (*benchmark_entry_.mutable_extras())[name].set_double_value(value); return Status::OK(); } -Status TestReporter::Initialize() { - if (fname_.empty()) { - return Status::OK(); - } - string mangled_fname = strings::StrCat( - fname_, absl::StrJoin(str_util::Split(test_name_, '/'), "__")); - Env* env = Env::Default(); - if (env->FileExists(mangled_fname).ok()) { - return errors::InvalidArgument("Cannot create TestReporter, file exists: ", - mangled_fname); - } - TF_RETURN_IF_ERROR(env->NewWritableFile(mangled_fname, &log_file_)); - TF_RETURN_IF_ERROR(log_file_->Flush()); - - benchmark_entry_.set_name(test_name_); - closed_ = false; - return Status::OK(); -} +Status TestReporter::Initialize() { return report_file_.Initialize(); } } // namespace tensorflow diff --git a/tensorflow/core/util/reporter.h b/tensorflow/core/util/reporter.h index e551e2e4f57..51d7502701c 100644 --- a/tensorflow/core/util/reporter.h +++ b/tensorflow/core/util/reporter.h @@ -29,6 +29,34 @@ limitations under the License. namespace tensorflow { +// The TestReportFile provides a file abstraction for TF tests to use. +class TestReportFile { + public: + // Create a TestReportFile with the test name 'test_name'. + TestReportFile(const string& fname, const string& test_name); + + // Initialize the TestReportFile. If the reporting env flag is set, + // try to create the reporting file. Fails if the file already exists. + Status Initialize(); + + // Append the report file w/ 'content'. + Status Append(const string& content); + + // Close the report file. + Status Close(); + + bool IsClosed() const { return closed_; } + + ~TestReportFile() { Close().IgnoreError(); } // Autoclose in destructor. + + private: + bool closed_; + string fname_; + string test_name_; + std::unique_ptr<WritableFile> log_file_; + TF_DISALLOW_COPY_AND_ASSIGN(TestReportFile); +}; + // The TestReporter writes test / benchmark output to binary Protobuf files when // the environment variable "TEST_REPORT_FILE_PREFIX" is defined. // @@ -91,10 +119,7 @@ class TestReporter { const char* fname_ptr = getenv(kTestReporterEnv); return (fname_ptr != nullptr) ? fname_ptr : ""; } - bool closed_; - string fname_; - string test_name_; - std::unique_ptr<WritableFile> log_file_; + TestReportFile report_file_; BenchmarkEntry benchmark_entry_; TF_DISALLOW_COPY_AND_ASSIGN(TestReporter); };