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
 53import sys
 54from typing import Literal, Optional, Union
 55
 56from model_signing import manifest
 57from model_signing._hashing import hashing
 58from model_signing._hashing import io
 59from model_signing._hashing import memory
 60from model_signing._serialization import file
 61from model_signing._serialization import file_shard
 62
 63
 64if sys.version_info >= (3, 11):
 65    from typing import Self
 66else:
 67    from typing_extensions import Self
 68
 69
 70# `TypeAlias` only exists from Python 3.10
 71# `TypeAlias` is deprecated in Python 3.12 in favor of `type`
 72if sys.version_info >= (3, 10):
 73    from typing import TypeAlias
 74else:
 75    from typing_extensions import TypeAlias
 76
 77
 78# Type alias to support `os.PathLike`, `str` and `bytes` objects in the API
 79# When Python 3.12 is the minimum supported version we can use `type`
 80# When Python 3.11 is the minimum supported version we can use `|`
 81PathLike: TypeAlias = Union[str, bytes, os.PathLike]
 82
 83
 84def hash(model_path: PathLike) -> manifest.Manifest:
 85    """Hashes a model using the default configuration.
 86
 87    Hashing is the shared part between signing and verification and is also
 88    expected to be the slowest component. When serializing a model, we need to
 89    spend time proportional to the model size on disk.
 90
 91    This method returns a "manifest" of the model. A manifest is a collection of
 92    every object in the model, paired with the corresponding hash. Currently, we
 93    consider an object in the model to be either a file or a shard of the file.
 94    Large models with large files will be hashed much faster when every shard is
 95    hashed in parallel, at the cost of generating a larger payload for the
 96    signature. In future releases we could support hashing individual tensors or
 97    tensor slices for further speed optimizations for very large models.
 98
 99    Args:
100        model_path: The path to the model to hash.
101
102    Returns:
103        A manifest of the hashed model.
104    """
105    return Config().hash(model_path)
106
107
108class Config:
109    """Configuration to use when hashing models.
110
111    Hashing is the shared part between signing and verification and is also
112    expected to be the slowest component. When serializing a model, we need to
113    spend time proportional to the model size on disk.
114
115    Hashing builds a "manifest" of the model. A manifest is a collection of
116    every object in the model, paired with the corresponding hash. Currently, we
117    consider an object in the model to be either a file or a shard of the file.
118    Large models with large files will be hashed much faster when every shard is
119    hashed in parallel, at the cost of generating a larger payload for the
120    signature. In future releases we could support hashing individual tensors or
121    tensor slices for further speed optimizations for very large models.
122
123    This configuration class supports configuring the hashing granularity. By
124    default, we hash at file level granularity.
125
126    This configuration class also supports configuring the hash method used to
127    generate the hash for every object in the model. We currently support SHA256
128    and BLAKE2, with SHA256 being the default.
129
130    This configuration class also supports configuring which paths from the
131    model directory should be ignored. These are files that doesn't impact the
132    behavior of the model, or files that won't be distributed with the model. By
133    default, only files that are associated with a git repository (`.git`,
134    `.gitattributes`, `.gitignore`, etc.) are ignored.
135    """
136
137    def __init__(self):
138        """Initializes the default configuration for hashing."""
139        self._ignored_paths = frozenset()
140        self._ignore_git_paths = True
141        self.use_file_serialization()
142        self._allow_symlinks = False
143
144    def hash(self, model_path: PathLike) -> manifest.Manifest:
145        """Hashes a model using the current configuration."""
146        # All paths in ignored_paths must have model_path as prefix
147        ignored_paths = []
148        for p in self._ignored_paths:
149            rp = os.path.relpath(p, model_path)
150            # rp may start with "../" if it is not relative to model_path
151            if not rp.startswith("../"):
152                ignored_paths.append(pathlib.Path(os.path.join(model_path, rp)))
153
154        if self._ignore_git_paths:
155            ignored_paths.extend(
156                [
157                    os.path.join(model_path, p)
158                    for p in [
159                        ".git/",
160                        ".gitattributes",
161                        ".github/",
162                        ".gitignore",
163                    ]
164                ]
165            )
166
167        self._serializer.set_allow_symlinks(self._allow_symlinks)
168
169        return self._serializer.serialize(
170            pathlib.Path(model_path), ignore_paths=ignored_paths
171        )
172
173    def _build_stream_hasher(
174        self, hashing_algorithm: Literal["sha256", "blake2"] = "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        # TODO: Once Python 3.9 support is deprecated revert to using `match`
185        if hashing_algorithm == "sha256":
186            return memory.SHA256()
187        if hashing_algorithm == "blake2":
188            return memory.BLAKE2()
189
190        raise ValueError(f"Unsupported hashing method {hashing_algorithm}")
191
192    def _build_file_hasher_factory(
193        self,
194        hashing_algorithm: Literal["sha256", "blake2"] = "sha256",
195        chunk_size: int = 1048576,
196    ) -> Callable[[pathlib.Path], io.SimpleFileHasher]:
197        """Builds the hasher factory for a serialization by file.
198
199        Args:
200            hashing_algorithm: The hashing algorithm to use to hash a file.
201            chunk_size: The amount of file to read at once. Default is 1MB. A
202              special value of 0 signals to attempt to read everything in a
203              single call.
204
205        Returns:
206            The hasher factory that should be used by the active serialization
207            method.
208        """
209
210        def _factory(path: pathlib.Path) -> io.SimpleFileHasher:
211            hasher = self._build_stream_hasher(hashing_algorithm)
212            return io.SimpleFileHasher(path, hasher, chunk_size=chunk_size)
213
214        return _factory
215
216    def _build_sharded_file_hasher_factory(
217        self,
218        hashing_algorithm: Literal["sha256", "blake2"] = "sha256",
219        chunk_size: int = 1048576,
220        shard_size: int = 1_000_000_000,
221    ) -> Callable[[pathlib.Path, int, int], io.ShardedFileHasher]:
222        """Builds the hasher factory for a serialization by file shards.
223
224        Args:
225            hashing_algorithm: The hashing algorithm to use to hash a shard.
226            chunk_size: The amount of file to read at once. Default is 1MB. A
227              special value of 0 signals to attempt to read everything in a
228              single call.
229            shard_size: The size of a file shard. Default is 1 GB.
230
231        Returns:
232            The hasher factory that should be used by the active serialization
233            method.
234        """
235
236        def _factory(
237            path: pathlib.Path, start: int, end: int
238        ) -> io.ShardedFileHasher:
239            hasher = self._build_stream_hasher(hashing_algorithm)
240            return io.ShardedFileHasher(
241                path,
242                hasher,
243                start=start,
244                end=end,
245                chunk_size=chunk_size,
246                shard_size=shard_size,
247            )
248
249        return _factory
250
251    def use_file_serialization(
252        self,
253        *,
254        hashing_algorithm: Literal["sha256", "blake2"] = "sha256",
255        chunk_size: int = 1048576,
256        max_workers: Optional[int] = None,
257        allow_symlinks: bool = False,
258        ignore_paths: Iterable[pathlib.Path] = frozenset(),
259    ) -> Self:
260        """Configures serialization to build a manifest of (file, hash) pairs.
261
262        The serialization method in this configuration is changed to one where
263        every file in the model is paired with its digest and a manifest
264        containing all these pairings is being built.
265
266        Args:
267            hashing_algorithm: The hashing algorithm to use to hash a file.
268            chunk_size: The amount of file to read at once. Default is 1MB. A
269              special value of 0 signals to attempt to read everything in a
270              single call.
271            max_workers: Maximum number of workers to use in parallel. Default
272              is to defer to the `concurrent.futures` library to select the best
273              value for the current machine.
274            allow_symlinks: Controls whether symbolic links are included. If a
275              symlink is present but the flag is `False` (default) the
276              serialization would raise an error.
277
278        Returns:
279            The new hashing configuration with the new serialization method.
280        """
281        self._serializer = file.Serializer(
282            self._build_file_hasher_factory(hashing_algorithm, chunk_size),
283            max_workers=max_workers,
284            allow_symlinks=allow_symlinks,
285            ignore_paths=ignore_paths,
286        )
287        return self
288
289    def use_shard_serialization(
290        self,
291        *,
292        hashing_algorithm: Literal["sha256", "blake2"] = "sha256",
293        chunk_size: int = 1048576,
294        shard_size: int = 1_000_000_000,
295        max_workers: Optional[int] = None,
296        allow_symlinks: bool = False,
297        ignore_paths: Iterable[pathlib.Path] = frozenset(),
298    ) -> Self:
299        """Configures serialization to build a manifest of (shard, hash) pairs.
300
301        The serialization method in this configuration is changed to one where
302        every file in the model is sharded in equal sized shards, every shard is
303        paired with its digest and a manifest containing all these pairings is
304        being built.
305
306        Args:
307            hashing_algorithm: The hashing algorithm to use to hash a shard.
308            chunk_size: The amount of file to read at once. Default is 1MB. A
309              special value of 0 signals to attempt to read everything in a
310              single call.
311            shard_size: The size of a file shard. Default is 1 GB.
312            max_workers: Maximum number of workers to use in parallel. Default
313              is to defer to the `concurrent.futures` library to select the best
314              value for the current machine.
315            allow_symlinks: Controls whether symbolic links are included. If a
316              symlink is present but the flag is `False` (default) the
317              serialization would raise an error.
318            ignore_paths: Paths of files to ignore.
319
320        Returns:
321            The new hashing configuration with the new serialization method.
322        """
323        self._serializer = file_shard.Serializer(
324            self._build_sharded_file_hasher_factory(
325                hashing_algorithm, chunk_size, shard_size
326            ),
327            max_workers=max_workers,
328            allow_symlinks=allow_symlinks,
329            ignore_paths=ignore_paths,
330        )
331        return self
332
333    def set_ignored_paths(
334        self, *, paths: Iterable[PathLike], ignore_git_paths: bool = True
335    ) -> Self:
336        """Configures the paths to be ignored during serialization of a model.
337
338        If the model is a single file, there are no paths that are ignored. If
339        the model is a directory, all paths are considered as relative to the
340        model directory, since we never look at files outside of it.
341
342        If an ignored path is a directory, serialization will ignore both the
343        path and any of its children.
344
345        Args:
346            paths: The paths to ignore.
347            ignore_git_paths: Whether to ignore git related paths (default) or
348              include them in the signature.
349
350        Returns:
351            The new hashing configuration with a new set of ignored paths.
352        """
353        # Use relpath to possibly fix weird paths like '../a/b' -> 'b'
354        # when '../a/' is a no-op
355        self._ignored_paths = frozenset(
356            {pathlib.Path(p).resolve() for p in paths}
357        )
358        self._ignore_git_paths = ignore_git_paths
359        return self
360
361    def add_ignored_paths(
362        self, *, model_path: PathLike, paths: Iterable[PathLike]
363    ) -> None:
364        """Add more paths to ignore to existing set of paths.
365
366        Args:
367            model_path: The path to the model
368            paths: Additional paths to ignore. All path must be relative to
369                   the model directory.
370        """
371        newset = set(self._ignored_paths)
372        newset.update([os.path.join(model_path, p) for p in paths])
373        self._ignored_paths = newset
374
375    def set_allow_symlinks(self, allow_symlinks: bool) -> Self:
376        """Set whether following symlinks is allowed."""
377        self._allow_symlinks = allow_symlinks
378        return self