Coverage for src/mlopus/mlflow/api/mv.py: 88%

1from pathlib import Path 

2from typing import Callable, TypeVar, Mapping 

3 

4from mlopus.utils import pydantic 

5from . import entity, contract 

6from .common import schema, decorators, transfer 

7 

8A = TypeVar("A") # Any type 

9 

10 

11class ModelVersionApi(schema.ModelVersion, entity.EntityApi): 

12 """Model version metadata with MLflow API handle.""" 

13 

14 def __init__(self, **kwargs): 

15 super().__init__(**kwargs) 

16 from .run import RunApi 

17 from .model import ModelApi 

18 

19 self.run: RunApi = RunApi(**self.run) 

20 self.model: ModelApi = ModelApi(**self.model) 

21 

22 def using(self, api: contract.MlflowApiContract) -> "ModelVersionApi": 

23 super().using(api) 

24 self.run.using(api) 

25 self.model.using(api) 

26 return self 

27 

28 def _get_latest_data(self) -> schema.ModelVersion: 

29 """Get latest data for this entity. Used for self update after methods with the `require_update` decorator.""" 

30 return self.api.get_model_version(self) 

31 

32 @property 

33 def url(self) -> str: 

34 """Get model version URL.""" 

35 return self.api.get_model_version_url(self) 

36 

37 @pydantic.validate_arguments 

38 def clean_cached_artifact(self) -> "ModelVersionApi": 

39 """Clean cached artifact for this model version.""" 

40 self.api.clean_cached_model_artifact(self) 

41 return self 

42 

43 def cache(self): 

44 """Cache metadata and artifact for this model version.""" 

45 self.cache_meta() 

46 self.cache_artifact() 

47 

48 def export(self, target: Path): 

49 """Export metadata and artifact cache of this model version to target path. 

50 

51 :param target: Cache export path. 

52 """ 

53 self.export_meta(target) 

54 self.export_artifact(target) 

55 

56 def cache_meta(self) -> "ModelVersionApi": 

57 """Fetch latest metadata for this model version and save it to cache.""" 

58 return self._use_values_from(self.api.cache_model_version_meta(self)) 

59 

60 def export_meta(self, target: Path) -> "ModelVersionApi": 

61 """Export model version metadata cache to target. 

62 

63 :param target: Cache export path. 

64 """ 

65 return self._use_values_from(self.api.export_model_version_meta(self, target)) 

66 

67 @decorators.require_update 

68 def set_tags(self, tags: Mapping) -> "ModelVersionApi": 

69 """Set tags on this model version. 

70 

71 :param tags: See :attr:`schema.Model.tags`. 

72 """ 

73 self.api.set_tags_on_model_version(self, tags) 

74 return self 

75 

76 def cache_artifact(self) -> Path: 

77 """Pull artifact of this model version from MLflow server to local cache.""" 

78 return self.api.cache_model_artifact(self) 

79 

80 @pydantic.validate_arguments 

81 def export_artifact(self, target: Path) -> Path: 

82 """Export model version artifact cache to target. 

83 

84 See also: 

85 - :meth:`mlopus.mlflow.BaseMlflowApi.export_model_artifact` 

86 

87 :param target: Cache export path. 

88 """ 

89 return self.api.export_model_artifact(self, target) 

90 

91 @pydantic.validate_arguments 

92 def list_artifacts(self, path_suffix: str = "") -> transfer.LsResult: 

93 """List artifacts in this model version. 

94 

95 :param path_suffix: Plain relative path inside model artifact dir (e.g.: `a/b/c`). 

96 """ 

97 return self.api.list_model_artifact(self, path_suffix) 

98 

99 def get_artifact(self) -> Path: 

100 """Get local path to model artifact. 

101 

102 See also: 

103 - :meth:`mlopus.mlflow.BaseMlflowApi.get_model_artifact` 

104 """ 

105 return self.api.get_model_artifact(self) 

106 

107 @pydantic.validate_arguments 

108 def place_artifact(self, target: Path, overwrite: bool = False, link: bool = True) -> "ModelVersionApi": 

109 """Place model version artifact on target path. 

110 

111 See also: 

112 - :meth:`mlopus.mlflow.BaseMlflowApi.place_model_artifact` 

113 

114 :param target: Target path. 

115 :param overwrite: Overwrite target path if exists. 

116 :param link: Use symbolic link instead of copy. 

117 """ 

118 self.api.place_model_artifact(self, target, overwrite, link) 

119 return self 

120 

121 @pydantic.validate_arguments 

122 def load_artifact(self, loader: Callable[[Path], A]) -> A: 

123 """Load model version artifact. 

124 

125 See also: 

126 - :meth:`mlopus.mlflow.BaseMlflowApi.load_model_artifact` 

127 

128 :param loader: Loader callback. 

129 """ 

130 return self.api.load_model_artifact(self, loader)