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 e03b27b09ec..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,28 +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 _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) -> None: + 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") - try: - with open(path, "wb") as 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}") - except Exception as e: - logging.error(f"Error while saving to {path}: {e}") + logging.info(f"Saved exported program to {path}") diff --git a/exir/program/test/BUCK b/exir/program/test/BUCK index f316eb83f6a..b9e0f627786 100644 --- a/exir/program/test/BUCK +++ b/exir/program/test/BUCK @@ -18,12 +18,14 @@ 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", "//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..519b582b96f 100644 --- a/exir/program/test/test_program.py +++ b/exir/program/test/test_program.py @@ -7,11 +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, 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, @@ -31,6 +46,8 @@ 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.flat_tensor.serialize.serialize import FlatTensorSerializer from executorch.extension.pybindings.portable_lib import ( _load_for_executorch_from_buffer, ) @@ -43,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__() @@ -1196,6 +1221,559 @@ 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 _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.""" 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