Add missing docstrings and refactor a little (#261)

* Add missing docstrings and refactor a little
closes #70

* fix lint and
closes #257

* fix batch and dataset types

* fix test

mlem/api/commands.py

@@ -4,7 +4,6 @@
 import posixpath
 from typing import Any, Dict, Iterable, List, Optional, Type, Union
 
-import numpy as np
 from fsspec import AbstractFileSystem
 from fsspec.implementations.local import LocalFileSystem
 
@@ -34,10 +33,10 @@
     MlemLink,
     MlemModel,
     MlemObject,
+    MlemPackager,
 )
-from mlem.pack import Packager
-from mlem.runtime.client.base import BaseClient
-from mlem.runtime.server.base import Server
+from mlem.runtime.client import Client
+from mlem.runtime.server import Server
 from mlem.ui import (
     EMOJI_APPLY,
     EMOJI_COPY,
@@ -91,8 +90,7 @@ def apply(
             batch_dataset = get_dataset_value(part, batch_size)
             for chunk in batch_dataset:
                 preds = w.call_method(resolved_method, chunk.data)
-                res += [*preds]
-        res = [np.array(res)]
+                res += [*preds]  # TODO: merge results
     else:
         res = [
             w.call_method(resolved_method, get_dataset_value(part))
@@ -103,17 +101,12 @@ def apply(
             return res[0]
         return res
     if len(res) == 1:
-        return save(
-            res[0], output, repo=target_repo, external=external, index=index
-        )
-
-    raise NotImplementedError(
-        "Saving several input data objects is not implemented yet"
-    )
+        res = res[0]
+    return save(res, output, repo=target_repo, external=external, index=index)
 
 
 def apply_remote(
-    client: Union[str, BaseClient],
+    client: Union[str, Client],
     *data: Union[str, MlemDataset, Any],
     method: str = None,
     output: str = None,
@@ -124,7 +117,7 @@ def apply_remote(
     """Apply provided model against provided data
 
     Args:
-        client (BaseClient): The client to access methods of deployed model.
+        client (Client): The client to access methods of deployed model.
         data (Any): Input to the model.
         method (str, optional): Which model method to use.
             If None, use the only method model has.
@@ -138,7 +131,7 @@ def apply_remote(
             Otherwise returns None.
 
     """
-    client = ensure_mlem_object(BaseClient, client, **client_kwargs)
+    client = ensure_mlem_object(Client, client, **client_kwargs)
     if method is not None:
         try:
             resolved_method = getattr(client, method)
@@ -154,11 +147,8 @@ def apply_remote(
             return res[0]
         return res
     if len(res) == 1:
-        return save(res[0], output, repo=target_repo, index=index)
-
-    raise NotImplementedError(
-        "Saving several input data objects is not implemented yet"
-    )
+        res = res[0]
+    return save(res, output, repo=target_repo, index=index)
 
 
 def clone(
@@ -317,21 +307,21 @@ def link(
 
 
 def pack(
-    packager: Union[str, Packager],
+    packager: Union[str, MlemPackager],
     model: Union[str, MlemModel],
     **packager_kwargs,
 ):
     """Pack model in docker-build-ready folder or directly build a docker image.
 
     Args:
-        packager (Union[str, Packager]): Packager to use.
+        packager (Union[str, MlemPackager]): Packager to use.
             Out-of-the-box supported string values are "docker_dir" and "docker".
         model (Union[str, MlemModel]): The model to pack.
     """
     model = get_model_meta(model)
-    return ensure_mlem_object(Packager, packager, **packager_kwargs).package(
-        model
-    )
+    return ensure_mlem_object(
+        MlemPackager, packager, **packager_kwargs
+    ).package(model)
 
 
 def serve(model: MlemModel, server: Union[Server, str], **server_kwargs):
@@ -341,7 +331,7 @@ def serve(model: MlemModel, server: Union[Server, str], **server_kwargs):
         model (MlemModel): The model to serve.
         server (Union[Server, str]): Out-of-the-box supported one is "fastapi".
     """
-    from mlem.runtime.interface.base import ModelInterface
+    from mlem.runtime.interface import ModelInterface
 
     model.load_value()
     interface = ModelInterface(model_type=model.model_type)

mlem/api/utils.py

@@ -9,7 +9,7 @@
 
 def get_dataset_value(dataset: Any, batch_size: Optional[int] = None) -> Any:
     if isinstance(dataset, str):
-        return load(dataset)
+        return load(dataset, batch_size=batch_size)
     if isinstance(dataset, MlemDataset):
         # TODO: https://github.com/iterative/mlem/issues/29
         #  fix discrepancies between model and data meta objects

mlem/cli/apply.py

@@ -26,7 +26,7 @@
 from mlem.core.metadata import load_meta
 from mlem.core.objects import MlemDataset, MlemModel
 from mlem.ext import list_implementations
-from mlem.runtime.client.base import BaseClient
+from mlem.runtime.client import Client
 from mlem.ui import set_echo
 
 
@@ -123,7 +123,7 @@ def apply(
 def apply_remote(
     subtype: str = Argument(
         "",
-        help=f"Type of client. Choices: {list_implementations(BaseClient)}",
+        help=f"Type of client. Choices: {list_implementations(Client)}",
         show_default=False,
     ),
     data: str = Argument(..., help="Path to dataset object"),
@@ -146,7 +146,7 @@ def apply_remote(
         Apply hosted mlem model to local mlem dataset
         $ mlem apply-remote http mydataset -c host="0.0.0.0" -c port=8080 --output myprediction
     """
-    client = config_arg(BaseClient, load, subtype, conf, file_conf)
+    client = config_arg(Client, load, subtype, conf, file_conf)
 
     with set_echo(None if json else ...):
         result = run_apply_remote(
@@ -163,7 +163,7 @@ def apply_remote(
 
 
 def run_apply_remote(
-    client: BaseClient,
+    client: Client,
     data: str,
     repo,
     rev,

mlem/cli/package.py

@@ -12,17 +12,16 @@
     option_rev,
 )
 from mlem.core.metadata import load_meta
-from mlem.core.objects import MlemModel
+from mlem.core.objects import MlemModel, MlemPackager
 from mlem.ext import list_implementations
-from mlem.pack import Packager
 
 
 @mlem_command("pack", section="runtime")
 def pack(
     model: str = Argument(..., help="Path to model"),
     subtype: str = Argument(
         "",
-        help=f"Type of packing. Choices: {list_implementations(Packager)}",
+        help=f"Type of packing. Choices: {list_implementations(MlemPackager)}",
         show_default=False,
     ),
     repo: Optional[str] = option_repo,
@@ -45,6 +44,6 @@ def pack(
     from mlem.api.commands import pack
 
     pack(
-        config_arg(Packager, load, subtype, conf, file_conf),
+        config_arg(MlemPackager, load, subtype, conf, file_conf),
         load_meta(model, repo, rev, force_type=MlemModel),
     )

mlem/cli/serve.py

@@ -14,7 +14,7 @@
 from mlem.core.metadata import load_meta
 from mlem.core.objects import MlemModel
 from mlem.ext import list_implementations
-from mlem.runtime.server.base import Server
+from mlem.runtime.server import Server
 
 
 @mlem_command("serve", section="runtime")

mlem/config.py

@@ -52,6 +52,8 @@ def inner(settings: BaseSettings) -> Dict[str, Any]:
 
 
 class MlemConfigBase(BaseSettings):
+    """Special base for mlem settings to be able to read them from files"""
+
     config_path: str = ""
     config_fs: Optional[AbstractFileSystem] = None
 
@@ -78,6 +80,8 @@ def customise_sources(
 
 
 class MlemConfig(MlemConfigBase):
+    """Base Mlem Config"""
+
     GITHUB_USERNAME: Optional[str] = Field(default=None, env="GITHUB_USERNAME")
     GITHUB_TOKEN: Optional[str] = Field(default=None, env="GITHUB_TOKEN")
     ADDITIONAL_EXTENSIONS_RAW: str = Field(

mlem/contrib/docker/base.py

@@ -20,9 +20,8 @@
 )
 from mlem.core.base import MlemABC
 from mlem.core.errors import DeploymentError
-from mlem.core.objects import MlemModel
-from mlem.pack import Packager
-from mlem.runtime.server.base import Server
+from mlem.core.objects import MlemModel, MlemPackager
+from mlem.runtime.server import Server
 from mlem.ui import EMOJI_BUILD, EMOJI_OK, EMOJI_UPLOAD, echo
 
 logger = logging.getLogger(__name__)
@@ -302,7 +301,7 @@ class _DockerPackMixin(BaseModel):
     args: DockerBuildArgs = DockerBuildArgs()
 
 
-class DockerDirPackager(Packager, _DockerPackMixin):
+class DockerDirPackager(MlemPackager, _DockerPackMixin):
     type: ClassVar[str] = "docker_dir"
     target: str
 
@@ -318,7 +317,7 @@ def package(self, obj: MlemModel):
         return docker_dir
 
 
-class DockerImagePackager(Packager, _DockerPackMixin):
+class DockerImagePackager(MlemPackager, _DockerPackMixin):
     type: ClassVar[str] = "docker"
     image: DockerImage
     env: DockerEnv = DockerEnv()

mlem/contrib/docker/context.py

@@ -15,7 +15,7 @@
 import mlem
 from mlem.core.objects import MlemModel
 from mlem.core.requirements import Requirements, UnixPackageRequirement
-from mlem.runtime.server.base import Server
+from mlem.runtime.server import Server
 from mlem.ui import EMOJI_BUILD, EMOJI_PACK, echo, no_echo
 from mlem.utils.module import get_python_version
 from mlem.utils.templates import TemplateModel

mlem/contrib/docker/helpers.py

@@ -1,5 +1,5 @@
 from mlem.core.objects import MlemModel
-from mlem.runtime.server.base import Server
+from mlem.runtime.server import Server
 
 from . import DockerImagePackager
 from .base import DockerBuildArgs, DockerEnv, DockerImage

mlem/contrib/fastapi.py

@@ -11,8 +11,8 @@
 
 from mlem.core.model import Signature
 from mlem.core.requirements import LibRequirementsMixin
-from mlem.runtime.interface.base import Interface
-from mlem.runtime.server.base import Server
+from mlem.runtime.interface import Interface
+from mlem.runtime.server import Server
 from mlem.ui import EMOJI_NAILS, echo
 
 logger = logging.getLogger(__name__)

mlem/contrib/heroku/meta.py

@@ -4,7 +4,7 @@
 from pydantic import BaseModel
 
 from mlem.core.objects import DeployState, DeployStatus, MlemDeploy, MlemEnv
-from mlem.runtime.client.base import BaseClient, HTTPClient
+from mlem.runtime.client import Client, HTTPClient
 
 from ...core.errors import DeploymentError
 from ...ui import EMOJI_OK, echo
@@ -39,7 +39,7 @@ def ensured_app(self) -> HerokuAppMeta:
             raise ValueError("App is not created yet")
         return self.app
 
-    def get_client(self) -> BaseClient:
+    def get_client(self) -> Client:
         return HTTPClient(
             host=urlparse(self.ensured_app.web_url).netloc, port=80
         )

mlem/contrib/pip/base.py

@@ -11,9 +11,8 @@
 
 import mlem
 from mlem.core.meta_io import get_fs, get_uri
-from mlem.core.objects import MlemModel
+from mlem.core.objects import MlemModel, MlemPackager
 from mlem.core.requirements import InstallableRequirement
-from mlem.pack import Packager
 from mlem.ui import EMOJI_PACK, echo, no_echo
 from mlem.utils.module import get_python_version
 from mlem.utils.templates import TemplateModel
@@ -77,7 +76,7 @@ def make_distr(self, obj: MlemModel, root: str, fs: AbstractFileSystem):
         )
 
 
-class PipPackager(Packager, PipMixin):
+class PipPackager(MlemPackager, PipMixin):
     type: ClassVar = "pip"
     target: str
 
@@ -86,7 +85,7 @@ def package(self, obj: MlemModel):
         self.make_distr(obj, root, fs)
 
 
-class WhlPackager(Packager, PipMixin):
+class WhlPackager(MlemPackager, PipMixin):
     type: ClassVar = "whl"
     target: str
 

mlem/contrib/rabbitmq.py

@@ -10,9 +10,9 @@
 from mlem.core.errors import MlemError
 from mlem.core.model import Signature
 from mlem.runtime import Interface
-from mlem.runtime.client.base import BaseClient
-from mlem.runtime.interface.base import InterfaceDescriptor
-from mlem.runtime.server.base import Server
+from mlem.runtime.client import Client
+from mlem.runtime.interface import InterfaceDescriptor
+from mlem.runtime.server import Server
 from mlem.ui import EMOJI_NAILS, echo
 
 REQUEST = "_request"
@@ -95,7 +95,7 @@ def serve(self, interface: Interface):
         self.channel.start_consuming()
 
 
-class RabbitMQClient(BaseClient, RabbitMQMixin):
+class RabbitMQClient(Client, RabbitMQMixin):
     type: ClassVar = "rmq"
     timeout: float = 0