Skip to content
This repository has been archived by the owner on Sep 13, 2023. It is now read-only.

Add docstrings to extensions #413

Merged
merged 12 commits into from
Oct 6, 2022
Prev Previous commit
Next Next commit
Add docstrings to extensions
mike0sv committed Sep 16, 2022
commit 7873ca437b6d59f084a074954f92691b7b1874f8
7 changes: 5 additions & 2 deletions mlem/api/commands.py
Original file line number Diff line number Diff line change
@@ -331,7 +331,9 @@ def build(
)


def serve(model: MlemModel, server: Union[Server, str], **server_kwargs):
def serve(
model: Union[str, MlemModel], server: Union[Server, str], **server_kwargs
):
"""Serve model via HTTP/HTTPS.

Args:
@@ -340,7 +342,8 @@ def serve(model: MlemModel, server: Union[Server, str], **server_kwargs):
"""
from mlem.runtime.interface import ModelInterface

model.load_value()
model = get_model_meta(model, load_value=True)

interface = ModelInterface(model_type=model.model_type)

server_obj = ensure_mlem_object(Server, server, **server_kwargs)
5 changes: 5 additions & 0 deletions mlem/contrib/bitbucketfs.py
Original file line number Diff line number Diff line change
@@ -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
5 changes: 5 additions & 0 deletions mlem/contrib/callable.py
Original file line number Diff line number Diff line change
@@ -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
5 changes: 5 additions & 0 deletions mlem/contrib/catboost.py
Original file line number Diff line number Diff line change
@@ -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
4 changes: 3 additions & 1 deletion mlem/contrib/docker/__init__.py
Original file line number Diff line number Diff line change
@@ -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
"""
5 changes: 5 additions & 0 deletions mlem/contrib/dvc.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,8 @@
"""DVC Support
Extension type: storage

Support for storing artifacts with DVC
"""
import contextlib
import os.path
import posixpath
5 changes: 5 additions & 0 deletions mlem/contrib/fastapi.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,8 @@
"""FastAPI serving
Extension type: serving

FastAPIServer implementation
"""
import logging
from collections.abc import Callable
from types import ModuleType
5 changes: 5 additions & 0 deletions mlem/contrib/github.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,8 @@
"""Github URI support
Extension type: uri

Implementation of `GithubResolver`
"""
import pathlib
import posixpath
from typing import ClassVar, Dict, Optional
5 changes: 5 additions & 0 deletions mlem/contrib/gitlabfs.py
Original file line number Diff line number Diff line change
@@ -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
5 changes: 5 additions & 0 deletions mlem/contrib/heroku/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
"""Heroku Deployments support
Extension type: deployment

Implements MlemEnv, MlemDeployment and DeployState to work with heroku.com
"""
3 changes: 3 additions & 0 deletions mlem/contrib/kubernetes/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
"""Kubernetes Deployments support
Extension type: deployment
"""
6 changes: 6 additions & 0 deletions mlem/contrib/lightgbm.py
Original file line number Diff line number Diff line change
@@ -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
5 changes: 5 additions & 0 deletions mlem/contrib/numpy.py
Original file line number Diff line number Diff line change
@@ -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

5 changes: 5 additions & 0 deletions mlem/contrib/onnx.py
Original file line number Diff line number Diff line change
@@ -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
6 changes: 6 additions & 0 deletions mlem/contrib/pandas.py
Original file line number Diff line number Diff line change
@@ -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
6 changes: 6 additions & 0 deletions mlem/contrib/pip/__init__.py
Original file line number Diff line number Diff line change
@@ -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
"""
5 changes: 5 additions & 0 deletions mlem/contrib/rabbitmq.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,8 @@
"""RabbitMQ serving
Extension type: serving

RabbitMQServer implementation
"""
import json
from time import time
from typing import Callable, ClassVar, Optional
5 changes: 5 additions & 0 deletions mlem/contrib/sagemaker/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
"""Sagemaker Deployments support
Extension type: deployment

Implements MlemEnv, MlemDeployment and DeployState to work with AWS SageMaker
"""
5 changes: 5 additions & 0 deletions mlem/contrib/sklearn.py
Original file line number Diff line number Diff line change
@@ -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
6 changes: 6 additions & 0 deletions mlem/contrib/tensorflow.py
Original file line number Diff line number Diff line change
@@ -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
7 changes: 7 additions & 0 deletions mlem/contrib/torch.py
Original file line number Diff line number Diff line change
@@ -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
6 changes: 6 additions & 0 deletions mlem/contrib/xgboost.py
Original file line number Diff line number Diff line change
@@ -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
20 changes: 20 additions & 0 deletions mlem/ext.py
Original file line number Diff line number Diff line change
@@ -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
#
7 changes: 0 additions & 7 deletions setup.py
Original file line number Diff line number Diff line change
@@ -164,13 +164,6 @@
"env.heroku = mlem.contrib.heroku.meta:HerokuEnv",
"deploy_state.heroku = mlem.contrib.heroku.meta:HerokuState",
"server._heroku = mlem.contrib.heroku.server:HerokuServer",
"deployment.kubernetes = mlem.contrib.kubernetes.base:K8sDeployment",
"deploy_state.kubernetes = mlem.contrib.kubernetes.base:K8sDeploymentState",
"env.kubernetes = mlem.contrib.kubernetes.base:K8sEnv",
"builder.kubernetes = mlem.contrib.kubernetes.base:K8sYamlBuilder",
"k8s_service_type.clusterip = mlem.contrib.kubernetes.service:ClusterIPService",
"k8s_service_type.loadbalancer = mlem.contrib.kubernetes.service:LoadBalancerService",
"k8s_service_type.nodeport = mlem.contrib.kubernetes.service:NodePortService",
"data_reader.lightgbm = mlem.contrib.lightgbm:LightGBMDataReader",
"data_type.lightgbm = mlem.contrib.lightgbm:LightGBMDataType",
"data_writer.lightgbm = mlem.contrib.lightgbm:LightGBMDataWriter",
35 changes: 34 additions & 1 deletion tests/test_ext.py
Original file line number Diff line number Diff line change
@@ -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",
}