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

1import typing 

2from pathlib import Path 

3from typing import Iterator, Mapping 

4 

5from mlopus.utils import dicts, pydantic, mongo, urls 

6from . import entity, contract 

7from .common import schema, decorators 

8from .run import RunApi 

9 

10RunIdentifier = contract.RunIdentifier 

11 

12 

13class ExpApi(schema.Experiment, entity.EntityApi): 

14 """Experiment metadata with MLflow API handle.""" 

15 

16 def _get_latest_data(self) -> schema.Experiment: 

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

18 return self.api.get_exp(self) 

19 

20 @property 

21 def url(self) -> str: 

22 """This experiment's URL.""" 

23 return self.api.get_exp_url(self) 

24 

25 @pydantic.validate_arguments 

26 def find_runs(self, query: mongo.Query | None = None, sorting: mongo.Sorting | None = None) -> Iterator[RunApi]: 

27 """Search runs belonging to this experiment with query in MongoDB query language. 

28 

29 :param query: Query in MongoDB query language. 

30 :param sorting: Sorting criteria (e.g.: `[("asc_field", 1), ("desc_field", -1)]`). 

31 """ 

32 results = self.api.find_runs(dicts.set_reserved_key(query, key="exp.id", val=self.id), sorting) 

33 return typing.cast(Iterator[RunApi], results) 

34 

35 def cache_meta(self) -> "ExpApi": 

36 """Fetch latest metadata for this experiment and save it to cache.""" 

37 return self._use_values_from(self.api.cache_exp_meta(self)) 

38 

39 def export_meta(self, target: Path) -> "ExpApi": 

40 """Export experiment metadata cache to target. 

41 

42 :param target: Cache export path. 

43 """ 

44 return self._use_values_from(self.api.export_exp_meta(self, target)) 

45 

46 @pydantic.validate_arguments 

47 def create_run( 

48 self, 

49 name: str | None = None, 

50 tags: Mapping | None = None, 

51 repo: str | urls.Url | None = None, 

52 parent: RunIdentifier | None = None, 

53 ) -> RunApi: 

54 """Declare a new run in this experiment to be used later. 

55 

56 :param name: See :attr:`schema.Run.name`. 

57 :param tags: See :attr:`schema.Run.tags`. 

58 :param repo: See :paramref:`~mlopus.mlflow.BaseMlflowApi.create_run.repo`. 

59 :param parent: Parent run ID or object. 

60 """ 

61 return typing.cast(RunApi, self.api.create_run(self, name, tags, repo, parent)) 

62 

63 @pydantic.validate_arguments 

64 def start_run( 

65 self, 

66 name: str | None = None, 

67 tags: Mapping | None = None, 

68 repo: str | urls.Url | None = None, 

69 parent: RunIdentifier | None = None, 

70 ) -> RunApi: 

71 """Start a new run in this experiment. 

72 

73 :param name: See :attr:`schema.Run.name`. 

74 :param tags: See :attr:`schema.Run.tags`. 

75 :param repo: See :paramref:`~mlopus.mlflow.BaseMlflowApi.start_run.repo`. 

76 :param parent: Parent run ID or object. 

77 """ 

78 return typing.cast(RunApi, self.api.start_run(self, name, tags, repo, parent)) 

79 

80 @decorators.require_update 

81 def set_tags(self, tags: Mapping) -> "ExpApi": 

82 """Set tags on this experiment. 

83 

84 :param tags: See :attr:`schema.Experiment.tags`. 

85 """ 

86 self.api.set_tags_on_exp(self, tags) 

87 return self