Add docstrings to extensions

mlem/cli/main.py

@@ -47,6 +47,9 @@
 PATH_METAVAR = "path"
 COMMITISH_METAVAR = "commitish"
 
+PATH_METAVAR = "path"
+COMMITISH_METAVAR = "commitish"
+
 
 class MlemFormatter(HelpFormatter):
     def write_heading(self, heading: str) -> None:

mlem/cli/serve.py

@@ -26,8 +26,7 @@
 
 serve = Typer(
     name="serve",
-    help="""Create an API from model methods using a server implementation.
-    """,
+    help="""Create an API from model methods using a server implementation.""",
     cls=mlem_group("runtime"),
     subcommand_metavar="server",
 )

mlem/contrib/bitbucketfs.py

@@ -1,3 +1,8 @@
+"""BitBucket URI support
+Extension type: uri
+
+Implementation of `BitbucketFileSystem` and `BitbucketResolver`
+"""
 import posixpath
 from typing import ClassVar, List, Optional
 from urllib.parse import quote_plus, urljoin, urlparse, urlsplit

mlem/contrib/callable.py

@@ -1,3 +1,8 @@
+"""MLEM Models from arbitraty callables
+Extension type: model
+
+ModelType implementation to turn any python callable into MLEM Model
+"""
 import posixpath
 from collections import defaultdict
 from importlib import import_module

mlem/contrib/catboost.py

@@ -1,3 +1,8 @@
+"""Catboost Models Support
+Extension type: model
+
+Implementations of ModelType and ModelIO for `CatBoostClassifier` and `CatBoostRegressor`
+"""
 import os
 import posixpath
 import tempfile

mlem/contrib/docker/__init__.py

@@ -1,4 +1,6 @@
-"""
+"""Docker builds support
+Extension type: deployment
+
 Building docker images from the model
 or packing all necessary things to do that in a folder
 """

mlem/contrib/dvc.py

@@ -1,3 +1,8 @@
+"""DVC Support
+Extension type: storage
+
+Support for storing artifacts with DVC
+"""
 import contextlib
 import os.path
 import posixpath

mlem/contrib/fastapi.py

@@ -1,3 +1,8 @@
+"""FastAPI serving
+Extension type: serving
+
+FastAPIServer implementation
+"""
 import logging
 from collections.abc import Callable
 from types import ModuleType

mlem/contrib/github.py

@@ -1,3 +1,8 @@
+"""Github URI support
+Extension type: uri
+
+Implementation of `GithubResolver`
+"""
 import pathlib
 import posixpath
 import re

mlem/contrib/gitlabfs.py

@@ -1,3 +1,8 @@
+"""Gitlab URI support
+Extension type: uri
+
+Implementation of `GitlabFileSystem` and `GitlabResolver`
+"""
 import posixpath
 from typing import ClassVar, Optional
 from urllib.parse import quote_plus, urlparse, urlsplit

mlem/contrib/heroku/__init__.py

@@ -0,0 +1,5 @@
+"""Heroku Deployments support
+Extension type: deployment
+
+Implements MlemEnv, MlemDeployment and DeployState to work with heroku.com
+"""

mlem/contrib/kubernetes/__init__.py

@@ -0,0 +1,3 @@
+"""Kubernetes Deployments support
+Extension type: deployment
+"""

mlem/contrib/lightgbm.py

@@ -1,3 +1,9 @@
+"""LightGBM models support
+Extension type: model
+
+ModelType and ModelIO implementations for `lightgbm.Booster` as well as
+LightGBMDataType with Reader and Writer for `lightgbm.Dataset`
+"""
 import os
 import posixpath
 import tempfile

mlem/contrib/numpy.py

@@ -1,3 +1,8 @@
+"""Numpy data types support
+Extension type: data
+
+DataType, Reader and Writer implementations for `np.ndarray` and `np.number` primitives
+"""
 from types import ModuleType
 from typing import Any, ClassVar, Iterator, List, Optional, Tuple, Type, Union
 

mlem/contrib/onnx.py

@@ -1,3 +1,8 @@
+"""ONNX models support
+Extension type: model
+
+ModelType and ModelIO implementations for `onnx.ModelProto`
+"""
 from typing import Any, ClassVar, List, Optional, Union
 
 import numpy as np

mlem/contrib/pandas.py

@@ -1,3 +1,9 @@
+"""Pandas data types support
+Extension type: data
+
+DataType, Reader and Writer implementations for `pd.DataFrame` and `pd.Series`
+ImportHook implementation for files saved with pandas
+"""
 import os.path
 import posixpath
 import re

mlem/contrib/pip/__init__.py

@@ -0,0 +1,6 @@
+"""Python Package builds support
+Extension type: build
+
+Contains two Builder implementations: `pip` to create a directory with
+Python Package from model and `whl` to create a wheel file with Python Package
+"""

mlem/contrib/rabbitmq.py

@@ -1,3 +1,8 @@
+"""RabbitMQ serving
+Extension type: serving
+
+RabbitMQServer implementation
+"""
 import json
 from time import time
 from typing import Callable, ClassVar, Optional

mlem/contrib/sagemaker/__init__.py

@@ -0,0 +1,5 @@
+"""Sagemaker Deployments support
+Extension type: deployment
+
+Implements MlemEnv, MlemDeployment and DeployState to work with AWS SageMaker
+"""

mlem/contrib/sklearn.py

@@ -1,3 +1,8 @@
+"""Scikit-Learn models support
+Extension type: model
+
+ModelType implementations for any sklearn-compatible classes as well as `Pipeline`
+"""
 from typing import Any, ClassVar, List, Optional, Union
 
 import sklearn

mlem/contrib/tensorflow.py

@@ -1,3 +1,9 @@
+"""Tensorflow models support
+Extension type: model
+
+ModelType and ModelIO implementations for `tf.keras.Model`
+DataType, Reader and Writer implementations for `tf.Tensor`
+"""
 import posixpath
 import tempfile
 from typing import Any, ClassVar, Iterator, List, Optional, Tuple

mlem/contrib/torch.py

@@ -1,3 +1,10 @@
+"""Torch models support
+Extension type: model
+
+ModelType and ModelIO implementations for `torch.nn.Module`
+ImportHook for importing files saved with `torch.save`
+DataType, Reader and Writer implementations for `torch.Tensor`
+"""
 from typing import Any, ClassVar, Iterator, List, Optional, Tuple
 
 import torch

mlem/contrib/xgboost.py

@@ -1,3 +1,9 @@
+"""XGBoost models support
+Extension type: model
+
+ModelType and ModelIO implementations for `xgboost.Booster` as well as
+DataType, Reader and Writer implementations for `xgboost.DMatrix`
+"""
 import os
 import posixpath
 import tempfile

mlem/ext.py

@@ -4,6 +4,7 @@
 """
 import importlib
 import logging
