Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/python/saved_model/fingerprinting_utils.py: 37%

1# Copyright 2022 The TensorFlow Authors. All Rights Reserved. 

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"""Utilities for SavedModel fingerprinting. 

16 

17This module contains utility classes and functions for working with the 

18SavedModel fingerprint. 

19""" 

20 

21from absl import logging 

22 

23from tensorflow.core.config import flags 

24from tensorflow.core.protobuf import fingerprint_pb2 

25from tensorflow.python.lib.io import file_io 

26from tensorflow.python.saved_model import fingerprinting 

27from tensorflow.python.saved_model.pywrap_saved_model import constants 

28from tensorflow.python.saved_model.pywrap_saved_model import fingerprinting as fingerprinting_pywrap 

29from tensorflow.python.saved_model.pywrap_saved_model import metrics 

30from tensorflow.python.util import compat 

31 

32 

33def write_fingerprint(export_dir, saved_model_serialized): 

34 """Write fingerprint protobuf, if requested. 

35 

36 Writes a `tf.saved_model.experimental.Fingerprint` object to a 

37 `fingerprint.pb` file in the `export_dir`. 

38 

39 Args: 

40 export_dir: The directory in which to write the fingerprint. 

41 saved_model_serialized: The serialized SavedModel proto. 

42 """ 

43 

44 if flags.config().saved_model_fingerprinting.value(): 

45 fingerprint_path = file_io.join( 

46 compat.as_str(export_dir), 

47 compat.as_str(constants.FINGERPRINT_FILENAME)) 

48 logging.info("Writing fingerprint to %s", fingerprint_path) 

49 try: 

50 fingerprint_serialized = fingerprinting_pywrap.CreateFingerprintDef( 

51 saved_model_serialized, export_dir) 

52 except fingerprinting_pywrap.FingerprintException as e: 

53 raise ValueError(e) from None 

54 file_io.atomic_write_string_to_file(fingerprint_path, 

55 fingerprint_serialized) 

56 # We need to deserialize the fingerprint in order to send its values. 

57 fingerprint_proto = fingerprint_pb2.FingerprintDef() 

58 fingerprint_proto.ParseFromString(fingerprint_serialized) 

59 metrics.SetWriteFingerprint(fingerprint=fingerprint_serialized) 

60 fingerprint = fingerprinting.Fingerprint.from_proto(fingerprint_serialized) 

61 metrics.SetWritePathAndSingleprint(path=export_dir, 

62 singleprint=fingerprint.singleprint()) 

63 

64 

65def singleprint_from_saved_model(export_dir, saved_model_serialized): 

66 """Returns the singleprint of a SavedModel in `export_dir`. 

67 

68 Args: 

69 export_dir: The directory that contains the SavedModel. 

70 saved_model_serialized: The serialized SavedModel proto. 

71 

72 Returns: 

73 A string containing the singleprint of the SavedModel. 

74 

75 Raises: 

76 ValueError: If a valid singleprint cannot be constructed from a SavedModel. 

77 """ 

78 try: 

79 fingerprint_serialized = fingerprinting_pywrap.CreateFingerprintDef( 

80 saved_model_serialized, export_dir) 

81 except fingerprinting_pywrap.FingerprintException as e: 

82 raise ValueError(e) from None 

83 fingerprint = fingerprinting.Fingerprint.from_proto(fingerprint_serialized) 

84 return fingerprint.singleprint() 

85 

86 

87def to_proto(fingerprint): 

88 if not isinstance(fingerprint, fingerprinting.Fingerprint): 

89 raise TypeError("Supplied value is not a Fingerprint.") 

90 return fingerprint_pb2.FingerprintDef( 

91 saved_model_checksum=fingerprint.saved_model_checksum, 

92 graph_def_program_hash=fingerprint.graph_def_program_hash, 

93 signature_def_hash=fingerprint.signature_def_hash, 

94 saved_object_graph_hash=fingerprint.saved_object_graph_hash, 

95 checkpoint_hash=fingerprint.checkpoint_hash)