[torchlib] [ci] Fix failing external data tests for Windows

onnxscript/ir/_core.py

@@ -671,6 +671,13 @@ def tobytes(self) -> bytes:
         length = self._length or self.nbytes
         return self.raw[offset : offset + length]
 
+    def release(self) -> None:
+        """Delete all references to the memory buffer and close the memory-mapped file."""
+        self._array = None
+        if self.raw is not None:
+            self.raw.close()
+            self.raw = None
+
     @property
     def metadata_props(self) -> dict[str, str]:
         if self._metadata_props is None:

onnxscript/ir/_core_test.py

@@ -244,6 +244,26 @@ def test_initialize(self):
         # Ensure repeated reads are consistent
         np.testing.assert_equal(tensor, self.data)
 
+    def test_release_does_not_invalidate_tensor(self):
+        external_tensor = self.model.graph.initializer[0]
+        external_info = onnx.external_data_helper.ExternalDataInfo(external_tensor)
+        tensor = _core.ExternalTensor(
+            external_info.location,
+            offset=external_info.offset,
+            length=external_info.length,
+            dtype=ir.DataType.FLOAT,
+            base_dir=self.base_path,
+            name="input",
+            shape=_core.Shape(external_tensor.dims),
+        )
+        self.assertEqual(tensor.dtype, ir.DataType.FLOAT)
+        self.assertEqual(tensor.tobytes(), self.data.tobytes())
+        # Release tensor
+        tensor.release()
+        self.assertEqual(tensor.raw, None)
+        # Tensor can be re-loaded after release
+        self.assertEqual(tensor.tobytes(), self.data.tobytes())
+
     def test_initialize_with_relative_path(self):
         external_tensor = self.model.graph.initializer[0]
         external_info = onnx.external_data_helper.ExternalDataInfo(external_tensor)

onnxscript/ir/_external_data.py

@@ -100,6 +100,7 @@ def _load_external_data_file(
             if os.path.samefile(tensor.path, os.path.join(base_path, relative_path)):
                 # Copy the data as the .numpy() call references data from a file whose data is eventually modified
                 tensor_data = external_tensor.numpy().copy()
+                external_tensor.release()
                 tensor = _core.Tensor(
                     tensor_data, name=external_tensor.name, dtype=external_tensor.dtype
                 )
@@ -165,6 +166,8 @@ def _save_external_data(
             current_offset = tensor_info.offset
             assert tensor is not None
             raw_data = tensor.tobytes()
+            if isinstance(tensor, _core.ExternalTensor):
+                tensor.release()
             # Pad file to required offset if needed
             file_size = data_file.tell()
             if current_offset > file_size:
@@ -223,6 +226,7 @@ def convert_tensors_to_external(
     path = os.path.join(base_path, relative_path)
     # Check if file path is valid, and create subsequent subdirectories within the path if they don't exist
     os.makedirs(os.path.dirname(path), exist_ok=True)
+    tmp_file_created = False
     # Check if file exists. Load pre-existing external data if it does.
     if os.path.exists(path):
         # Check if any tensor in the model is using the destination file
@@ -241,6 +245,7 @@ def convert_tensors_to_external(
                 os.makedirs(tmp_path, exist_ok=True)
                 # If exisiting external tensors are not loaded to memory, copy the external data to a temporary location
                 os.rename(path, os.path.join(tmp_path, relative_path))
+                tmp_file_created = True
                 for tensor in tensors:
                     if (
                         isinstance(tensor, _core.ExternalTensor)
@@ -270,6 +275,12 @@ def convert_tensors_to_external(
         external_tensors[i]
         for i in sorted(range(len(external_tensors)), key=lambda i: sorted_indices[i])
     ]
+
+    # Clean-up temporary file if it is created
+    tmp_path = os.path.join(base_path, "tmp", relative_path)
+    if os.path.exists(tmp_path) and tmp_file_created:
+        os.remove(tmp_path)
+
     return external_tensors
 
 

onnxscript/ir/_external_data_test.py

@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 import os
+import sys
 import tempfile
 import typing
 import unittest
@@ -115,7 +116,10 @@ class OffloadExternalTensorTest(unittest.TestCase):
 
     def setUp(self):
         # File paths
-        self.temp_dir = tempfile.TemporaryDirectory()  # pylint: disable=consider-using-with
+        if sys.version_info[:2] >= (3, 10):
+            self.temp_dir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)  # pylint: disable=consider-using-with
+        else:
+            self.temp_dir = tempfile.TemporaryDirectory()  # pylint: disable=consider-using-with
         self.external_data_name = "external_tensors.bin"
         self.base_path = self.temp_dir.name
         self.ext_data_1 = "external_data_1.bin"
@@ -136,7 +140,15 @@ def setUp(self):
         self.model_with_mixed_external_data = self._model_with_mixed_external_data()
 
     def tearDown(self) -> None:
-        self.temp_dir.cleanup()
+        # Handle exceptions for windows and python versions < 3.10
+        try:
+            self.temp_dir.cleanup()
+        except PermissionError as e:
+            print(f"PermissionError: {e}")
+        except FileNotFoundError as e:
+            print(f"FileNotFoundError: {e}")
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            print(f"An unexpected error occurred: {e}")
 
     def _simple_model(self) -> ir.Model:
         tensor1 = ir.Tensor(