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"""High level API for the hashing interface of `model_signing` library.
17Hashing is used both for signing and verification and users should ensure that
18the same configuration is used in both cases.
20The module could also be used to just hash a single model, without signing it:
22```python
23model_signing.hashing.hash(model_path)
24```
26This module allows setting up the hashing configuration to a single variable and
27then sharing it between signing and verification.
29```python
30hashing_config = model_signing.hashing.Config().set_ignored_paths(
31 paths=["README.md"], ignore_git_paths=True
32)
34signing_config = (
35 model_signing.signing.Config()
36 .use_elliptic_key_signer(private_key="key")
37 .set_hashing_config(hashing_config)
38)
40verifying_config = (
41 model_signing.verifying.Config()
42 .use_elliptic_key_verifier(public_key="key.pub")
43 .set_hashing_config(hashing_config)
44)
45```
47The API defined here is stable and backwards compatible.
48"""
50from collections.abc import Callable, Iterable
51import os
52import pathlib
53from typing import Literal, Self, TypeAlias
55import blake3
57from model_signing import manifest
58from model_signing._hashing import hashing
59from model_signing._hashing import io
60from model_signing._hashing import memory
61from model_signing._serialization import file
62from model_signing._serialization import file_shard
65# Type alias to support `os.PathLike`, `str` and `bytes` objects in the API
66# `TypeAlias` is deprecated in Python 3.12 in favor of `type`, which we can use
67# once Python 3.12 is the minimum supported version
68PathLike: TypeAlias = str | bytes | os.PathLike
71def hash(model_path: PathLike) -> manifest.Manifest:
72 """Hashes a model using the default configuration.
74 Hashing is the shared part between signing and verification and is also
75 expected to be the slowest component. When serializing a model, we need to
76 spend time proportional to the model size on disk.
78 This method returns a "manifest" of the model. A manifest is a collection of
79 every object in the model, paired with the corresponding hash. Currently, we
80 consider an object in the model to be either a file or a shard of the file.
81 Large models with large files will be hashed much faster when every shard is
82 hashed in parallel, at the cost of generating a larger payload for the
83 signature. In future releases we could support hashing individual tensors or
84 tensor slices for further speed optimizations for very large models.
86 Args:
87 model_path: The path to the model to hash.
89 Returns:
90 A manifest of the hashed model.
91 """
92 return Config().hash(model_path)
95class Config:
96 """Configuration to use when hashing models.
98 Hashing is the shared part between signing and verification and is also
99 expected to be the slowest component. When serializing a model, we need to
100 spend time proportional to the model size on disk.
102 Hashing builds a "manifest" of the model. A manifest is a collection of
103 every object in the model, paired with the corresponding hash. Currently, we
104 consider an object in the model to be either a file or a shard of the file.
105 Large models with large files will be hashed much faster when every shard is
106 hashed in parallel, at the cost of generating a larger payload for the
107 signature. In future releases we could support hashing individual tensors or
108 tensor slices for further speed optimizations for very large models.
110 This configuration class supports configuring the hashing granularity. By
111 default, we hash at file level granularity.
113 This configuration class also supports configuring the hash method used to
114 generate the hash for every object in the model. We currently support
115 SHA256, BLAKE2, and BLAKE3, with SHA256 being the default.
117 This configuration class also supports configuring which paths from the
118 model directory should be ignored. These are files that doesn't impact the
119 behavior of the model, or files that won't be distributed with the model. By
120 default, only files that are associated with a git repository (`.git`,
121 `.gitattributes`, `.gitignore`, etc.) are ignored.
122 """
124 def __init__(self):
125 """Initializes the default configuration for hashing."""
126 self._ignored_paths = frozenset()
127 self._ignore_git_paths = True
128 self.use_file_serialization()
129 self._allow_symlinks = False
131 def hash(
132 self,
133 model_path: PathLike,
134 *,
135 files_to_hash: Iterable[PathLike] | None = None,
136 ) -> manifest.Manifest:
137 """Hashes a model using the current configuration."""
138 # All paths in ``_ignored_paths`` are expected to be relative to the
139 # model directory. Join them to ``model_path`` and ensure they do not
140 # escape it.
141 model_path = pathlib.Path(model_path)
142 ignored_paths = []
143 for p in self._ignored_paths:
144 full = model_path / p
145 try:
146 full.relative_to(model_path)
147 except ValueError:
148 continue
149 ignored_paths.append(full)
151 if self._ignore_git_paths:
152 ignored_paths.extend(
153 [
154 model_path / p
155 for p in [
156 ".git/",
157 ".gitattributes",
158 ".github/",
159 ".gitignore",
160 ]
161 ]
162 )
164 self._serializer.set_allow_symlinks(self._allow_symlinks)
166 return self._serializer.serialize(
167 pathlib.Path(model_path),
168 ignore_paths=ignored_paths,
169 files_to_hash=files_to_hash,
170 )
172 def _build_stream_hasher(
173 self,
174 hashing_algorithm: Literal["sha256", "blake2", "blake3"] = "sha256",
175 ) -> hashing.StreamingHashEngine:
176 """Builds a streaming hasher from a constant string.
178 Args:
179 hashing_algorithm: The hashing algorithm to use.
181 Returns:
182 An instance of the requested hasher.
183 """
184 match hashing_algorithm:
185 case "sha256":
186 return memory.SHA256()
187 case "blake2":
188 return memory.BLAKE2()
189 case "blake3":
190 return memory.BLAKE3()
191 case _:
192 raise ValueError(
193 f"Unsupported hashing method {hashing_algorithm}"
194 )
196 def _build_file_hasher_factory(
197 self,
198 hashing_algorithm: Literal["sha256", "blake2", "blake3"] = "sha256",
199 chunk_size: int = 1048576,
200 max_workers: int | None = None,
201 ) -> Callable[[pathlib.Path], io.FileHasher]:
202 """Builds the hasher factory for a serialization by file.
204 Args:
205 hashing_algorithm: The hashing algorithm to use to hash a file.
206 chunk_size: The amount of file to read at once. Default is 1MB. A
207 special value of 0 signals to attempt to read everything in a
208 single call. This is ignored for BLAKE3.
209 max_workers: Maximum number of workers to use in parallel. Defaults
210 to the number of logical cores. Only relevant for BLAKE3.
212 Returns:
213 The hasher factory that should be used by the active serialization
214 method.
215 """
216 if max_workers is None:
217 max_workers = blake3.blake3.AUTO
219 def _factory(path: pathlib.Path) -> io.FileHasher:
220 if hashing_algorithm == "blake3":
221 return io.Blake3FileHasher(path, max_threads=max_workers)
222 hasher = self._build_stream_hasher(hashing_algorithm)
223 return io.SimpleFileHasher(path, hasher, chunk_size=chunk_size)
225 return _factory
227 def _build_sharded_file_hasher_factory(
228 self,
229 hashing_algorithm: Literal["sha256", "blake2"] = "sha256",
230 chunk_size: int = 1048576,
231 shard_size: int = 1_000_000_000,
232 ) -> Callable[[pathlib.Path, int, int], io.ShardedFileHasher]:
233 """Builds the hasher factory for a serialization by file shards.
235 This is not recommended for BLAKE3 because it is not necessary. BLAKE3
236 already operates in parallel.
238 Args:
239 hashing_algorithm: The hashing algorithm to use to hash a shard.
240 chunk_size: The amount of file to read at once. Default is 1MB. A
241 special value of 0 signals to attempt to read everything in a
242 single call.
243 shard_size: The size of a file shard. Default is 1 GB.
245 Returns:
246 The hasher factory that should be used by the active serialization
247 method.
248 """
250 def _factory(
251 path: pathlib.Path, start: int, end: int
252 ) -> io.ShardedFileHasher:
253 hasher = self._build_stream_hasher(hashing_algorithm)
254 return io.ShardedFileHasher(
255 path,
256 hasher,
257 start=start,
258 end=end,
259 chunk_size=chunk_size,
260 shard_size=shard_size,
261 )
263 return _factory
265 def use_file_serialization(
266 self,
267 *,
268 hashing_algorithm: Literal["sha256", "blake2", "blake3"] = "sha256",
269 chunk_size: int = 1048576,
270 max_workers: int | None = None,
271 allow_symlinks: bool = False,
272 ignore_paths: Iterable[pathlib.Path] = frozenset(),
273 ) -> Self:
274 """Configures serialization to build a manifest of (file, hash) pairs.
276 The serialization method in this configuration is changed to one where
277 every file in the model is paired with its digest and a manifest
278 containing all these pairings is being built.
280 Args:
281 hashing_algorithm: The hashing algorithm to use to hash a file.
282 chunk_size: The amount of file to read at once. Default is 1MB. A
283 special value of 0 signals to attempt to read everything in a
284 single call. Ignored for BLAKE3.
285 max_workers: Maximum number of workers to use in parallel. Default
286 is to defer to the `concurrent.futures` library to select the best
287 value for the current machine, or the number of logical cores
288 when doing BLAKE3 hashing. When reading files off of slower
289 hardware like an HDD rather than an SSD, and using BLAKE3,
290 setting max_workers to 1 may improve performance.
291 allow_symlinks: Controls whether symbolic links are included. If a
292 symlink is present but the flag is `False` (default) the
293 serialization would raise an error.
294 ignore_paths: Paths of files to ignore during serialization.
296 Returns:
297 The new hashing configuration with the new serialization method.
298 """
299 self._serializer = file.Serializer(
300 self._build_file_hasher_factory(
301 hashing_algorithm, chunk_size, max_workers
302 ),
303 max_workers=max_workers,
304 allow_symlinks=allow_symlinks,
305 ignore_paths=ignore_paths,
306 )
307 return self
309 def use_shard_serialization(
310 self,
311 *,
312 hashing_algorithm: Literal["sha256", "blake2", "blake3"] = "sha256",
313 chunk_size: int = 1048576,
314 shard_size: int = 1_000_000_000,
315 max_workers: int | None = None,
316 allow_symlinks: bool = False,
317 ignore_paths: Iterable[pathlib.Path] = frozenset(),
318 ) -> Self:
319 """Configures serialization to build a manifest of (shard, hash) pairs.
321 For BLAKE3 this is equivalent to not sharding. Sharding is bypassed
322 because BLAKE3 already operates in parallel. This means the chunk_size
323 and shard_size args are ignored.
325 The serialization method in this configuration is changed to one where
326 every file in the model is sharded in equal sized shards, every shard is
327 paired with its digest and a manifest containing all these pairings is
328 being built.
330 Args:
331 hashing_algorithm: The hashing algorithm to use to hash a shard.
332 chunk_size: The amount of file to read at once. Default is 1MB. A
333 special value of 0 signals to attempt to read everything in a
334 single call.
335 shard_size: The size of a file shard. Default is 1 GB.
336 max_workers: Maximum number of workers to use in parallel. Default
337 is to defer to the `concurrent.futures` library to select the best
338 value for the current machine.
339 allow_symlinks: Controls whether symbolic links are included. If a
340 symlink is present but the flag is `False` (default) the
341 serialization would raise an error.
342 ignore_paths: Paths of files to ignore.
344 Returns:
345 The new hashing configuration with the new serialization method.
346 """
347 if hashing_algorithm == "blake3":
348 return self.use_file_serialization(
349 hashing_algorithm=hashing_algorithm,
350 chunk_size=chunk_size,
351 max_workers=max_workers,
352 allow_symlinks=allow_symlinks,
353 ignore_paths=ignore_paths,
354 )
356 self._serializer = file_shard.Serializer(
357 self._build_sharded_file_hasher_factory(
358 hashing_algorithm, chunk_size, shard_size
359 ),
360 max_workers=max_workers,
361 allow_symlinks=allow_symlinks,
362 ignore_paths=ignore_paths,
363 )
364 return self
366 def set_ignored_paths(
367 self, *, paths: Iterable[PathLike], ignore_git_paths: bool = True
368 ) -> Self:
369 """Configures the paths to be ignored during serialization of a model.
371 If the model is a single file, there are no paths that are ignored. If
372 the model is a directory, all paths are considered as relative to the
373 model directory, since we never look at files outside of it.
375 If an ignored path is a directory, serialization will ignore both the
376 path and any of its children.
378 Args:
379 paths: The paths to ignore.
380 ignore_git_paths: Whether to ignore git related paths (default) or
381 include them in the signature.
383 Returns:
384 The new hashing configuration with a new set of ignored paths.
385 """
386 # Preserve the user-provided relative paths; they are resolved against
387 # the model directory later when hashing.
388 self._ignored_paths = frozenset(pathlib.Path(p) for p in paths)
389 self._ignore_git_paths = ignore_git_paths
390 return self
392 def add_ignored_paths(
393 self, *, model_path: PathLike, paths: Iterable[PathLike]
394 ) -> None:
395 """Add more paths to ignore to existing set of paths.
397 Args:
398 model_path: The path to the model
399 paths: Additional paths to ignore. All path must be relative to
400 the model directory.
401 """
402 newset = set(self._ignored_paths)
403 model_path = pathlib.Path(model_path)
404 for p in paths:
405 candidate = pathlib.Path(p)
406 full = model_path / candidate
407 try:
408 full.relative_to(model_path)
409 except ValueError:
410 continue
411 newset.add(candidate)
412 self._ignored_paths = newset
414 def set_allow_symlinks(self, allow_symlinks: bool) -> Self:
415 """Set whether following symlinks is allowed."""
416 self._allow_symlinks = allow_symlinks
417 return self