fix(stream_io): Finalize temporary files correctly on __exit__

The temporary file finalizer's check that the file is not already closed
was causing the upload to be skipped if the stream was used as a context
manager without an explicit call to .close(), because CPython's
implementation of .__exit__() for NamedTemporaryFile objects closes
the underlying file before calling the wrapper's own .close() method.
After changing to `weakref.finalize` in commit 3fd7f1e, uploading
temporary files became inherently idempotent, so the check that the file
is not already closed is no longer necessary anyway.
This change deletes that check.

Furthermore, changes to the NamedTemporaryFile implementation with the
addition of its delete_on_close parameter in Python 3.12
(python/cpython#97015) make its .__exit__() method no longer call
its .close() method at all, so this change also embeds the finalizer
as a hook in a wrapper around both .close() and .__exit__() separately.

Author: Eta0

CHANGELOG.md

@@ -5,6 +5,21 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+
+### Fixed
+
+- File objects opened with `stream_io.open_stream("s3://...", "wb")` for writing
+  to object storage now correctly upload their content when closed implicitly
+  at the end of a `with` block, without requiring an explicit call to their
+  `.close()` method
+  - Since `TensorSerializer` objects already call `.close()` explicitly on
+    their output file objects, either when `TensorSerializer.close()` is invoked
+    or when the `TensorSerializer` is garbage collected, this bug mainly applies
+    to manual usage of `stream_io.open_stream()` for object storage uploads
+    not involving a `TensorSerializer`
+
+
 ## [2.7.1] - 2023-12-06
 
 ### Fixed
@@ -265,7 +280,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   - `get_gpu_name`
   - `no_init_or_tensor`
 
-[2.7.0]: https://github.com/coreweave/tensorizer/compare/v2.7.0...v2.7.1
+[Unreleased]: https://github.com/coreweave/tensorizer/compare/v2.7.1...HEAD
+[2.7.1]: https://github.com/coreweave/tensorizer/compare/v2.7.0...v2.7.1
 [2.7.0]: https://github.com/coreweave/tensorizer/compare/v2.6.0...v2.7.0
 [2.6.0]: https://github.com/coreweave/tensorizer/compare/v2.5.1...v2.6.0
 [2.5.1]: https://github.com/coreweave/tensorizer/compare/v2.5.0...v2.5.1

tensorizer/_version.py

@@ -1 +1 @@
-__version__ = "2.7.1"
+__version__ = "2.7.2"

tensorizer/stream_io.py

@@ -1053,9 +1053,9 @@ def _infer_credentials(
 def _temp_file_closer(file: io.IOBase, file_name: str, *upload_args):
     """
     Close, upload by name, and then delete the file.
-    Meant to replace .close() on a particular instance
-    of a temporary file-like wrapper object, as an unbound
-    callback to a weakref.finalize() registration on the wrapper.
+    Meant to be placed as a hook before both .close() and .__exit__()
+    on a particular instance of a temporary file-like wrapper object,
+    as a callback to a weakref.finalize() registration on the wrapper.
 
     The reason this implementation is necessary is really complicated.
 
@@ -1077,17 +1077,6 @@ def _temp_file_closer(file: io.IOBase, file_name: str, *upload_args):
     so they have to buffer it all in memory.
     """
 
-    if file.closed:
-        # Makes closure idempotent.
-
-        # If the file object is used as a context
-        # manager, close() is called twice (once in the
-        # serializer code, once after, when leaving the
-        # context).
-
-        # Without this check, this would trigger two
-        # separate uploads.
-        return
     try:
         file.close()
         s3_upload(file_name, *upload_args)
@@ -1281,6 +1270,9 @@ def open_stream(
             # with primitive temporary file support (e.g. Windows)
             temp_file = tempfile.NamedTemporaryFile(mode="wb+", delete=False)
 
+            # Attach a callback to upload the temporary file when it closes.
+            # weakref finalizers are idempotent, so this upload callback
+            # is guaranteed to run at most once.
             guaranteed_closer = weakref.finalize(
                 temp_file,
                 _temp_file_closer,
@@ -1291,7 +1283,35 @@ def open_stream(
                 s3_secret_access_key,
                 s3_endpoint,
             )
-            temp_file.close = guaranteed_closer
+
+            # Always run the close + upload procedure
+            # before any code from Python's NamedTemporaryFile wrapper.
+            # It isn't safe to call a bound method from a weakref finalizer,
+            # but calling a weakref finalizer alongside a bound method
+            # creates no problems, other than that the code outside the
+            # finalizer is not guaranteed to be run at any point.
+            # In this case, the weakref finalizer performs all necessary
+            # cleanup itself, but the original NamedTemporaryFile methods
+            # are invoked as well, just in case.
+            wrapped_close = temp_file.close
+
+            def close_wrapper():
+                guaranteed_closer()
+                return wrapped_close()
+
+            # Python 3.12+ doesn't call NamedTemporaryFile.close() during
+            # .__exit__(), so it must be wrapped separately.
+            # Since guaranteed_closer is idempotent, it's fine to call it in
+            # both methods, even if both are called back-to-back.
+            wrapped_exit = temp_file.__exit__
+
+            def exit_wrapper(exc, value, tb):
+                guaranteed_closer()
+                return wrapped_exit(exc, value, tb)
+
+            temp_file.close = close_wrapper
+            temp_file.__exit__ = exit_wrapper
+
             return temp_file
         else:
             s3_endpoint = s3_endpoint or default_s3_read_endpoint