Add doc

lineapy/api/models/linea_artifact.py

@@ -123,13 +123,15 @@ def _get_storage_path(self) -> Optional[str]:
     @lru_cache(maxsize=None)
     def get_metadata(self, lineapy_only: bool = False) -> ArtifactInfo:
         """
-        Get artifact metadata
+        Get artifact backend storage metadata
 
         :param lineapy_only: If ``False``, will include both LineaPy related
             metadata and metadata from storage backend(if it is not LineaPy).
             If ``True``, will only return LineaPy related metadata no matter
             which storage backend is using.
 
+        :return: Metadata for artifact backend storage.
+
         """
 
         if self._artifact_id is None or self.date_created is None:

lineapy/data/types.py

@@ -539,6 +539,10 @@ class LineaArtifactDef(TypedDict):
 
 @dataclass
 class LineaArtifactInfo:
+    """
+    Backend storage metadata for LineaPy
+    """
+
     artifact_id: int
     name: str
     version: int
@@ -552,6 +556,10 @@ class LineaArtifactInfo:
 
 @dataclass
 class MLflowArtifactInfo:
+    """
+    Backend storage metadata for MLflow
+    """
+
     id: int
     artifact_id: int
     tracking_uri: str
@@ -561,5 +569,12 @@ class MLflowArtifactInfo:
 
 
 class ArtifactInfo(TypedDict):
+    """
+    Artifact backend storage metadata
+    - lineapy : storage backend for LineaPy
+    - mlflow : storage backend metadata for MLflow (only exists when the
+        artifact is saved in MLflow)
+    """
+
     lineapy: LineaArtifactInfo
     mlflow: NotRequired[MLflowArtifactInfo]

lineapy/db/db.py

@@ -343,13 +343,17 @@ def write_artifact(self, artifact: Artifact) -> None:
     def write_mlflow_artifactmetadata(
         self, artifactorm: ArtifactORM, modelinfo: ModelInfo
     ) -> None:
+        """
+        Write MLflow metadata for the artifact
+        """
         model_flavors = [
             flavor
             for flavor in modelinfo.flavors.keys()
             if flavor != "python_function"
         ]
         if len(model_flavors) > 1:
-            raise
+            msg = "Currently, only one MLflow model flavor(other than python_function) is supported."
+            raise NotImplementedError(msg)
 
         mlflowmetadataorm = MLflowArtifactMetadataORM(
             artifact_id=artifactorm.id,
@@ -778,7 +782,10 @@ def delete_artifact_by_name(
             self.session.delete(res)
         self.renew_session()
 
-    def delete_mlflow_metadata_by_artifact_id(self, artifact_id: int):
+    def delete_mlflow_metadata_by_artifact_id(self, artifact_id: int) -> None:
+        """
+        Delete MLflow metadata for the artifact
+        """
         res_query = self.session.query(MLflowArtifactMetadataORM).filter(
             MLflowArtifactMetadataORM.artifact_id == artifact_id
         )
@@ -853,7 +860,7 @@ def get_mlflowartifactmetadataorm_by_artifact_id(
         self, artifact_id: int
     ) -> MLflowArtifactMetadataORM:
         """
-        Gets the most recent artifact with a certain name.
+        Get MLflow metadata for the artifact
         """
 
         res_query = self.session.query(MLflowArtifactMetadataORM).filter(

tests/test_mlflow.py

@@ -11,7 +11,11 @@
 
 def test_save_and_get_to_mlflow(execute):
     """
-    Test lineapy.save to mlflow, passing `registered_model_name` to test **kwargs.
+    Test mlflow as backend storage
+
+    Test lineapy.save for mlflow, passing `registered_model_name` to test
+    **kwargs. Test Artifact.get_value() for mlflow and validate model
+    exists in MLflow.
     """
     code = """
 import lineapy