Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Copyright 2024 The Sigstore Authors
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
15"""Model serializers that operate at file shard level of granularity."""
17from collections.abc import Callable, Iterable
18import concurrent.futures
19import itertools
20import os
21import pathlib
23from typing_extensions import override
25from model_signing import manifest
26from model_signing._hashing import io
27from model_signing._serialization import serialization
30def _endpoints(step: int, end: int) -> Iterable[int]:
31 """Yields numbers from `step` to `end` inclusive, spaced by `step`.
33 Last value is always equal to `end`, even when `end` is not a multiple of
34 `step`. There is always a value returned.
36 Examples:
37 ```python
38 >>> list(_endpoints(2, 8))
39 [2, 4, 6, 8]
40 >>> list(_endpoints(2, 9))
41 [2, 4, 6, 8, 9]
42 >>> list(_endpoints(2, 2))
43 [2]
44 >>> list(_endpoints(2, 0))
45 [0]
47 Yields:
48 Values in the range, from `step` and up to `end`.
49 """
50 yield from range(step, end, step)
51 yield end
54class Serializer(serialization.Serializer):
55 """Model serializer that produces a manifest recording every file shard.
57 Traverses the model directory and creates digests for every file found,
58 sharding the file in equal shards and computing the digests in parallel.
59 """
61 def __init__(
62 self,
63 sharded_hasher_factory: Callable[
64 [pathlib.Path, int, int], io.ShardedFileHasher
65 ],
66 *,
67 max_workers: int | None = None,
68 allow_symlinks: bool = False,
69 ignore_paths: Iterable[pathlib.Path] = frozenset(),
70 ):
71 """Initializes an instance to serialize a model with this serializer.
73 Args:
74 sharded_hasher_factory: A callable to build the hash engine used to
75 hash every shard of the files in the model. Because each shard is
76 processed in parallel, every thread needs to call the factory to
77 start hashing. The arguments are the file, and the endpoints of
78 the shard.
79 max_workers: Maximum number of workers to use in parallel. Default
80 is to defer to the `concurrent.futures` library.
81 allow_symlinks: Controls whether symbolic links are included. If a
82 symlink is present but the flag is `False` (default) the
83 serialization would raise an error.
84 ignore_paths: The paths of files to ignore.
85 """
86 self._hasher_factory = sharded_hasher_factory
87 self._max_workers = max_workers
88 self._allow_symlinks = allow_symlinks
89 self._ignore_paths = ignore_paths
91 # Precompute some private values only once by using a mock file hasher.
92 # None of the arguments used to build the hasher are used.
93 hasher = sharded_hasher_factory(pathlib.Path(), 0, 1)
94 self._shard_size = hasher.shard_size
95 self._serialization_description = manifest._ShardSerialization(
96 # Here we need the internal hasher name, not the mangled name.
97 # This name is used when guessing the hashing configuration.
98 hasher._content_hasher.digest_name,
99 self._shard_size,
100 self._allow_symlinks,
101 self._ignore_paths,
102 )
104 def set_allow_symlinks(self, allow_symlinks: bool) -> None:
105 """Set whether following symlinks is allowed."""
106 self._allow_symlinks = allow_symlinks
107 hasher = self._hasher_factory(pathlib.Path(), 0, 1)
108 self._serialization_description = manifest._ShardSerialization(
109 hasher._content_hasher.digest_name,
110 self._shard_size,
111 self._allow_symlinks,
112 self._ignore_paths,
113 )
115 @override
116 def serialize(
117 self,
118 model_path: pathlib.Path,
119 *,
120 ignore_paths: Iterable[pathlib.Path] = frozenset(),
121 files_to_hash: Iterable[pathlib.Path] | None = None,
122 ) -> manifest.Manifest:
123 """Serializes the model given by the `model_path` argument.
125 Args:
126 model_path: The path to the model.
127 ignore_paths: The paths to ignore during serialization. If a
128 provided path is a directory, all children of the directory are
129 ignored.
130 files_to_hash: Optional list of files to hash; ignore all others
132 Returns:
133 The model's serialized manifest.
135 Raises:
136 ValueError: The model contains a symbolic link, but the serializer
137 was not initialized with `allow_symlinks=True`.
138 """
139 shards = []
140 file_count = 0
141 # TODO: github.com/sigstore/model-transparency/issues/200 - When
142 # Python3.12 is the minimum supported version, the glob can be replaced
143 # with `pathlib.Path.walk` for a clearer interface, and some speed
144 # improvement.
145 if files_to_hash is None:
146 files_to_hash = itertools.chain(
147 (model_path,), model_path.glob("**/*")
148 )
149 for path in files_to_hash:
150 if serialization.should_ignore(path, ignore_paths):
151 continue
152 if (
153 self._allow_symlinks
154 and path.is_symlink()
155 and not serialization.check_symlink(path, model_path)
156 ):
157 continue
158 serialization.check_file_or_directory(
159 path, allow_symlinks=self._allow_symlinks
160 )
161 if path.is_file():
162 file_count += 1
163 shards.extend(self._get_shards(path))
164 elif path != model_path and path.is_symlink():
165 # The walk never descends into a symlink to a directory, so
166 # accepting it would leave every file under it out of the
167 # manifest without any error.
168 raise ValueError(
169 f"Cannot use '{path}' because it is a symlink to a"
170 " directory. Files under it would not be serialized."
171 )
173 if file_count == 0:
174 raise ValueError(
175 "Model contains no regular files after applying exclusions"
176 )
178 if not shards:
179 raise ValueError(
180 "Model contains only empty files which produce no shards"
181 )
183 manifest_items = []
184 with concurrent.futures.ThreadPoolExecutor(
185 max_workers=self._max_workers
186 ) as tpe:
187 futures = [
188 tpe.submit(self._compute_hash, model_path, path, start, end)
189 for path, start, end in shards
190 ]
191 for future in concurrent.futures.as_completed(futures):
192 manifest_items.append(future.result())
194 # Recreate serialization_description for new ignore_paths
195 if ignore_paths:
196 rel_ignore_paths = []
197 for p in ignore_paths:
198 rp = os.path.relpath(p, model_path)
199 # rp may start with "../" if it is not relative to model_path
200 if not rp.startswith("../"):
201 rel_ignore_paths.append(pathlib.Path(rp))
203 hasher = self._hasher_factory(pathlib.Path(), 0, 1)
204 self._serialization_description = manifest._ShardSerialization(
205 hasher._content_hasher.digest_name,
206 self._shard_size,
207 self._allow_symlinks,
208 frozenset(list(self._ignore_paths) + rel_ignore_paths),
209 )
211 model_name = model_path.name
212 if not model_name or model_name == "..":
213 model_name = os.path.basename(model_path.resolve())
215 return manifest.Manifest(
216 model_name, manifest_items, self._serialization_description
217 )
219 def _get_shards(
220 self, path: pathlib.Path
221 ) -> list[tuple[pathlib.Path, int, int]]:
222 """Determines the shards of a given file path.
224 An empty file gets a single empty shard, so that it is still recorded in
225 the manifest and covered by the signature.
226 """
227 shards = []
228 path_size = path.stat().st_size
229 start = 0
230 for end in _endpoints(self._shard_size, path_size):
231 shards.append((path, start, end))
232 start = end
233 return shards
235 def _compute_hash(
236 self, model_path: pathlib.Path, path: pathlib.Path, start: int, end: int
237 ) -> manifest.ShardedFileManifestItem:
238 """Produces the manifest item of the file given by `path`.
240 Args:
241 model_path: The path to the model.
242 path: Path to the file in the model, that is currently transformed
243 to a manifest item.
244 start: The start offset of the shard (included).
245 end: The end offset of the shard (not included).
247 Returns:
248 The itemized manifest.
249 """
250 relative_path = path.relative_to(model_path)
251 digest = self._hasher_factory(path, start, end).compute()
252 return manifest.ShardedFileManifestItem(
253 path=relative_path, digest=digest, start=start, end=end
254 )