No public description

PiperOrigin-RevId: 715008416

model/orbax/model/core/python/file_utils.py

@@ -0,0 +1,30 @@
+# Copyright 2024 The Orbax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""File utilities."""
+
+import contextlib
+
+
+_file_opener = open
+
+
+@contextlib.contextmanager
+def open_file(filename: str, mode: str):
+  """Opens a file with the given filename and mode."""
+  f = _file_opener(filename, mode)
+  try:
+    yield f
+  finally:
+    f.close()

model/orbax/model/core/python/save.py

@@ -25,6 +25,7 @@
 from absl import logging
 from orbax.experimental.model.core.checkpoint.tensor_bundle import _tensor_bundle_api
 from orbax.experimental.model.core.python import constants
+from orbax.experimental.model.core.python import file_utils
 from orbax.experimental.model.core.python import manifest_constants
 from orbax.experimental.model.core.python import module
 from orbax.experimental.model.core.python import unstructured_data
@@ -168,7 +169,7 @@ def save(
     manifest_proto = build_manifest_proto(
         m, path, supplemental_info=supplemental_info,
     )
-    with open(
+    with file_utils.open_file(
         os.path.join(path, manifest_constants.MANIFEST_FILENAME), 'wb'
     ) as f:
       f.write(manifest_proto.SerializeToString())
@@ -190,7 +191,9 @@ def save(
           'function_aliases is set but no named function exists in the graph.'
       )
   proto = builder.build()
-  with open(os.path.join(path, constants.SAVED_MODEL_FILENAME_PB), 'wb') as f:
+  with file_utils.open_file(
+      os.path.join(path, constants.SAVED_MODEL_FILENAME_PB), 'wb'
+  ) as f:
     f.write(proto.SerializeToString())
 
   # Write checkpoint. The `builder.variables_by_name` to used to ensure

model/orbax/model/core/python/unstructured_data.py

@@ -18,6 +18,7 @@
 import os
 
 from orbax.experimental.model.core.protos.manifest_pb2 import UnstructuredData  # pylint: disable=g-importing-member
+from orbax.experimental.model.core.python import file_utils
 
 
 @dataclasses.dataclass
@@ -77,7 +78,7 @@ def write_inlined_data_to_file(
   else:
     data_to_write = proto.inlined_bytes
     io_mode = "wb"
-  with open(os.path.join(dirname, filename), io_mode) as f:
+  with file_utils.open_file(os.path.join(dirname, filename), io_mode) as f:
     f.write(data_to_write)
   result = UnstructuredData()
   result.file_system_location.string_path = filename