model_signing.hashing

  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.
 14
 15"""High level API for the hashing interface of `model_signing` library.
 16
 17Hashing is used both for signing and verification and users should ensure that
 18the same configuration is used in both cases.
 19
 20The module could also be used to just hash a single model, without signing it:
 21
 22```python
 23model_signing.hashing.hash(model_path)
 24```
 25
 26This module allows setting up the hashing configuration to a single variable and
 27then sharing it between signing and verification.
 28
 29```python
 30hashing_config = model_signing.hashing.Config().set_ignored_paths(
 31    paths=["README.md"], ignore_git_paths=True
 32)
 33
 34signing_config = (
 35    model_signing.signing.Config()
 36    .use_elliptic_key_signer(private_key="key")
 37    .set_hashing_config(hashing_config)
 38)
 39
 40verifying_config = (
 41    model_signing.verifying.Config()
 42    .use_elliptic_key_verifier(public_key="key.pub")
 43    .set_hashing_config(hashing_config)
 44)
 45```
 46
 47The API defined here is stable and backwards compatible.
 48"""
 49
 50from collections.abc import Callable, Iterable
 51import os
 52import pathlib
 53from typing import Literal, Self, TypeAlias
 54
 55import blake3
 56
 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
 63
 64
 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
 69
 70
 71def hash(model_path: PathLike) -> manifest.Manifest:
 72    """Hashes a model using the default configuration.
 73
 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.
 77
 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.
 85
 86    Args:
 87        model_path: The path to the model to hash.
 88
 89    Returns:
 90        A manifest of the hashed model.
 91    """
 92    return Config().hash(model_path)
 93
 94
 95class Config:
 96    """Configuration to use when hashing models.
 97
 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.
101
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.
109
110    This configuration class supports configuring the hashing granularity. By
111    default, we hash at file level granularity.
112
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.
116
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    """
123
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
130
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)
150
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            )
163
164        self._serializer.set_allow_symlinks(self._allow_symlinks)
165
166        return self._serializer.serialize(
167            pathlib.Path(model_path),
168            ignore_paths=ignored_paths,
169            files_to_hash=files_to_hash,
170        )
171
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.
177
178        Args:
179            hashing_algorithm: The hashing algorithm to use.
180
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                )
195
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.
203
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.
211
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
218
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)
224
225        return _factory
226
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.
234
235        This is not recommended for BLAKE3 because it is not necessary. BLAKE3
236        already operates in parallel.
237
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.
244
245        Returns:
246            The hasher factory that should be used by the active serialization
247            method.
248        """
249
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            )
262
263        return _factory
264
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.
275
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.
279
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.
295
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
308
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.
320
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.
324
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.
329
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.
343
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            )
355
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
365
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.
370
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.
374
375        If an ignored path is a directory, serialization will ignore both the
376        path and any of its children.
377
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.
382
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
391
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.
396
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
413
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