Skip to content

Commit

Permalink
Add file name and size to the serialization metadata logging (#113077)
Browse files Browse the repository at this point in the history
Summary:

To be able to get more info on serialization/deserialization events, adding these two files to the metadata logging.
- file_name
- file_size

Test Plan: buck2 test mode/dev caffe2/caffe2/serialize:inline_container_test

Reviewed By: davidberard98

Differential Revision: D51040426
  • Loading branch information
atannous authored and facebook-github-bot committed Nov 7, 2023
1 parent 65304d8 commit e732f7f
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 10 deletions.
19 changes: 12 additions & 7 deletions caffe2/serialize/inline_container.cc
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
#include <thread>

#include <c10/core/Allocator.h>
#include <c10/core/Backend.h>
#include <c10/core/CPUAllocator.h>
#include <c10/core/Backend.h>
#include <c10/util/Exception.h>
Expand Down Expand Up @@ -88,7 +89,7 @@ static std::string basename(const std::string& name) {

static std::string parentdir(const std::string& name) {
size_t end = name.find_last_of('/');
if(end == std::string::npos) {
if (end == std::string::npos) {
end = name.find_last_of('\\');
}

Expand Down Expand Up @@ -179,7 +180,9 @@ void PyTorchStreamReader::init() {
}
c10::LogAPIUsageMetadata(
"pytorch.stream.reader.metadata",
{{"serialization_id", serialization_id_}});
{{"serialization_id", serialization_id_},
{"file_name", archive_name_},
{"file_size", str(mz_zip_get_archive_size(ar_.get()))}});

// version check
at::DataPtr version_ptr;
Expand Down Expand Up @@ -700,8 +703,8 @@ void PyTorchStreamWriter::writeEndOfFile() {
~Finalizer() {
var_ = true;
}
private:
bool& var_;
private:
bool& var_;
} f(finalized_);

auto allRecords = getAllWrittenRecords();
Expand Down Expand Up @@ -736,6 +739,11 @@ void PyTorchStreamWriter::writeEndOfFile() {
mz_zip_writer_finalize_archive(ar_.get());
mz_zip_writer_end(ar_.get());
valid("writing central directory for archive ", archive_name_.c_str());
c10::LogAPIUsageMetadata(
"pytorch.stream.writer.metadata",
{{"serialization_id", serialization_id_},
{"file_name", archive_name_},
{"file_size", str(mz_zip_get_archive_size(ar_.get()))}});
if (file_stream_.is_open()) {
file_stream_.close();
}
Expand Down Expand Up @@ -779,9 +787,6 @@ void PyTorchStreamWriter::writeSerializationId() {
kSerializationIdRecordName,
serialization_id_.c_str(),
serialization_id_.size());
c10::LogAPIUsageMetadata(
"pytorch.stream.writer.metadata",
{{"serialization_id", serialization_id_}});
}
}

Expand Down
9 changes: 7 additions & 2 deletions caffe2/serialize/inline_container_test.cc
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
#include <array>
#include <cstdio>
#include <cstring>
#include <string>

#include <gtest/gtest.h>
Expand Down Expand Up @@ -408,9 +409,13 @@ TEST(PytorchStreamWriterAndReader, LogAPIUsageMetadata) {
ASSERT_EQ(logs.size(), 2);
std::map<std::string, std::map<std::string, std::string>> expected_logs = {
{"pytorch.stream.writer.metadata",
{{"serialization_id", writer.serializationId()}}},
{{"serialization_id", writer.serializationId()},
{"file_name", "archive"},
{"file_size", str(oss.str().length())}}},
{"pytorch.stream.reader.metadata",
{{"serialization_id", writer.serializationId()}}}
{{"serialization_id", writer.serializationId()},
{"file_name", "archive"},
{"file_size", str(iss.str().length())}}}
};
ASSERT_EQ(expected_logs, logs);

Expand Down
5 changes: 4 additions & 1 deletion torch/package/package_importer.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,10 @@ def __init__(

torch._C._log_api_usage_metadata(
"torch.package.PackageImporter.metadata",
{"serialization_id": self.zip_reader.serialization_id()},
{
"serialization_id": self.zip_reader.serialization_id(),
"file_name": self.filename,
},
)

self.root = _PackageNode(None)
Expand Down

0 comments on commit e732f7f

Please sign in to comment.