+import re
 import sys
 from types import ModuleType
 from typing import Callable, Dict, List, Optional, Union
@@ -109,6 +110,12 @@ class ExtensionLoader:
         Extension("mlem.contrib.gitlabfs", [], True),
         Extension("mlem.contrib.bitbucketfs", [], True),
         Extension("mlem.contrib.sagemaker", ["sagemaker", "boto3"], False),
+        Extension("mlem.contrib.dvc", ["dvc"], False),
+        Extension(
+            "mlem.contrib.heroku", ["fastapi", "uvicorn", "docker"], False
+        ),
+        Extension("mlem.contrib.pip", [], False),
+        Extension("mlem.contrib.kubernetes", ["kubernetes", "docker"], False),
     )
 
     _loaded_extensions: Dict[Extension, ModuleType] = {}
@@ -258,6 +265,19 @@ def load_extensions(*exts: str):
         ExtensionLoader.load(ext)
 
 
+def get_ext_type(ext: Union[str, Extension]):
+    if isinstance(ext, Extension):
+        ext_module = ext.module
+    else:
+        ext_module = ext
+
+    doc = import_module(ext_module).__doc__ or ""
+    search = re.search(r"Extension type: (\w*)", doc)
+    if search is None:
+        raise ValueError(f"{ext_module} extension doesnt define it's type")
+    return search.group(1)
+
+
 # Copyright 2019 Zyfra
 # Copyright 2021 Iterative
 #

tests/test_ext.py

@@ -1,8 +1,12 @@
+import os
 import re
+from importlib import import_module
 from pathlib import Path
 
-from mlem import ExtensionLoader
+import pytest
+
 from mlem.config import MlemConfig, MlemConfigBase
+from mlem.ext import ExtensionLoader, get_ext_type
 from mlem.utils.entrypoints import (
     MLEM_CONFIG_ENTRY_POINT,
     MLEM_ENTRY_POINT,
@@ -85,3 +89,32 @@ def test_all_ext_has_pip_extra():
         assert name in extras
         ext_extras = extras[name]
         assert set(reqs) == {re.split("[~=]", r)[0] for r in ext_extras}
+
+
+def test_all_ext_registered():
+    from mlem import contrib
+
+    files = os.listdir(os.path.dirname(contrib.__file__))
+    ext_sources = {
+        name[: -len(".py")] if name.endswith(".py") else name
+        for name in files
+        if not name.startswith("__")
+    }
+    assert set(ExtensionLoader.builtin_extensions) == {
+        f"mlem.contrib.{name}" for name in ext_sources
+    }
+
+
+@pytest.mark.parametrize("mod", ExtensionLoader.builtin_extensions.keys())
+def test_all_ext_docstring(mod):
+    module = import_module(mod)
+    assert module.__doc__ is not None
+    assert get_ext_type(mod) in {
+        "model",
+        "deployment",
+        "data",
+        "serving",
+        "build",
+        "uri",
+        "storage",
+    }