Switch to json for snapshot metadata (#120)

Summary:
Pull Request resolved: #120

When the number of entries in the snapshot metadata is very large, yaml
serialization becomes a bottleneck and there's little we can do to optimize.
Fortunately, when the metadata is serialized via json it is still valid yaml.
Technically, the snapshot metadata format is still yaml.

Reviewed By: raypeng

Differential Revision: D40626156

fbshipit-source-id: 013e81fd42dd4dd1debf9579c4bee854374e1ca2

Author: Yifu Wang

tests/test_manifest.py

@@ -5,13 +5,14 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import json
 from dataclasses import asdict
 from typing import Dict, Generator
 from unittest.mock import patch
 
 import pytest
 
+import yaml
+
 from _pytest.fixtures import SubRequest  # @manual
 
 from torchsnapshot.manifest import (
@@ -28,6 +29,12 @@
 )
 from torchsnapshot.manifest_ops import _insert_entry, get_manifest_for_rank
 
+try:
+    from yaml import CSafeDumper as Dumper
+except ImportError:
+    from yaml import SafeDumper as Dumper
+
+
 _WORLD_SIZE = 2
 _MANIFEST_0: Dict[str, Entry] = {
     "0/foo": DictEntry(
@@ -227,17 +234,15 @@
 @pytest.fixture(params=[True, False])
 def use_cyaml(request: SubRequest) -> Generator[None, None, None]:
     if request.param:
-        from yaml import CSafeDumper, CSafeLoader
+        from yaml import CSafeLoader
 
-        with patch("torchsnapshot.manifest.Dumper", CSafeDumper):
-            with patch("torchsnapshot.manifest.Loader", CSafeLoader):
-                yield
+        with patch("torchsnapshot.manifest.Loader", CSafeLoader):
+            yield
     else:
-        from yaml import SafeDumper, SafeLoader
+        from yaml import SafeLoader
 
-        with patch("torchsnapshot.manifest.Dumper", SafeDumper):
-            with patch("torchsnapshot.manifest.Loader", SafeLoader):
-                yield
+        with patch("torchsnapshot.manifest.Loader", SafeLoader):
+            yield
 
 
 @pytest.mark.usefixtures("use_cyaml")
@@ -255,26 +260,19 @@ def test_manifest_yaml_serialization(manifest: Dict[str, Entry]) -> None:
 
 @pytest.mark.usefixtures("use_cyaml")
 @pytest.mark.parametrize("manifest", [_MANIFEST_0, _MANIFEST_1])
-def test_manifest_json_serialization(manifest: Dict[str, Entry]) -> None:
+def test_manifest_yaml_dumper(manifest: Dict[str, Entry]) -> None:
     """
-    Verify that when the metadata is serialized via json, it is load-able with
-    the yaml loader.
-
-    When the number of entries in the snapshot metadata is very large, yaml
-    serialization becomes a bottleneck and there's little we can do to
-    optimize. We likely need to switch to json to overcome this. Fortunately,
-    when our metadata is serialized via json, it is compatible with the yaml
-    loader, so we can make the switch in a backward compatible fashion. This
-    test makes sure that we don't do anything crazy with the metadata to break
-    this compatibility.
+    :func:`SnapshotMetadata.to_yaml` switched to :func:`json.dumps`` to help
+    with the serialization performance. This test verifies that old snapshot
+    metadata serialized with :func:`yaml.dump` are still loadable.
     """
     metadata = SnapshotMetadata(
         version="0.0.0",
         world_size=_WORLD_SIZE,
         manifest=manifest,
     )
-    yaml_str = metadata.to_yaml()
-    json_str = json.dumps(asdict(metadata))
+    yaml_str = yaml.dump(asdict(metadata), sort_keys=False, Dumper=Dumper)
+    json_str = metadata.to_yaml()
     metadata_from_yaml = SnapshotMetadata.from_yaml(yaml_str=yaml_str)
     metadata_from_json = SnapshotMetadata.from_yaml(yaml_str=json_str)
     assert metadata_from_json == metadata_from_yaml

torchsnapshot/manifest.py

@@ -8,6 +8,7 @@
 # pyre-ignore-all-errors[2]: Allow `Any` in type annotations
 
 import base64
+import json
 import logging
 import struct
 from dataclasses import asdict, dataclass
@@ -17,9 +18,9 @@
 import yaml
 
 try:
-    from yaml import CSafeDumper as Dumper, CSafeLoader as Loader
+    from yaml import CSafeLoader as Loader
 except ImportError:
-    from yaml import SafeDumper as Dumper, SafeLoader as Loader
+    from yaml import SafeLoader as Loader
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -280,7 +281,12 @@ class SnapshotMetadata:
     manifest: Manifest
 
     def to_yaml(self) -> str:
-        return yaml.dump(asdict(self), sort_keys=False, Dumper=Dumper)
+        # When the number of entries in the snapshot metadata is large, yaml
+        # serialization becomes slow and there's little room for optimization.
+        # Since the snapshot metadata can be dumped as json and json is a
+        # subset of yaml, using json.dumps() here to help with the
+        # serialization performance without needing to deprecate yaml.
+        return json.dumps(asdict(self), sort_keys=False, indent=2)
 
     @classmethod
     def from_yaml(cls, yaml_str: str) -> "SnapshotMetadata":

torchsnapshot/snapshot.py

@@ -31,24 +31,17 @@
 from .dist_store import get_or_create_store, LinearBarrier
 
 from .flatten import flatten, inflate
-from .io_preparer import (
-    ObjectBufferConsumer,
-    prepare_read,
-    prepare_write,
-    TensorIOPreparer,
-)
+from .io_preparer import prepare_read, prepare_write
 from .io_types import ReadIO, ReadReq, StoragePlugin, WriteIO, WriteReq
 from .knobs import is_batching_disabled
 
 from .manifest import (
-    ChunkedTensorEntry,
     Entry,
     is_container_entry,
     Manifest,
     PrimitiveEntry,
     ShardedTensorEntry,
     SnapshotMetadata,
-    TensorEntry,
 )
 from .manifest_ops import get_manifest_for_rank
 from .partitioner import consolidate_replicated_entries, partition_write_reqs