Coverage for src/mlopus/mlflow/api/common/schema.py: 100%

1from datetime import datetime 

2from enum import Enum 

3from typing import Dict, Any 

4 

5from mlopus.utils import pydantic, urls 

6from . import patterns 

7 

8 

9class RunStatus(Enum): 

10 """Run status values.""" 

11 

12 FAILED = "FAILED" 

13 RUNNING = "RUNNING" 

14 FINISHED = "FINISHED" 

15 SCHEDULED = "SCHEDULED" 

16 

17 

18class BaseEntity(pydantic.BaseModel, pydantic.MappingMixin): 

19 """Base class for entity schemas.""" 

20 

21 tags: Dict[str, Any] = pydantic.Field( 

22 repr=False, 

23 description=( 

24 "Nested tags dict with JSON-serializable leaf-values and keys " 

25 f"in the pattern: `{patterns.TAG_PARAM_OR_METRIC_KEY.pattern}`" 

26 ), 

27 ) 

28 

29 

30class Experiment(BaseEntity): 

31 """Type of Experiment used by MLOpus in generic MLflow-compliant APIs.""" 

32 

33 id: str = pydantic.Field(description=f"Experiment ID in the pattern `{patterns.EXP_ID.pattern}`") 

34 name: str = pydantic.Field(description="Experiment Name.") 

35 

36 

37class Run(BaseEntity): 

38 """Type of Run used by MLOpus in generic MLflow-compliant APIs.""" 

39 

40 id: str = pydantic.Field(description=f"Run ID in the pattern `{patterns.RUN_ID.pattern}`") 

41 name: str = pydantic.Field(description="Run name.") 

42 exp: Experiment = pydantic.Field(repr=False, description="Parent experiment.") 

43 repo: str = pydantic.Field(repr=False, description="Artifacts repo URL.") 

44 params: Dict[str, Any] = pydantic.Field( 

45 repr=False, 

46 description="Nested params dict with JSON-serializable leaf-values.", 

47 ) 

48 metrics: Dict[str, Any] = pydantic.Field( 

49 repr=False, 

50 description="Nested metrics dict with float leaf-values.", 

51 ) 

52 status: RunStatus | None = pydantic.Field(repr=False, description="Run status.") 

53 end_time: datetime | None = pydantic.Field( 

54 repr=False, 

55 description="Start time in local timezone.", 

56 ) 

57 start_time: datetime | None = pydantic.Field( 

58 repr=False, 

59 description="End time in local timezone.", 

60 ) 

61 

62 @property 

63 def repo_url(self) -> urls.Url: 

64 """Artifacts repo URL.""" 

65 return urls.parse_url(self.repo) 

66 

67 

68class Model(BaseEntity): 

69 """Type of registered Model used by MLOpus in generic MLflow-compliant APIs.""" 

70 

71 name: str = pydantic.Field(description=f"Model vame in the pattern `{patterns.MODEL_NAME.pattern}`") 

72 

73 

74class ModelVersion(BaseEntity): 

75 """Type of ModelVersion used by MLOpus in generic MLflow-compliant APIs.""" 

76 

77 version: str = pydantic.Field(description=f"Model version in the pattern `{patterns.MODEL_VERSION.pattern}`") 

78 model: Model = pydantic.Field(description="Parent model.") 

79 run: Run = pydantic.Field(repr=False, description="Parent run.") 

80 path_in_run: str = pydantic.Field(repr=False, description="Path inside run artifacts.")