From 06b46f746901da3b024c1475df8509abe06466fa Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sun, 15 Nov 2026 02:05:52 -0800 Subject: [PATCH 1/2] Make save and save_pte_program raise when a write fails `ExecutorchProgramManager.save` caught any error from the write, logged it and returned normally. `save_pte_program` did the same for the `.pte` and the `.ptd` files, and then returned the path of a file that might be empty or cut short. A caller had no way to notice, short of reading the log. Now both let the error reach the caller. `save` still logs the path when it succeeds. `save_pte_program` now writes the `.pte` through `save` instead of repeating its code, so both paths behave the same way. Behavior change for callers: - `save` and `save_pte_program` raise on a failed write, where they used to log and return. - `save_pte_program` needs an object with `save`, as its type hint already says (`ExecutorchProgramManager`). The legacy `ExecutorchProgram` that the deprecated `exir.capture(...).to_edge().to_executorch()` path returns has no `save`, so passing it now raises `AttributeError` instead of writing the file. Such code can call the legacy object's own `write_to_file` instead. Test plan: `exir/program/test/test_program.py`. New tests check that `save` raises when the write fails, and that `save_pte_program` writes the `.pte` and `.ptd` files, returns the path, and raises when either write fails. Both fail on the code before this change. --- exir/program/_program.py | 9 +++---- exir/program/test/BUCK | 1 + exir/program/test/test_program.py | 39 +++++++++++++++++++++++++++++++ extension/export_util/utils.py | 11 ++------- 4 files changed, 45 insertions(+), 15 deletions(-) diff --git a/exir/program/_program.py b/exir/program/_program.py index e03b27b09ec..927976a82d7 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -2064,9 +2064,6 @@ def save(self, path: str) -> None: if path[-4:] != ".pte": logging.error(f"Path {path} does not end with .pte") raise ValueError(f"Path {path} does not end with .pte") - try: - with open(path, "wb") as file: - self.write_to_file(file) - logging.info(f"Saved exported program to {path}") - except Exception as e: - logging.error(f"Error while saving to {path}: {e}") + with open(path, "wb") as file: + self.write_to_file(file) + logging.info(f"Saved exported program to {path}") diff --git a/exir/program/test/BUCK b/exir/program/test/BUCK index f316eb83f6a..eddfdf89e96 100644 --- a/exir/program/test/BUCK +++ b/exir/program/test/BUCK @@ -24,6 +24,7 @@ fbcode_target(_kind = runtime.python_test, "//executorch/exir/program:lib", "//executorch/exir/tests:lib", "//executorch/exir/tests:models", + "//executorch/extension/export_util:export_util", "//executorch/extension/pybindings:portable_lib", # @manual ], ) diff --git a/exir/program/test/test_program.py b/exir/program/test/test_program.py index 413498166ec..916d8efd7a8 100644 --- a/exir/program/test/test_program.py +++ b/exir/program/test/test_program.py @@ -7,11 +7,15 @@ # pyre-unsafe import copy +import os +import tempfile import unittest from typing import Any, Dict +from unittest import mock import torch from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig +from executorch.exir._serialize._cord import Cord from executorch.exir.backend.test.op_partitioner_demo import ( AddMulPartitionerDemo, NonDecompTestPartitioner, @@ -31,6 +35,7 @@ from executorch.exir.tracer import _default_decomposition_table from executorch.exir.verification.verifier import EXIREdgeDialectVerifier +from executorch.extension.export_util.utils import save_pte_program from executorch.extension.pybindings.portable_lib import ( _load_for_executorch_from_buffer, ) @@ -1196,6 +1201,40 @@ def test_save_fails(self): with self.assertRaises(ValueError): _ = et.save("/tmp/test_save.pt") + def test_save_raises_when_the_write_fails(self): + model = TestLinear() + et = to_edge( + torch.export.export(model, model._get_random_inputs(), strict=True) + ).to_executorch() + with tempfile.TemporaryDirectory() as directory: + with mock.patch.object( + Cord, "write_to_file", side_effect=OSError("disk full") + ): + with self.assertRaises(OSError): + et.save(os.path.join(directory, "model.pte")) + + def test_save_pte_program_writes_the_data_files_and_raises_on_failure(self): + model = TestLinear() + et = to_edge( + torch.export.export(model, model._get_random_inputs(), strict=True) + ).to_executorch() + et._tensor_data = {"data": Cord(b"data")} + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "model.pte") + self.assertEqual(path, save_pte_program(et, path, directory)) + with open(os.path.join(directory, "data.ptd"), "rb") as f: + self.assertEqual(b"data", f.read()) + failures = { + "pte": (ExecutorchProgramManager, "write_to_file"), + "ptd": (ExecutorchProgramManager, "write_tensor_data_to_file"), + } + for name, (owner, method) in failures.items(): + with self.subTest(write=name), mock.patch.object( + owner, method, side_effect=OSError("disk full") + ): + with self.assertRaises(OSError): + save_pte_program(et, path, directory) + def test__transform_override_verifiers(self): """Test that _transform can override verifiers in the exported program.""" diff --git a/extension/export_util/utils.py b/extension/export_util/utils.py index e7534b22448..13248343bc2 100644 --- a/extension/export_util/utils.py +++ b/extension/export_util/utils.py @@ -141,14 +141,7 @@ def save_pte_program( else: filename = os.path.join(output_dir, f"{model_name}.pte") - try: - # Write program to file. - with open(filename, "wb") as file: - prog.write_to_file(file) - logging.info(f"Saved exported program to {filename}") - # Write data to file/s. - prog.write_tensor_data_to_file(outdir=output_dir) - except Exception as e: - logging.error(f"Error while saving to {filename}: {e}") + prog.save(filename) + prog.write_tensor_data_to_file(outdir=output_dir) return filename From d7106cd16127a57c4f1e37b04f3ca14b9cfdb5fe Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sun, 15 Nov 2026 07:44:40 -0800 Subject: [PATCH 2/2] Let save move file-backed data into the output instead of copying it `FileBackedData` keeps a large blob, such as a delegate's compiled engine, in a file until the program is written. Writing it copies the file into the `.pte` or `.ptd`. So until the caller deletes the file afterwards, the blob is on disk twice. For an engine of tens of gigabytes, that copy alone can make an export run out of disk. This adds a keyword-only `consume=True` to `ExecutorchProgramManager.save` and `write_tensor_data_to_file`. With it, data that owns its file (`cleanup=True`, as from `move_from`) is moved into the output: - Chunks are copied from the end of the data back to the start, each to its final offset. The source is truncated after each chunk, once the chunk was written in full, flushed and synced. - The part of the output not yet written stays a hole (a sparse file), so for a new output the disk holds the moved data once, plus one chunk. On a file system without sparse files the peak is the same as a copy. - It trades time for space: the sync per chunk makes a move several times slower than a copy (3 to 9 times in two measurements on btrfs). - Which buffers move is decided before the output is opened. Only when some do is the output a temporary file next to the final one (`..partial`), renamed onto the final name only when the write succeeds. A save that fails or is killed partway leaves the final name as it was, never a full-length file with missing data. If data was already moved, a note on the error names the temporary file, which holds it, and says so when only the rename failed; otherwise the temporary file is deleted. Such data can't be written again, so the export has to be rerun. An old file at the final name keeps its space until the rename, so overwriting one saves no space. The rename gives the final name a new file, owned by the caller's user and group. A symlink at the final name is followed, and the file it names keeps its mode; a hard link at the final name is replaced, as any rename replaces it. - Moved data can't be read or written again, by the program, by the lowered program it came from, or by a deep copy of either, so any ETRecord has to be saved first. Reading a `FileBackedData` whose file lost bytes raises, because the check is on the file's size, not a flag. `save` and `write_tensor_data_to_file` run that check before they open any output, so writing again to the same path can't wipe the only copy. They also refuse to write over the data's own file. - A buffer that deduplication put in both the `.pte` and a `.ptd`, or twice in one file, is copied, because no single write knows it is the last one to need it. - A deep copy of a `FileBackedData` is the same object, so a program made from a deep copy of an edge program still owns its data and can move it. - From its first truncation until the move is done, the data's file is not deleted when the object is collected, so an interrupt can't lose the part still in it. `consume` is offered only where the library opens the output itself. `write_to_file` methods, which take a file the caller opened, always copy as before. The source file is opened once, without following links, and judged on what was opened. Anything else gets a copy: data that does not own its file, a link, a file with another hard link or that is not writable, a source on another file system than the output (where the move would free nothing), and Windows, where files are not sparse by default. When some buffers are copied, one warning gives their number and the reasons. Data moves into a new file at the final name or into an existing writable regular file. With nothing to move, anywhere else (a pipe or a device, for example), when the final name can't be checked, or when no temporary file can be created next to it, the save is a plain one, exactly as without `consume`: same file, no temporary file. Where data moves, `save` writes the program itself and does not call an overridden `write_to_file`. A move that fails before its first truncation leaves the data intact. One that fails later leaves that source file in place and adds a note to the error that says which bytes are still in it and where the rest is in the temporary file. If earlier buffers or `.ptd` files were already done, notes list them, since they now hold the only copy of what moved into them. Other API changes: `FileBackedData` gains `check()` and a `path` property, and `Cord` gains `file_backed_buffers()`. Test plan: `exir/_serialize/test/test_cord.py` and `exir/program/test/test_program.py`. New tests check: - a test backend whose data is file-backed, lowered through `to_edge_transform_and_lower`: a consuming save and `.ptd` write, in several chunks, empty its file, and the deserialized `.pte` (or `.ptd` when external) holds the data; the same from a deep copy of the edge program; - the source still holds each range when it reaches the output, and allocated blocks stay well under twice the data (a copy's peak) while the move runs; - a consuming save moves owned data into the `.pte`, and a consuming `.ptd` write after it moves more; a save without `consume` copies and can be repeated; - a consuming save or `.ptd` write stopped partway leaves the final name as it was (absent, or the old file), the note names the temporary file, and one stopped before anything moved leaves no temporary file; a failed multi-file `.ptd` write lists the files already written; a consuming save follows a symlink and keeps the file's mode, and its temporary file never allows more than that mode; it copies into a pipe, into a read-only file (which raises, as without `consume`), into a directory where no temporary file can be created, and where a check of the final name fails (a check a save without `consume` never makes); on Windows it writes the final name directly; with nothing to move it writes the old file in place, without a sync; it warns once about the buffers it copies, and not when all move; a failed rename keeps the complete temporary file and the note says so; - writing again to the same path after a consuming write, also through `save_pte_program`, raises and leaves the first output whole; writing over the data's own file raises, with or without `consume`, and so does a move into the source file itself; - `consume` can't be passed by position; - the copy fallbacks, each checked on both the source and the output: unowned data, a link (whose target stays whole), a second hard link, a read-only file, another file system, Windows, and buffers the write was not asked to move; - a short write, a short read, a failed sync, a failed truncation, a failure on a later chunk (the note's offsets are checked against the bytes), an interrupt right after a truncation, a failed final seek, and a failure in a second buffer after the first was moved; the data's file is not left to be deleted while the move shrinks it; - every read of a moved or changed file raises; a buffer shared by the `.pte` and a `.ptd`, in either order, and one repeated in one file are copied and stay whole; - `save` with nothing to move still calls a subclass's own `write_to_file`, and a moving one does not. The tests of the move itself are skipped on Windows, which always copies. `test_named_data_store.py` and `exir/tests/test_serde.py` also pass. --- exir/_serialize/_cord.py | 204 ++++++++++- exir/_serialize/test/test_cord.py | 420 +++++++++++++++++++++++ exir/program/_program.py | 233 ++++++++++++- exir/program/test/BUCK | 1 + exir/program/test/test_program.py | 541 +++++++++++++++++++++++++++++- 5 files changed, 1373 insertions(+), 26 deletions(-) diff --git a/exir/_serialize/_cord.py b/exir/_serialize/_cord.py index c4bd4415ae9..6644860dc8b 100644 --- a/exir/_serialize/_cord.py +++ b/exir/_serialize/_cord.py @@ -8,13 +8,31 @@ import io import os import shutil +import stat +import sys import tempfile import weakref -from typing import List, Optional, Union +from typing import Container, Dict, IO, List, Optional, Union class FileBackedData: - """A byte buffer that stays on disk until explicitly closed.""" + """A byte buffer that stays on disk until explicitly closed or moved. + + ``ExecutorchProgramManager.save(path, consume=True)`` moves data that owns its file + (``cleanup=True``, as from ``move_from``) into the file it writes instead of copying + it. Chunks are copied from the end back to the start, and the file is truncated after + each chunk is written and synced to the output; the part of the output not yet + written stays a hole. On a file system with sparse files, the data is then on disk + once, plus one chunk; on one without, the move needs as much space as a copy. The + syncs make a move several times slower than a copy. The move needs a regular file + that is not a link and has no other name, on the output's file system. Anything else, + and Windows (where files are not sparse by default), gets a copy. After a move the + file is empty, and reading this or any other ``FileBackedData`` on it raises. + + A move that fails before its first truncation leaves the data intact. One that fails + later leaves the file in place and says, in a note on the error, which bytes are + still in it and where the rest is in the output. + """ _COPY_CHUNK_SIZE = 8 * 1024 * 1024 @@ -56,27 +74,166 @@ def move_from(cls, path: str, sha256: Optional[bytes] = None) -> "FileBackedData def __len__(self) -> int: return self._size + def __deepcopy__(self, memo: Dict[int, object]) -> "FileBackedData": + # One object per file: of two that owned it, the first collected would delete the + # file under the other, and one that doesn't own it can't move it. + return self + + @property + def path(self) -> str: + return self._path + + def _check_size(self, size: int) -> None: + # The size, not a flag, shows a moved file, so every object on it sees it. + if size != self._size: + raise RuntimeError( + f"{self._path} has {size} bytes, expected {self._size}; it changed after " + "it was wrapped, for example because a consuming save moved its data" + ) + + def check(self) -> None: + """Raise if the file no longer holds the data, as after a move.""" + self._check_size(os.stat(self._path).st_size) + + def _open(self) -> IO[bytes]: + f = open(self._path, "rb") + try: + self._check_size(os.fstat(f.fileno()).st_size) + except BaseException: + f.close() + raise + return f + def prefix(self, size: int) -> bytes: - with open(self._path, "rb") as f: + with self._open() as f: return f.read(size) def sha256(self) -> bytes: if self._sha256 is None: digest = hashlib.sha256() - with open(self._path, "rb") as f: + with self._open() as f: while chunk := f.read(self._COPY_CHUNK_SIZE): digest.update(chunk) self._sha256 = digest.digest() return self._sha256 def to_bytes(self) -> bytes: - with open(self._path, "rb") as f: + with self._open() as f: return f.read() def write_to_file(self, outfile: io.BufferedIOBase) -> None: - with open(self._path, "rb") as f: + with self._open() as f: shutil.copyfileobj(f, outfile, length=self._COPY_CHUNK_SIZE) + def _move_to(self, outfile: io.BufferedIOBase) -> Optional[int]: + """Move the data into ``outfile``, a regular file that its caller created for + this write (see the class docstring). Returns the offset in ``outfile`` it was + moved to, or None, having written nothing, if it has to be copied instead.""" + source = self._open_to_move(outfile) + if source is None: + return None + try: + return self._move_chunks(outfile, source) + finally: + os.close(source) + + def _copy_reason( + self, output_device: int, source: Optional[os.stat_result] = None + ) -> Optional[str]: + """Why a move into a file on ``output_device`` would copy the data instead, or + None if it can move. Judges ``source``, the stat of the opened file, if given, and + else the file at the path.""" + if self._finalizer is None or not self._finalizer.alive: + return "not owned" + if sys.platform == "win32": + return "on Windows" + if source is None: + try: + source = os.lstat(self._path) + except OSError: + return "unreadable" + if not stat.S_ISREG(source.st_mode) or not os.access(self._path, os.W_OK): + return "a link or read-only" + # Truncating changes every name of the file, and on another file system it frees + # nothing for the output. + if source.st_nlink != 1: + return "with another hard link" + if source.st_dev != output_device: + return "on another file system" + if source.st_size != self._size: + return "changed in size" + return None + + def _open_to_move(self, outfile: io.BufferedIOBase) -> Optional[int]: + output_stat = os.fstat(outfile.fileno()) + if self._copy_reason(output_stat.st_dev) is not None: + return None + try: + # O_NOFOLLOW refuses a link, and the checks below look at the file opened. + source = os.open(self._path, os.O_RDWR | os.O_NOFOLLOW) + except OSError: + return None + try: + source_stat = os.fstat(source) + if os.path.samestat(source_stat, output_stat): + raise ValueError(f"Can't move {self._path} into itself") + except BaseException: + os.close(source) + raise + if self._copy_reason(output_stat.st_dev, source_stat) is not None: + os.close(source) + return None + return source + + def _move_chunks(self, outfile: io.BufferedIOBase, source: int) -> int: + start = outfile.tell() + end = self._size + detached = False + try: + while end > 0: + begin = max(end - self._COPY_CHUNK_SIZE, 0) + chunk = os.pread(source, end - begin, begin) + if len(chunk) != end - begin: + raise OSError( + f"read {len(chunk)} of {end - begin} bytes at offset {begin} " + f"of {self._path}" + ) + outfile.seek(start + begin) + written = outfile.write(chunk) + if written != len(chunk): + raise OSError( + f"wrote {written} of {len(chunk)} bytes at offset " + f"{start + begin}" + ) + # The chunk has to be on the output's disk before the source gives it up. + outfile.flush() + os.fsync(outfile.fileno()) + if not detached and self._finalizer is not None: + # From the first truncation, part of the data is found only through + # the error note, so no interrupt may let the finalizer delete it. + detached = self._finalizer.detach() is not None + end = begin + os.ftruncate(source, begin) + if detached: + self._finalizer = weakref.finalize(self, self._remove, self._path) + outfile.seek(start + self._size) + except BaseException as e: + try: + # The truncation itself may have failed and left more bytes in the file. + end = os.fstat(source).st_size + except OSError: + pass + if detached and end == self._size: + # Nothing left the file, so this object still owns it. + self._finalizer = weakref.finalize(self, self._remove, self._path) + elif end < self._size: + e.add_note( + f"Bytes [0, {end}) are still in {self._path}; bytes " + f"[{end}, {self._size}) are at offset {start + end} of the output." + ) + raise + return start + def close(self) -> None: if self._finalizer is not None: self._finalizer() @@ -131,10 +288,35 @@ def append(self, data: Union[CordBuffer, "Cord"]) -> None: f"Can only append bytes, FileBackedData, or Cords, received {type(data)}" ) + def file_backed_buffers(self) -> List[FileBackedData]: + """The file-backed buffers of the Cord, in order, with repeats.""" + return [item for item in self._buffers if isinstance(item, FileBackedData)] + def write_to_file(self, outfile: io.BufferedIOBase) -> None: """Write the Cord to a file.""" - for item in self._buffers: - if isinstance(item, bytes): - outfile.write(item) - else: - item.write_to_file(outfile) + self._write_to_file(outfile) + + def _write_to_file( + self, outfile: io.BufferedIOBase, move: Container[FileBackedData] = () + ) -> None: + """Write the Cord to ``outfile``, moving the buffers in ``move`` into it when they + can be (see ``FileBackedData``). Each of them must appear once. If the write + fails, a note on the error lists the buffers already moved, whose only copy is now + in ``outfile``, and where they are.""" + moved: List[str] = [] + try: + for item in self._buffers: + if isinstance(item, bytes): + outfile.write(item) + continue + offset = item._move_to(outfile) if item in move else None + if offset is None: + item.write_to_file(outfile) + else: + moved.append(f"{item.path} ({len(item)} bytes) at offset {offset}") + except BaseException as e: + if moved: + e.add_note( + "Moved into the output before the failure: " + "; ".join(moved) + ) + raise diff --git a/exir/_serialize/test/test_cord.py b/exir/_serialize/test/test_cord.py index 85770015da8..80204042ecf 100644 --- a/exir/_serialize/test/test_cord.py +++ b/exir/_serialize/test/test_cord.py @@ -5,13 +5,21 @@ # LICENSE file in the root directory of this source tree. +import copy import io import os +import random +import sys import tempfile import unittest +from unittest import mock from executorch.exir._serialize._cord import Cord, FileBackedData +_needs_move = unittest.skipIf( + sys.platform == "win32", "Windows copies instead of moving" +) + class TestCord(unittest.TestCase): def test_cord_init(self) -> None: @@ -82,3 +90,415 @@ def test_file_backed_data(self) -> None: self.assertEqual(b"PrefixFileBacked", bytes(cord)) self.assertEqual([], os.listdir(directory)) + + # 10 full chunks and part of one, each different, so a chunk read from the wrong + # offset shows in the output. + _CHUNK = 8192 + _PAYLOAD = random.Random(0).randbytes(10 * _CHUNK + 1000) + + def _owned(self, directory: str) -> FileBackedData: + path = os.path.join(directory, "source.bin") + with open(path, "wb") as f: + f.write(self._PAYLOAD) + return FileBackedData.move_from(path) + + def _move(self, data: FileBackedData, out_path: str) -> None: + with open(out_path, "wb") as outfile: + Cord(data)._write_to_file(outfile, move={data}) + + def _assert_copied(self, data: FileBackedData, out_path: str) -> None: + self.assertEqual(self._PAYLOAD, data.to_bytes()) + with open(out_path, "rb") as result: + self.assertEqual(self._PAYLOAD, result.read()) + + @_needs_move + def test_a_move_keeps_the_data_on_disk_once(self) -> None: + # Random bytes, so a compressing file system can't hide a second copy. APFS fills + # a hole under about 16 MiB with zeros, so there a move can allocate up to 16 MiB + # more; with 64 MiB of data and 4 MiB chunks that stays under the bound. + chunk = 4 * 1024 * 1024 + payload = random.Random(0).randbytes(16 * chunk) + with tempfile.TemporaryDirectory() as directory: + source = os.path.join(directory, "source.bin") + with open(source, "wb") as f: + f.write(payload) + data = FileBackedData.move_from(source) + source = data.path + out_path = os.path.join(directory, "out.bin") + on_disk = [] + + class Recording(io.FileIO): + def write(self, b) -> int: + offset = self.tell() - len(b"Prefix") + size = super().write(b) + if 0 <= offset < len(payload): + # The source must still hold the range that reaches the output. + assert os.path.getsize(source) >= offset + len(b), offset + blocks = os.stat(source).st_blocks + on_disk.append( + (blocks + os.fstat(self.fileno()).st_blocks) * 512 + ) + return size + + cord = Cord(b"Prefix") + cord.append(data) + cord.append(b"Suffix") + with mock.patch.object(FileBackedData, "_COPY_CHUNK_SIZE", chunk): + # A buffer bigger than a chunk, so a chunk not flushed before the + # truncation reaches the file too late and fails the check above. + with io.BufferedWriter( + Recording(out_path, "wb"), buffer_size=2 * chunk + ) as outfile: + cord._write_to_file(outfile, move={data}) + + with open(out_path, "rb") as result: + self.assertEqual(b"Prefix" + payload + b"Suffix", result.read()) + # Allocated blocks, so the unwritten part of the output must stay a hole. A move + # needs the data plus a chunk, and a copy twice the data. + self.assertLess(max(on_disk), len(payload) + len(payload) // 2) + self.assertEqual(0, os.path.getsize(source)) + with self.assertRaises(RuntimeError): + data.to_bytes() + + @_needs_move + def test_only_a_move_empties_the_file(self) -> None: + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + out_path = os.path.join(directory, "out.bin") + for move in ((), {data}): + with self.subTest(move=bool(move)): + with open(out_path, "wb") as outfile: + Cord(data)._write_to_file(outfile, move=move) + with open(out_path, "rb") as result: + self.assertEqual(self._PAYLOAD, result.read()) + self.assertEqual(0, os.path.getsize(data.path)) + + @_needs_move + def test_reading_a_moved_file_raises(self) -> None: + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + second = FileBackedData(data.path) + self._move(data, os.path.join(directory, "out.bin")) + reads = { + "to_bytes": second.to_bytes, + "prefix": lambda: second.prefix(1), + "sha256": second.sha256, + "check": second.check, + "write_to_file": lambda: second.write_to_file(io.BytesIO()), + } + for name, read in reads.items(): + with self.subTest(read=name), self.assertRaises(RuntimeError): + read() + + def test_reading_a_file_that_grew_raises(self) -> None: + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + with open(data.path, "ab") as f: + f.write(b"more") + with self.assertRaisesRegex(RuntimeError, "changed after it was wrapped"): + data.to_bytes() + + def test_a_deep_copy_shares_the_data_and_its_owner(self) -> None: + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + self.assertIs(data, copy.deepcopy(data)) + self.assertIs(data, copy.deepcopy(Cord(data)).file_backed_buffers()[0]) + + @_needs_move + def test_a_move_into_the_file_itself_raises(self) -> None: + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + link = os.path.join(directory, "link.bin") + os.symlink(data.path, link) + for name in (data.path, link): + with self.subTest(output=os.path.basename(name)): + with open(name, "r+b") as outfile: + with self.assertRaises(ValueError): + Cord(data)._write_to_file(outfile, move={data}) + self.assertEqual(self._PAYLOAD, data.to_bytes()) + + def test_a_move_copies_data_that_does_not_own_its_file(self) -> None: + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "source.bin") + with open(path, "wb") as f: + f.write(self._PAYLOAD) + data = FileBackedData(path) + out_path = os.path.join(directory, "out.bin") + self._move(data, out_path) + self._assert_copied(data, out_path) + + def test_a_move_copies_a_file_with_another_name(self) -> None: + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + os.link(data.path, os.path.join(directory, "other_name.bin")) + out_path = os.path.join(directory, "out.bin") + self._move(data, out_path) + self._assert_copied(data, out_path) + + @_needs_move + def test_a_move_judges_the_file_it_opened(self) -> None: + lstat = os.lstat + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + os.link(data.path, os.path.join(directory, "other_name.bin")) + + def lstat_before_the_link(path, *args, **kwargs) -> os.stat_result: + st = lstat(path, *args, **kwargs) + if path != data.path: + return st + return os.stat_result((*st[:3], 1, *st[4:10])) + + out_path = os.path.join(directory, "out.bin") + with mock.patch.object(os, "lstat", lstat_before_the_link): + self._move(data, out_path) + self._assert_copied(data, out_path) + + def test_a_move_copies_on_windows(self) -> None: + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + out_path = os.path.join(directory, "out.bin") + with mock.patch.object(sys, "platform", "win32"): + self._move(data, out_path) + self._assert_copied(data, out_path) + + @_needs_move + def test_a_move_copies_a_symlink_and_keeps_its_target(self) -> None: + with tempfile.TemporaryDirectory() as directory: + target = os.path.join(directory, "target.bin") + with open(target, "wb") as f: + f.write(self._PAYLOAD) + link = os.path.join(directory, "link.bin") + os.symlink(target, link) + data = FileBackedData.move_from(link) + out_path = os.path.join(directory, "out.bin") + self._move(data, out_path) + self._assert_copied(data, out_path) + with open(target, "rb") as f: + self.assertEqual(self._PAYLOAD, f.read()) + + @_needs_move + def test_a_move_copies_into_another_file_system(self) -> None: + fstat = os.fstat + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + out_path = os.path.join(directory, "out.bin") + with open(out_path, "wb") as outfile: + + def other_device_fstat(fd: int) -> os.stat_result: + st = fstat(fd) + if fd != outfile.fileno(): + return st + return os.stat_result((*st[:2], st.st_dev + 1, *st[3:10])) + + with mock.patch.object(os, "fstat", other_device_fstat): + Cord(data)._write_to_file(outfile, move={data}) + self._assert_copied(data, out_path) + + @_needs_move + def test_a_move_copies_data_it_cannot_write_to(self) -> None: + if os.geteuid() == 0: + self.skipTest("root can write to any file") + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + os.chmod(data.path, 0o444) + out_path = os.path.join(directory, "out.bin") + self._move(data, out_path) + self._assert_copied(data, out_path) + + @_needs_move + def test_a_move_copies_buffers_not_listed_to_move(self) -> None: + with tempfile.TemporaryDirectory() as directory: + moved, copied = self._owned(directory), self._owned(directory) + cord = Cord(moved) + cord.append(copied) + out_path = os.path.join(directory, "out.bin") + with open(out_path, "wb") as outfile: + cord._write_to_file(outfile, move={moved}) + with open(out_path, "rb") as result: + self.assertEqual(self._PAYLOAD * 2, result.read()) + self.assertEqual(0, os.path.getsize(moved.path)) + self.assertEqual(self._PAYLOAD, copied.to_bytes()) + + @_needs_move + def test_a_move_that_fails_first_keeps_the_data(self) -> None: + class Short(io.FileIO): + def write(self, b) -> int: + return super().write(bytes(b)[:-1]) + + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + out_path = os.path.join(directory, "out.bin") + with Short(out_path, "wb") as outfile: + with self.assertRaises(OSError): + Cord(data)._write_to_file(outfile, move={data}) + self.assertEqual(self._PAYLOAD, data.to_bytes()) + self._move(data, out_path) + with open(out_path, "rb") as result: + self.assertEqual(self._PAYLOAD, result.read()) + + @_needs_move + def test_a_move_that_fails_later_says_where_the_bytes_are(self) -> None: + payload, prefix = self._PAYLOAD, b"Prefix" + + class FailsThird(io.FileIO): + writes = 0 + + def write(self, b) -> int: + # The prefix, then the first chunk moved, then the failure. + FailsThird.writes += 1 + if FailsThird.writes == 3: + raise OSError("disk full") + return super().write(b) + + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + out_path = os.path.join(directory, "out.bin") + cord = Cord(prefix) + cord.append(data) + with mock.patch.object(FileBackedData, "_COPY_CHUNK_SIZE", self._CHUNK): + with FailsThird(out_path, "wb") as outfile: + with self.assertRaises(OSError) as raised: + cord._write_to_file(outfile, move={data}) + # The first chunk moved is the last one, a full chunk. + left = len(payload) - self._CHUNK + path = data.path + del data, cord # The bytes the note points to outlive the object. + note = raised.exception.__notes__[0] + self.assertIn(f"Bytes [0, {left}) are still in {path}", note) + self.assertIn(f"at offset {len(prefix) + left} of the output", note) + with open(path, "rb") as source: + self.assertEqual(payload[:left], source.read()) + with open(out_path, "rb") as result: + result.seek(len(prefix) + left) + self.assertEqual(payload[left:], result.read()) + + @_needs_move + def test_an_interrupt_right_after_a_truncation_gets_a_true_note(self) -> None: + ftruncate = os.ftruncate + + def interrupted_ftruncate(fd: int, length: int) -> None: + ftruncate(fd, length) + raise KeyboardInterrupt + + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + with mock.patch.object(FileBackedData, "_COPY_CHUNK_SIZE", self._CHUNK): + with mock.patch.object(os, "ftruncate", interrupted_ftruncate): + with self.assertRaises(KeyboardInterrupt) as raised: + self._move(data, os.path.join(directory, "out.bin")) + left = len(self._PAYLOAD) - self._CHUNK + path = data.path + del data + self.assertEqual(left, os.path.getsize(path)) + self.assertIn(f"Bytes [0, {left})", raised.exception.__notes__[0]) + + @_needs_move + def test_a_failed_truncation_keeps_the_data_and_its_owner(self) -> None: + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + with mock.patch.object(os, "ftruncate", side_effect=OSError("read-only")): + with self.assertRaises(OSError) as raised: + self._move(data, os.path.join(directory, "out.bin")) + self.assertFalse(hasattr(raised.exception, "__notes__")) + self.assertEqual(self._PAYLOAD, data.to_bytes()) + path = data.path + data.close() + self.assertFalse(os.path.exists(path)) + + @_needs_move + def test_a_move_keeps_the_finalizer_off_while_the_file_shrinks(self) -> None: + ftruncate = os.ftruncate + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + alive = [] + + def recording_ftruncate(fd: int, length: int) -> None: + finalizer = data._finalizer + alive.append(finalizer is not None and finalizer.alive) + ftruncate(fd, length) + + with mock.patch.object(FileBackedData, "_COPY_CHUNK_SIZE", self._CHUNK): + with mock.patch.object(os, "ftruncate", recording_ftruncate): + self._move(data, os.path.join(directory, "out.bin")) + self.assertGreater(len(alive), 1) + self.assertFalse(any(alive)) + # Once the move is done, the object owns the emptied file again. + path = data.path + data.close() + self.assertFalse(os.path.exists(path)) + + @_needs_move + def test_a_failed_final_seek_still_says_where_the_bytes_are(self) -> None: + class FailsLastSeek(io.FileIO): + def seek(self, offset: int, whence: int = 0) -> int: + if offset == len(TestCord._PAYLOAD): + raise OSError("seek failed") + return super().seek(offset, whence) + + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + with FailsLastSeek(os.path.join(directory, "out.bin"), "wb") as outfile: + with self.assertRaises(OSError) as raised: + Cord(data)._write_to_file(outfile, move={data}) + self.assertIn( + f"bytes [0, {len(self._PAYLOAD)}) are at offset 0", + raised.exception.__notes__[0], + ) + + @_needs_move + def test_a_write_error_reported_by_the_sync_keeps_the_data(self) -> None: + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + with mock.patch.object(os, "fsync", side_effect=OSError("quota")): + with self.assertRaises(OSError): + self._move(data, os.path.join(directory, "out.bin")) + self.assertEqual(self._PAYLOAD, data.to_bytes()) + + @_needs_move + def test_a_short_read_stops_the_move_and_keeps_the_data(self) -> None: + pread = os.pread + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + with mock.patch.object( + os, "pread", lambda fd, n, offset: pread(fd, n - 1, offset) + ): + with self.assertRaises(OSError): + self._move(data, os.path.join(directory, "out.bin")) + self.assertEqual(self._PAYLOAD, data.to_bytes()) + + @_needs_move + def test_a_failed_cord_write_says_where_the_moved_buffers_are(self) -> None: + payload, prefix = self._PAYLOAD, b"Prefix" + second_starts = len(prefix) + len(payload) + + class FailsInSecond(io.FileIO): + def write(self, b) -> int: + if self.tell() >= second_starts: + raise OSError("disk full") + return super().write(b) + + with tempfile.TemporaryDirectory() as directory: + first, second = self._owned(directory), self._owned(directory) + cord = Cord(prefix) + cord.append(first) + cord.append(second) + out_path = os.path.join(directory, "out.bin") + with FailsInSecond(out_path, "wb") as outfile: + with self.assertRaises(OSError) as raised: + cord._write_to_file(outfile, move={first, second}) + self.assertIn( + f"{first.path} ({len(payload)} bytes) at offset {len(prefix)}", + raised.exception.__notes__[-1], + ) + with open(out_path, "rb") as result: + self.assertEqual(prefix + payload, result.read(second_starts)) + self.assertEqual(self._PAYLOAD, second.to_bytes()) + + @_needs_move + def test_moving_a_file_that_another_move_emptied_raises(self) -> None: + with tempfile.TemporaryDirectory() as directory: + data = self._owned(directory) + second = FileBackedData(data.path, cleanup=True) + self._move(data, os.path.join(directory, "out.bin")) + with self.assertRaises(RuntimeError): + self._move(second, os.path.join(directory, "out2.bin")) diff --git a/exir/program/_program.py b/exir/program/_program.py index 927976a82d7..9e0b80c38e4 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -6,16 +6,34 @@ # LICENSE file in the root directory of this source tree. # pyre-strict +import contextlib import copy import io import logging import os -from collections import defaultdict -from typing import Any, Dict, List, Optional, Sequence, Set, TextIO, Type, Union +import stat +import uuid +from collections import Counter, defaultdict +from typing import ( + Any, + Collection, + Container, + Dict, + Iterable, + Iterator, + List, + Optional, + Sequence, + Set, + TextIO, + Tuple, + Type, + Union, +) import torch import torch._export -from executorch.exir._serialize._cord import Cord +from executorch.exir._serialize._cord import Cord, FileBackedData from executorch.exir._serialize._named_data_store import ( NamedDataStore, NamedDataStoreOutput, @@ -1874,6 +1892,117 @@ def to_executorch( # noqa (FLAKE8) C901 return et_pm +_UNFIT_OUTPUT = "for an output that can't take moved data" + + +def _copy_reasons( + path: str, buffers: Collection[FileBackedData], shared: Container[FileBackedData] +) -> Tuple[str, Optional[os.stat_result], Dict[FileBackedData, Optional[str]]]: + """The real path behind `path`, its stat, and why each of `buffers` would be copied + rather than moved into it, or None for those that can move. `shared` ones are + copied.""" + # A link at `path` is followed, as a plain save does. + real = os.path.realpath(path) + try: + existing: Optional[os.stat_result] = os.stat(real) + except FileNotFoundError: + existing = None + # Data moves only into a new or writable regular file whose directory can hold the + # temporary file; anything else gets the plain copy. + if existing is not None and not ( + stat.S_ISREG(existing.st_mode) and os.access(real, os.W_OK) + ): + return real, existing, dict.fromkeys(buffers, _UNFIT_OUTPUT) + device = os.stat(os.path.dirname(real)).st_dev + return ( + real, + existing, + { + data: "used more than once" if data in shared else data._copy_reason(device) + for data in buffers + }, + ) + + +def _warn_about_copies(path: str, total: int, reasons: Iterable[str]) -> None: + copied = Counter(reasons) + if copied: + logging.warning( + f"Copying {sum(copied.values())} of {total} file-backed buffers into {path} " + "instead of moving them: " + + ", ".join(f"{count} {reason}" for reason, count in copied.items()) + ) + + +def _end_temporary_output( + error: BaseException, + temp: str, + real: str, + move: Set[FileBackedData], + complete: bool, +) -> None: + if all(os.path.getsize(data.path) == len(data) for data in move): + # Nothing was moved, so the temporary file holds no data of its own. + os.remove(temp) + elif complete: + error.add_note(f"{temp} is complete; renaming it onto {real} ends the write.") + else: + error.add_note(f"The output written so far, with the data moved, is in {temp}.") + + +@contextlib.contextmanager +def _open_output( + path: str, buffers: Collection[FileBackedData], shared: Container[FileBackedData] +) -> Iterator[Tuple[io.BufferedWriter, Set[FileBackedData]]]: + """Open `path` for writing and pick which of `buffers` move into it; `shared` ones are + copied. Only when some move, open a temporary file next to `path` instead and rename + it onto `path` when the write succeeds.""" + # A consuming write fills its output out of order, so one stopped partway would leave + # a full-length file with holes at `path`, which loads with wrong data. + real = temp = "" + existing: Optional[os.stat_result] = None + fd = None + move: Set[FileBackedData] = set() + if buffers: + reasons: Dict[FileBackedData, Optional[str]] = dict.fromkeys( + buffers, _UNFIT_OUTPUT + ) + try: + real, existing, reasons = _copy_reasons(path, buffers, shared) + if None in reasons.values(): + temp = f"{real}.{uuid.uuid4().hex[:8]}.partial" + mode = 0o666 if existing is None else stat.S_IMODE(existing.st_mode) + fd = os.open(temp, os.O_WRONLY | os.O_CREAT | os.O_EXCL, mode) + except OSError: + pass + if fd is not None: + move = {data for data, reason in reasons.items() if reason is None} + _warn_about_copies( + path, + len(buffers), + [reasons[data] or _UNFIT_OUTPUT for data in buffers if data not in move], + ) + if fd is None: + with open(path, "wb") as f: + yield f, move + return + complete = False + try: + with open(fd, "wb") as output: + if existing is not None: + # The umask applied above; a plain save keeps the old file's mode. + os.fchmod(fd, stat.S_IMODE(existing.st_mode)) + yield output, move + output.flush() + # So a crash right after the rename can't leave `path` with missing data. + os.fsync(output.fileno()) + complete = True + os.replace(temp, real) + except BaseException as e: + _end_temporary_output(e, temp, real, move, complete) + raise + + class ExecutorchProgramManager: """ Package of one or more `ExportedPrograms` in Execution dialect. Designed to simplify @@ -2045,25 +2174,101 @@ def write_to_file(self, open_file: io.BufferedIOBase) -> None: """ self._pte_data.write_to_file(open_file) - def write_tensor_data_to_file(self, outdir) -> None: + def write_tensor_data_to_file(self, outdir, *, consume: bool = False) -> None: """ - Writes the serialized ExecuTorch data files to the directory at `outdir`. + Writes the serialized ExecuTorch data files to the directory at `outdir`. `consume` + works as in `save`. If a consuming write fails after some data files are done, a + note on the error lists them, since they may hold the only copy of data moved + into them. Data moved before a failure can't be written again, so rerun the export. """ assert self._tensor_data is not None - for filename, cord in self._tensor_data.items(): - if not filename.endswith(".ptd"): - filename += ".ptd" - with open(os.path.join(outdir, f"{filename}"), "wb") as f: - logging.info(f"Writing data file to {filename}") - cord.write_to_file(f) + files = { + os.path.join(outdir, name if name.endswith(".ptd") else f"{name}.ptd"): cord + for name, cord in self._tensor_data.items() + } + self._check_file_backed_data(files.values(), files) + shared = self._shared_file_backed_data() + done: List[str] = [] + try: + for path, cord in files.items(): + logging.info(f"Writing data file to {os.path.basename(path)}") + buffers = set(cord.file_backed_buffers()) if consume else set() + with _open_output(path, buffers, shared) as (f, move): + cord._write_to_file(f, move) + if move: + done.append(path) + except BaseException as e: + if done: + e.add_note(f"These data files were written before the failure: {done}") + raise - def save(self, path: str) -> None: + def _check_file_backed_data( + self, cords: Iterable[Cord], paths: Iterable[str] + ) -> None: + # Called before any output is opened, so a write can't empty a file that holds the + # only copy of data that an earlier consuming write moved, or replace the data's + # own file. + sources = [] + for cord in cords: + for data in cord.file_backed_buffers(): + data.check() + sources.append(os.stat(data.path)) + for path in paths: + if os.path.exists(path): + output = os.stat(path) + if any(os.path.samestat(output, source) for source in sources): + raise ValueError(f"{path} is the file that holds the data to write") + + def _shared_file_backed_data(self) -> Set[FileBackedData]: + # Deduplication can put one buffer in the .pte and in .ptd files, or twice in one + # file, and a write can't know whether it is the last one to need it, so such a + # buffer is copied. + seen: Set[FileBackedData] = set() + shared: Set[FileBackedData] = set() + for cord in [self._pte_data, *(self._tensor_data or {}).values()]: + for data in cord.file_backed_buffers(): + if data in seen: + shared.add(data) + else: + seen.add(data) + return shared + + def save(self, path: str, *, consume: bool = False) -> None: """ Saves the serialized ExecuTorch binary to the file at `path`. + + With `consume`, data the program keeps in files it owns, such as large delegate + payloads, is moved into `path` instead of copied (see `FileBackedData`). That + trades time for disk space: on file systems with sparse files, a save to a new + `path` peaks at about one copy of the data instead of two, but each chunk is synced + before its source shrinks, so the save is several times slower. Moved data can't + be read or written again, by this program, by the lowered program it came from, or + by a deep copy of either, which shares its data (`buffer` still returns its bytes + if it was read before), so save any ETRecord first. Data that can't be moved, or + that the `.pte` and a `.ptd` share, is copied, and one warning says how many + buffers and why. Where data moves, the save writes the program itself and does not + call `write_to_file`. + + Data moves into a new file at `path` or an existing writable regular file. The + save then writes a temporary file next to `path` and renames it onto `path` only + when the write succeeds. If it fails, `path` is unchanged. The temporary file is + deleted if no data was moved yet; otherwise a note on the error names it, since it + holds the only copy of the moved data, which can't be written again, so rerun the + export. Until the rename, an old file at `path` still takes its space on disk. The + rename gives `path` a new file, owned by the caller's user and group; a link at + `path` is followed and the file it names keeps its mode. With nothing to move, or + anywhere else, such as a pipe, where no temporary file can be created next to + `path`, or on Windows, the save copies, as without `consume`. """ if path[-4:] != ".pte": logging.error(f"Path {path} does not end with .pte") raise ValueError(f"Path {path} does not end with .pte") - with open(path, "wb") as file: - self.write_to_file(file) + self._check_file_backed_data([self._pte_data], [path]) + buffers = set(self._pte_data.file_backed_buffers()) if consume else set() + shared = self._shared_file_backed_data() + with _open_output(path, buffers, shared) as (file, move): + if move: + self._pte_data._write_to_file(file, move) + else: + self.write_to_file(file) logging.info(f"Saved exported program to {path}") diff --git a/exir/program/test/BUCK b/exir/program/test/BUCK index eddfdf89e96..b9e0f627786 100644 --- a/exir/program/test/BUCK +++ b/exir/program/test/BUCK @@ -18,6 +18,7 @@ fbcode_target(_kind = runtime.python_test, "//executorch/exir:lib", "//executorch/exir:print_program", "//executorch/exir:schema", + "//executorch/exir/backend/test:backend_with_named_data_map", "//executorch/exir/backend/test:op_partitioner_demo", "//executorch/exir/emit:lib", "//executorch/exir/passes:const_prop_pass", diff --git a/exir/program/test/test_program.py b/exir/program/test/test_program.py index 916d8efd7a8..519b582b96f 100644 --- a/exir/program/test/test_program.py +++ b/exir/program/test/test_program.py @@ -7,15 +7,26 @@ # pyre-unsafe import copy +import errno import os +import stat +import sys import tempfile +import threading import unittest from typing import Any, Dict from unittest import mock import torch from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig -from executorch.exir._serialize._cord import Cord +from executorch.exir._serialize._cord import Cord, FileBackedData +from executorch.exir._serialize._named_data_store import NamedDataStore +from executorch.exir._serialize._program import deserialize_pte_binary +from executorch.exir.backend.backend_details import PreprocessResult +from executorch.exir.backend.test.backend_with_named_data_map import ( + BackendWithNamedDataMap, + BackendWithNDMPartitioner, +) from executorch.exir.backend.test.op_partitioner_demo import ( AddMulPartitionerDemo, NonDecompTestPartitioner, @@ -36,6 +47,7 @@ from executorch.exir.verification.verifier import EXIREdgeDialectVerifier from executorch.extension.export_util.utils import save_pte_program +from executorch.extension.flat_tensor.serialize.serialize import FlatTensorSerializer from executorch.extension.pybindings.portable_lib import ( _load_for_executorch_from_buffer, ) @@ -48,6 +60,14 @@ from torch.nn import functional as F +_needs_move = unittest.skipIf( + sys.platform == "win32", "Windows copies instead of moving" +) +_needs_posix = unittest.skipIf( + sys.platform == "win32", "needs symlinks and POSIX file modes" +) + + class TestLinear(torch.nn.Module): def __init__(self): super().__init__() @@ -1235,6 +1255,525 @@ def test_save_pte_program_writes_the_data_files_and_raises_on_failure(self): with self.assertRaises(OSError): save_pte_program(et, path, directory) + def _program_with_owned_data(self, directory, name="blob"): + model = TestLinear() + et = to_edge( + torch.export.export(model, model._get_random_inputs(), strict=True) + ).to_executorch() + return et, self._owned(directory, name) + + def _owned(self, directory, name): + path = os.path.join(directory, f"{name}.bin") + with open(path, "wb") as f: + f.write(name.encode() * 1024) + return FileBackedData.move_from(path) + + def test_a_save_without_consume_leaves_the_program_whole(self): + with tempfile.TemporaryDirectory() as directory: + et, owned = self._program_with_owned_data(directory) + et._pte_data = Cord(owned) + for name in ("first.pte", "second.pte"): + et.save(os.path.join(directory, name)) + with open(os.path.join(directory, name), "rb") as f: + self.assertEqual(b"blob" * 1024, f.read()) + self.assertEqual(b"blob" * 1024, et.buffer) + + @_needs_move + def test_consume_moves_an_owned_buffer_into_the_pte(self): + with tempfile.TemporaryDirectory() as directory: + et, owned = self._program_with_owned_data(directory) + et._pte_data = Cord(b"header") + et._pte_data.append(owned) + with self.assertNoLogs(level="WARNING"): + et.save(os.path.join(directory, "model.pte"), consume=True) + with open(os.path.join(directory, "model.pte"), "rb") as f: + self.assertEqual(b"header" + b"blob" * 1024, f.read()) + self.assertEqual(0, os.path.getsize(owned.path)) + self.assertEqual( + ["model.pte"], [n for n in os.listdir(directory) if "model" in n] + ) + + @_needs_move + def test_consume_moves_backend_data_from_a_deep_copied_edge_program(self): + with tempfile.TemporaryDirectory() as directory: + payload = bytes(range(256)) * 64 + owned = self._save_backend_data_with_consume( + directory, payload, None, deep_copy=True + ) + self.assertEqual([0], [os.path.getsize(d.path) for d in owned]) + self._assert_backend_payload(directory, payload, None) + + @_needs_move + def test_a_consuming_save_stopped_partway_leaves_the_old_file(self): + ftruncate = os.ftruncate + + def interrupted_ftruncate(fd, length): + ftruncate(fd, length) + raise KeyboardInterrupt + + for old in (None, b"old program"): + with self.subTest(old=old), tempfile.TemporaryDirectory() as directory: + et, owned = self._program_with_owned_data(directory) + et._pte_data = Cord(owned) + path = os.path.join(directory, "model.pte") + if old is not None: + with open(path, "wb") as f: + f.write(old) + with mock.patch.object(FileBackedData, "_COPY_CHUNK_SIZE", 1024): + with mock.patch.object(os, "ftruncate", interrupted_ftruncate): + with self.assertRaises(KeyboardInterrupt) as raised: + et.save(path, consume=True) + if old is None: + self.assertFalse(os.path.exists(path)) + else: + with open(path, "rb") as f: + self.assertEqual(old, f.read()) + (partial,) = [ + n for n in os.listdir(directory) if n.endswith(".partial") + ] + partial = os.path.join(os.path.realpath(directory), partial) + self.assertIn(partial, raised.exception.__notes__[-1]) + + @_needs_move + def test_a_consuming_ptd_write_stopped_before_moving_leaves_no_file(self): + with tempfile.TemporaryDirectory() as directory: + et, owned = self._program_with_owned_data(directory) + et._tensor_data = {"data": Cord(owned)} + with mock.patch.object(os, "fsync", side_effect=KeyboardInterrupt): + with self.assertRaises(KeyboardInterrupt): + et.write_tensor_data_to_file(directory, consume=True) + self.assertEqual([os.path.basename(owned.path)], os.listdir(directory)) + self.assertEqual(b"blob" * 1024, owned.to_bytes()) + + @_needs_move + def test_a_failed_consuming_ptd_write_lists_the_files_already_written(self): + fsync = os.fsync + calls = [] + + def fails_in_the_second_file(fd): + # The first file syncs its one chunk and then itself. + calls.append(fd) + if len(calls) == 3: + raise OSError("disk full") + fsync(fd) + + with tempfile.TemporaryDirectory() as directory: + et, first = self._program_with_owned_data(directory, "first") + second = self._owned(directory, "second") + et._tensor_data = {"a": Cord(first), "b": Cord(second)} + with mock.patch.object(os, "fsync", fails_in_the_second_file): + with self.assertRaises(OSError) as raised: + et.write_tensor_data_to_file(directory, consume=True) + done = os.path.join(directory, "a.ptd") + self.assertIn(repr(done), raised.exception.__notes__[-1]) + with open(done, "rb") as f: + self.assertEqual(b"first" * 1024, f.read()) + self.assertEqual(0, os.path.getsize(first.path)) + + @_needs_move + def test_consume_moves_into_the_pte_and_then_the_ptd(self): + with tempfile.TemporaryDirectory() as directory: + et, pte_data = self._program_with_owned_data(directory) + ptd_data = self._owned(directory, "data") + et._pte_data = Cord(pte_data) + et._tensor_data = {"data": Cord(ptd_data)} + et.save(os.path.join(directory, "model.pte"), consume=True) + et.write_tensor_data_to_file(directory, consume=True) + for name, data in (("model.pte", b"blob"), ("data.ptd", b"data")): + with open(os.path.join(directory, name), "rb") as f: + self.assertEqual(data * 1024, f.read()) + self.assertEqual(0, os.path.getsize(pte_data.path)) + self.assertEqual(0, os.path.getsize(ptd_data.path)) + + @_needs_move + def test_a_save_after_a_consuming_save_keeps_the_first_output(self): + with tempfile.TemporaryDirectory() as directory: + et, owned = self._program_with_owned_data(directory) + et._pte_data = Cord(owned) + path = os.path.join(directory, "model.pte") + et.save(path, consume=True) + for consume in (False, True): + with self.assertRaises(RuntimeError): + et.save(path, consume=consume) + with open(path, "rb") as f: + self.assertEqual(b"blob" * 1024, f.read()) + + @_needs_move + def test_a_ptd_write_after_a_consuming_one_keeps_the_first_output(self): + with tempfile.TemporaryDirectory() as directory: + et, owned = self._program_with_owned_data(directory) + et._tensor_data = {"data": Cord(owned)} + et.write_tensor_data_to_file(directory, consume=True) + with self.assertRaises(RuntimeError): + et.write_tensor_data_to_file(directory) + with open(os.path.join(directory, "data.ptd"), "rb") as f: + self.assertEqual(b"blob" * 1024, f.read()) + + @_needs_move + def test_save_pte_program_after_a_consuming_save_keeps_the_first_output(self): + with tempfile.TemporaryDirectory() as directory: + et, owned = self._program_with_owned_data(directory) + et._pte_data = Cord(owned) + path = os.path.join(directory, "model.pte") + et.save(path, consume=True) + with self.assertRaises(RuntimeError): + save_pte_program(et, path) + with open(path, "rb") as f: + self.assertEqual(b"blob" * 1024, f.read()) + + def test_writing_over_the_data_file_itself_raises(self): + with tempfile.TemporaryDirectory() as directory: + et, owned = self._program_with_owned_data(directory) + et._pte_data = Cord(owned) + et._tensor_data = {"data": Cord(owned)} + for name in ("model.pte", "data.ptd"): + os.link(owned.path, os.path.join(directory, name)) + for consume in (False, True): + with self.subTest(consume=consume): + with self.assertRaises(ValueError): + et.save(os.path.join(directory, "model.pte"), consume=consume) + with self.assertRaises(ValueError): + et.write_tensor_data_to_file(directory, consume=consume) + self.assertEqual(b"blob" * 1024, owned.to_bytes()) + + @_needs_posix + def test_a_consuming_save_writes_through_a_link_and_keeps_the_mode(self): + fchmod = os.fchmod + created = [] + + def recording_fchmod(fd, mode): + created.append(stat.S_IMODE(os.fstat(fd).st_mode)) + fchmod(fd, mode) + + for umask in (0o022, 0o077): + with self.subTest(umask=oct(umask)), tempfile.TemporaryDirectory() as d: + created.clear() + et, owned = self._program_with_owned_data(d) + et._pte_data = Cord(owned) + target = os.path.join(d, "target.pte") + with open(target, "wb") as f: + f.write(b"old program") + os.chmod(target, 0o640) + inode = os.stat(target).st_ino + link = os.path.join(d, "model.pte") + os.symlink(target, link) + old_umask = os.umask(umask) + try: + with mock.patch.object(os, "fchmod", recording_fchmod): + et.save(link, consume=True) + finally: + os.umask(old_umask) + self.assertTrue(os.path.islink(link)) + self.assertNotEqual(inode, os.stat(target).st_ino) + with open(target, "rb") as f: + self.assertEqual(b"blob" * 1024, f.read()) + self.assertEqual(0o640, stat.S_IMODE(os.stat(target).st_mode)) + # The temporary file never allows more than the file it replaces. + self.assertTrue(created) + self.assertFalse([m for m in created if m & ~0o640]) + self.assertEqual(0, os.path.getsize(owned.path)) + + def test_a_consuming_save_with_nothing_to_move_writes_the_output_in_place(self): + for case in ("not owned", "read-only"): + if case == "read-only" and (sys.platform == "win32" or os.geteuid() == 0): + continue + with self.subTest(case=case), tempfile.TemporaryDirectory() as directory: + et, owned = self._program_with_owned_data(directory) + if case == "not owned": + path = os.path.join(directory, "data.bin") + with open(path, "wb") as f: + f.write(b"blob" * 1024) + owned = FileBackedData(path) + else: + os.chmod(owned.path, 0o444) + et._pte_data = Cord(owned) + out = os.path.join(directory, "model.pte") + with open(out, "wb") as f: + f.write(b"old program") + inode = os.stat(out).st_ino + with mock.patch.object(os, "fsync", wraps=os.fsync) as fsync: + et.save(out, consume=True) + fsync.assert_not_called() + self.assertEqual(inode, os.stat(out).st_ino) + with open(out, "rb") as f: + self.assertEqual(b"blob" * 1024, f.read()) + self.assertEqual(b"blob" * 1024, owned.to_bytes()) + + @_needs_move + def test_a_consuming_save_warns_once_about_copied_buffers(self): + with tempfile.TemporaryDirectory() as directory: + et, owned = self._program_with_owned_data(directory) + et._pte_data = Cord(owned) + for name in ("first", "second"): + path = os.path.join(directory, f"{name}.bin") + with open(path, "wb") as f: + f.write(b"blob") + et._pte_data.append(FileBackedData(path)) + with self.assertLogs(level="WARNING") as logs: + et.save(os.path.join(directory, "model.pte"), consume=True) + (warning,) = logs.output + self.assertIn("Copying 2 of 3 file-backed buffers", warning) + self.assertIn("2 not owned", warning) + self.assertNotIn("first.bin", warning) + self.assertEqual(0, os.path.getsize(owned.path)) + + @_needs_move + def test_a_consuming_save_whose_rename_fails_keeps_the_complete_output(self): + with tempfile.TemporaryDirectory() as directory: + et, owned = self._program_with_owned_data(directory) + et._pte_data = Cord(owned) + path = os.path.join(directory, "model.pte") + with mock.patch.object(os, "replace", side_effect=OSError("busy")): + with self.assertRaises(OSError) as raised: + et.save(path, consume=True) + self.assertFalse(os.path.exists(path)) + (partial,) = [n for n in os.listdir(directory) if n.endswith(".partial")] + partial = os.path.join(os.path.realpath(directory), partial) + with open(partial, "rb") as f: + self.assertEqual(b"blob" * 1024, f.read()) + note = raised.exception.__notes__[-1] + self.assertIn(f"{partial} is complete", note) + self.assertIn("renaming it onto", note) + + def test_a_save_copies_where_the_output_path_cannot_be_checked(self): + calls = [] + + def failing(check): + def checked(p, *args, **kwargs): + if str(p).endswith("model.pte"): + calls.append(p) + raise OSError(errno.ENAMETOOLONG, "File name too long", p) + return check(p, *args, **kwargs) + + return checked + + for module, name in ((os.path, "realpath"), (os, "stat")): + for consume in (False, True): + with self.subTest(check=name, consume=consume): + with tempfile.TemporaryDirectory() as directory: + et, owned = self._program_with_owned_data(directory) + et._pte_data = Cord(owned) + path = os.path.join(directory, "model.pte") + check = failing(getattr(module, name)) + calls.clear() + with mock.patch.object(module, name, check): + et.save(path, consume=consume) + if name == "realpath" and not consume: + # A plain save doesn't look at the path at all. + self.assertEqual([], calls) + with open(path, "rb") as f: + self.assertEqual(b"blob" * 1024, f.read()) + self.assertEqual(b"blob" * 1024, owned.to_bytes()) + + def test_a_consuming_save_on_windows_writes_the_output_directly(self): + with tempfile.TemporaryDirectory() as directory: + et, owned = self._program_with_owned_data(directory) + et._pte_data = Cord(owned) + path = os.path.join(directory, "model.pte") + with mock.patch.object(sys, "platform", "win32"), mock.patch.object( + os, "replace", wraps=os.replace + ) as replace: + et.save(path, consume=True) + replace.assert_not_called() + with open(path, "rb") as f: + self.assertEqual(b"blob" * 1024, f.read()) + self.assertEqual(b"blob" * 1024, owned.to_bytes()) + + @_needs_posix + def test_a_consuming_save_into_a_pipe_copies(self): + with tempfile.TemporaryDirectory() as directory: + et, owned = self._program_with_owned_data(directory) + et._pte_data = Cord(owned) + path = os.path.join(directory, "model.pte") + os.mkfifo(path) + received = [] + + def read(): + with open(path, "rb") as f: + received.append(f.read()) + + reader = threading.Thread(target=read, daemon=True) + reader.start() + et.save(path, consume=True) + reader.join(timeout=60) + self.assertEqual([b"blob" * 1024], received) + self.assertTrue(stat.S_ISFIFO(os.stat(path).st_mode)) + self.assertEqual(b"blob" * 1024, owned.to_bytes()) + + @_needs_posix + def test_a_consuming_save_copies_where_no_temporary_file_fits(self): + if os.geteuid() == 0: + self.skipTest("root can write to any directory") + with tempfile.TemporaryDirectory() as directory: + data_dir, out_dir = (os.path.join(directory, n) for n in ("data", "out")) + os.mkdir(data_dir) + os.mkdir(out_dir) + et, owned = self._program_with_owned_data(data_dir) + et._pte_data = Cord(owned) + path = os.path.join(out_dir, "model.pte") + with open(path, "wb") as f: + f.write(b"old program") + os.chmod(out_dir, 0o555) + try: + et.save(path, consume=True) + self.assertEqual(["model.pte"], os.listdir(out_dir)) + finally: + os.chmod(out_dir, 0o755) + with open(path, "rb") as f: + self.assertEqual(b"blob" * 1024, f.read()) + self.assertEqual(b"blob" * 1024, owned.to_bytes()) + + @_needs_posix + def test_a_consuming_save_refuses_an_output_it_cannot_write(self): + if os.geteuid() == 0: + self.skipTest("root can write to any file") + with tempfile.TemporaryDirectory() as directory: + et, owned = self._program_with_owned_data(directory) + et._pte_data = Cord(owned) + path = os.path.join(directory, "model.pte") + with open(path, "wb") as f: + f.write(b"old program") + os.chmod(path, 0o444) + for consume in (False, True): + with self.subTest(consume=consume): + with self.assertRaises(PermissionError): + et.save(path, consume=consume) + with open(path, "rb") as f: + self.assertEqual(b"old program", f.read()) + self.assertEqual(b"blob" * 1024, owned.to_bytes()) + self.assertFalse([n for n in os.listdir(directory) if "partial" in n]) + + def _save_backend_data_with_consume( + self, directory, payload, external_tag, deep_copy=False + ): + class Sin(torch.nn.Module): + def forward(self, x): + return torch.sin(x) + + owned = [] + + def preprocess(edge_program, compile_specs): + path = os.path.join(directory, "payload.bin") + with open(path, "wb") as f: + f.write(payload) + owned.append(FileBackedData.move_from(path)) + store = NamedDataStore() + store.add_named_data("payload", owned[-1], external_tag=external_tag) + return PreprocessResult( + processed_bytes=b"", + debug_handle_map={}, + data_store_output=store.get_named_data_store_output(), + ) + + with mock.patch.object( + BackendWithNamedDataMap, "preprocess", staticmethod(preprocess) + ): + edge = to_edge_transform_and_lower( + torch.export.export(Sin(), (torch.randn(4),), strict=True), + partitioner=[BackendWithNDMPartitioner()], + ) + # Users deep-copy the edge program to keep one for an ETRecord; either may go + # through to_executorch. + et = (copy.deepcopy(edge) if deep_copy else edge).to_executorch() + # Small chunks, so the payload moves in several of them. + with mock.patch.object(FileBackedData, "_COPY_CHUNK_SIZE", 4096): + et.save(os.path.join(directory, "model.pte"), consume=True) + et.write_tensor_data_to_file(directory, consume=True) + return owned + + def _assert_backend_payload(self, directory, payload, external_tag): + if external_tag is None: + with open(os.path.join(directory, "model.pte"), "rb") as f: + named = deserialize_pte_binary(f.read()).named_data + entry = named.pte_data["payload"] + else: + with open(os.path.join(directory, "model.ptd"), "rb") as f: + named = FlatTensorSerializer().deserialize(Cord(f.read())) + entry = named.named_data["payload"] + self.assertEqual(payload, bytes(named.buffers[entry.buffer_index])) + + @_needs_move + def test_consume_moves_backend_data_through_real_lowering(self): + payload = bytes(range(256)) * 64 + for external_tag in (None, "model"): + with self.subTest(external_tag=external_tag): + with tempfile.TemporaryDirectory() as directory: + owned = self._save_backend_data_with_consume( + directory, payload, external_tag + ) + self.assertEqual([0], [os.path.getsize(d.path) for d in owned]) + self._assert_backend_payload(directory, payload, external_tag) + + def test_consume_is_keyword_only(self): + model = TestLinear() + et = to_edge( + torch.export.export(model, model._get_random_inputs(), strict=True) + ).to_executorch() + with tempfile.TemporaryDirectory() as directory: + with self.assertRaises(TypeError): + et.save(os.path.join(directory, "model.pte"), True) + with self.assertRaises(TypeError): + et.write_tensor_data_to_file(directory, True) + + def test_consume_copies_a_buffer_repeated_in_one_file(self): + with tempfile.TemporaryDirectory() as directory: + et, owned = self._program_with_owned_data(directory) + et._pte_data = Cord(owned) + et._pte_data.append(owned) + et.save(os.path.join(directory, "model.pte"), consume=True) + with open(os.path.join(directory, "model.pte"), "rb") as f: + self.assertEqual(b"blob" * 2048, f.read()) + self.assertEqual(b"blob" * 1024, owned.to_bytes()) + + def test_consume_copies_a_buffer_shared_by_the_pte_and_a_ptd(self): + for pte_first in (True, False): + with self.subTest(pte_first=pte_first): + with tempfile.TemporaryDirectory() as directory: + et, shared = self._program_with_owned_data(directory) + # What deduplication produces when the .pte and a .ptd hold the + # same data. + et._pte_data = Cord(shared) + et._tensor_data = {"data": Cord(shared)} + pte = os.path.join(directory, "model.pte") + if pte_first: + et.save(pte, consume=True) + et.write_tensor_data_to_file(directory, consume=True) + if not pte_first: + et.save(pte, consume=True) + for name in ("model.pte", "data.ptd"): + with open(os.path.join(directory, name), "rb") as f: + self.assertEqual(b"blob" * 1024, f.read()) + self.assertEqual(b"blob" * 1024, shared.to_bytes()) + + def test_save_works_with_a_subclass_that_overrides_write_to_file(self): + def write_to_file(self, open_file): + open_file.write(b"subclass") + + model = TestLinear() + et = to_edge( + torch.export.export(model, model._get_random_inputs(), strict=True) + ).to_executorch() + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "model.pte") + # With nothing to move, a consuming save is a plain one. + for consume in (False, True): + with mock.patch.object( + ExecutorchProgramManager, "write_to_file", write_to_file + ): + et.save(path, consume=consume) + with open(path, "rb") as f: + self.assertEqual(b"subclass", f.read()) + if sys.platform != "win32": + # Where data moves, the save writes the program itself. + moving, owned = self._program_with_owned_data(directory) + moving._pte_data = Cord(owned) + with mock.patch.object( + ExecutorchProgramManager, "write_to_file", write_to_file + ): + moving.save(path, consume=True) + with open(path, "rb") as f: + self.assertEqual(b"blob" * 1024, f.read()) + def test__transform_override_verifiers(self): """Test that _transform can override verifiers in the exported program."""