Resolve requirements installed from Git or local source (#645)

For a package installed with `pip install git+https://...` this will pin
down the URL to the package to be installed from Git.

For a package installed with `pip install local/path` this will collect
the package and add it as base64-encoded string to `.mlem` file.

PR also adds `--build_arg` option for Docker builder to pass args at
building time.

TODO:
- [ ] check that nothing breaks because of new fields in `requirements:
` (e.g. `source_url`)

For k8s deployment this works like:
```sh
$ mlem declare deployment kubernetes deployer \
  --image_name myimage --service_type loadbalancer --registry remote \
  --env docker --env.registry remote --registry.host localhost --namespace myns \
  --build_arg.0 GITHUB_TOKEN --build_arg.1 GITHUB_USERNAME=aguschin
💾 Saving deployment to deployer.mlem
```
Note that one ARG value is set, another isn't:

```yaml
$ cat deployer.mlem
build_arg:
- GITHUB_TOKEN
- GITHUB_USERNAME=aguschin
env:
  object_type: env
  registry:
    type: remote
  type: kubernetes
image_name: myimage
namespace: myns
object_type: deployment
registry:
  host: localhost
  type: remote
service_type:
  type: loadbalancer
type: kubernetes
```

Now at `mlem deploy --load deployer.mlem` MLEM will use `GITHUB_TOKEN`
env var if set. If var has a different name, this can be run like
```sh
$ GITHUB_TOKEN=$ENV_VAR_TO_USE mlem deploy --load deployer.mlem
```
Env var takes precedence, so if you set `GITHUB_USERNAME`, it's value
will be used instead of `aguschin`.

mlem/contrib/docker/base.py

@@ -24,7 +24,11 @@
 from pydantic import BaseModel
 
 from mlem.config import LOCAL_CONFIG, project_config
-from mlem.contrib.docker.context import DockerBuildArgs, DockerModelDirectory
+from mlem.contrib.docker.context import (
+    DockerBuildArgs,
+    DockerModelDirectory,
+    get_build_args,
+)
 from mlem.contrib.docker.utils import (
     build_image_with_logs,
     container_is_running,
@@ -592,6 +596,7 @@ def build_image(self, context_dir: str) -> DockerImage:
                     tag=tag,
                     rm=True,
                     platform=self.args.platform,
+                    buildargs=get_build_args(self.args.build_arg),
                 )
                 docker_image = DockerImage(**self.image.dict())
                 docker_image.image_id = image.id

mlem/contrib/docker/context.py

@@ -244,6 +244,8 @@ class Config:
     """a path to mlem .whl file. If it is empty, mlem will be installed from pip"""
     platform: Optional[str] = None
     """platform to build docker for, see docs.docker.com/desktop/multi-arch/"""
+    build_arg: List[str] = []
+    """args to use at build time https://docs.docker.com/engine/reference/commandline/build/#build-arg"""
 
     def get_base_image(self):
         if self.base_image is None:
@@ -270,6 +272,17 @@ def update(self, other: "DockerBuildArgs"):
                     setattr(self, field, value)
 
 
+def get_build_args(build_arg) -> Dict[str, Optional[str]]:
+    args = {}
+    for arg in build_arg:
+        if "=" in arg:
+            key, value = arg.split("=", 1)
+            args[key] = os.getenv(arg) or value
+        else:
+            args[arg] = os.getenv(arg)
+    return args
+
+
 class DockerModelDirectory(BaseModel):
     model: MlemModel
     server: Server
@@ -370,7 +383,9 @@ def write_dockerfile(self, requirements: Requirements):
             dockerfile = DockerfileGenerator(
                 **self.docker_args.dict()
             ).generate(
-                env=env, packages=[p.package_name for p in unix_packages or []]
+                env=env,
+                arg=get_build_args(self.docker_args.build_arg),
+                packages=[p.package_name for p in unix_packages or []],
             )
             df.write(dockerfile)
 

mlem/contrib/docker/dockerfile.j2

@@ -1,5 +1,7 @@
 FROM {{ base_image }}
 WORKDIR /app
+{% for name, value in arg.items() %}ARG {{ name }}
+{% endfor %}
 {% include "pre_install.j2" ignore missing %}
 {% include "install_req.j2" %}
 {% include "post_install.j2" ignore missing %}

mlem/contrib/docker/install_req.j2

@@ -1,3 +1,5 @@
+# install Git in case something in requirements.txt will be installed from Git repo
+RUN {{ package_install_cmd }} git {{ package_clean_cmd }}
 {% if packages %}RUN {{ package_install_cmd }} {{ packages|join(" ")  }} {{ package_clean_cmd }}{% endif %}
 COPY requirements.txt .
 RUN pip install -r requirements.txt && pip cache purge

mlem/contrib/kubernetes/base.py

@@ -78,6 +78,8 @@ class K8sDeployment(
     """Path for kube config file of the cluster"""
     templates_dir: List[str] = []
     """List of dirs where templates reside"""
+    build_arg: List[str] = []
+    """Args to use at build time https://docs.docker.com/engine/reference/commandline/build/#build-arg"""
 
     def load_kube_config(self):
         config.load_kube_config(
@@ -134,6 +136,7 @@ def deploy(self, model: MlemModel):
                     registry=self.get_registry(),
                     daemon=self.daemon,
                     server=self.get_server(),
+                    build_arg=self.build_arg,
                 )
                 state.update_model(model)
                 redeploy = True

mlem/contrib/kubernetes/build.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import List, Optional
 
 from mlem.core.objects import MlemModel
 from mlem.runtime.server import Server
@@ -16,6 +16,7 @@ def build_k8s_docker(
     server: Server,
     platform: Optional[str] = "linux/amd64",
     # runners usually do not support arm64 images built on Mac M1 devices
+    build_arg: Optional[List[str]] = None,
 ):
     echo(EMOJI_BUILD + f"Creating docker image {image_name}")
     with set_offset(2):
@@ -28,4 +29,5 @@ def build_k8s_docker(
             tag=meta.meta_hash(),
             force_overwrite=True,
             platform=platform,
+            build_arg=build_arg or [],
         )

mlem/core/requirements.py

@@ -25,6 +25,7 @@
     Union,
 )
 
+from importlib_metadata import Distribution, PackageNotFoundError
 from pydantic import BaseModel, root_validator
 
 from mlem.core.base import MlemABC
@@ -86,6 +87,12 @@ class InstallableRequirement(PythonRequirement):
     """Pip package name for this module, if it is different from module name"""
     extra_index: Optional[str] = None
     """Extra index to use for this package"""
+    source_url: Optional[str] = None
+    """URL to download package from"""
+    vcs: Optional[str] = None
+    """Version control system"""
+    vcs_commit: Optional[str] = None
+    """Commit hash to checkout from VCS"""
 
     @root_validator
     def set_package_name(cls, values):  # pylint: disable=no-self-argument
@@ -112,6 +119,8 @@ def get_repr(self):
         """
         if self.version is not None:
             return f"{self.package}=={self.version}"
+        if self.vcs_commit is not None:
+            return f"{self.vcs}+{self.source_url}@{self.vcs_commit}"
         return self.package
 
     @classmethod
@@ -134,10 +143,37 @@ def from_module(
         """
         from mlem.utils.module import get_module_version
 
+        # similar package resolution exists in mlem/contrib/docker/context.py
+        source_url = None
+        vcs = None
+        rev = None
+
+        try:
+            pkg_info = json.loads(
+                Distribution.from_name(mod.__name__).read_text(
+                    "direct_url.json"
+                )
+            )
+        except (PackageNotFoundError, json.JSONDecodeError, TypeError):
+            pass
+        else:
+            if "url" in pkg_info:
+                source_url = pkg_info["url"]
+                # installed from local source:
+                if source_url.startswith("file://"):
+                    return CustomRequirement.from_module(mod)  # type: ignore
+                # installed from git branch (probably):
+                if "vcs_info" in pkg_info:
+                    rev = pkg_info["vcs_info"]["commit_id"]
+                    vcs = pkg_info["vcs_info"]["vcs"]
+
         return InstallableRequirement(
             module=mod.__name__,
-            version=get_module_version(mod),
+            version=None if rev else get_module_version(mod),
             package_name=package_name,
+            source_url=source_url,
+            vcs_commit=rev,
+            vcs=vcs,
         )
 
     @classmethod

mlem/utils/module.py

@@ -331,7 +331,7 @@ def get_module_as_requirement(
     if validate_pypi:
         mod_name = get_package_name(mod)
         check_pypi_module(mod_name, mod_version, raise_on_error=True)
-    return InstallableRequirement(module=mod.__name__, version=mod_version)
+    return InstallableRequirement.from_module(mod)
 
 
 def get_local_module_reqs(mod) -> List[ModuleType]:

tests/contrib/test_docker/test_context.py

@@ -21,6 +21,8 @@ def test_dockerfile_generator_custom_python_version():
     dockerfile = _cut_empty_lines(
         f"""FROM python:AAAA-slim
 WORKDIR /app
+# install Git in case something in requirements.txt will be installed from Git repo
+RUN apt-get update && apt-get -y upgrade && apt-get install --no-install-recommends -y git && apt-get clean && rm -rf /var/lib/apt/lists/*
 COPY requirements.txt .
 RUN pip install -r requirements.txt && pip cache purge
 COPY mlem_requirements.txt .
@@ -40,6 +42,8 @@ def test_dockerfile_generator_unix_packages():
     dockerfile = _cut_empty_lines(
         f"""FROM python:3.6-slim
 WORKDIR /app
+# install Git in case something in requirements.txt will be installed from Git repo
+RUN kek git lol
 RUN kek aaa bbb lol
 COPY requirements.txt .
 RUN pip install -r requirements.txt && pip cache purge
@@ -72,6 +76,8 @@ def test_dockerfile_generator_super_custom():
         f"""FROM my-python:3.6
 WORKDIR /app
 RUN echo "pre_install"
+# install Git in case something in requirements.txt will be installed from Git repo
+RUN apt-get update && apt-get -y upgrade && apt-get install --no-install-recommends -y git && apt-get clean && rm -rf /var/lib/apt/lists/*
 COPY requirements.txt .
 RUN pip install -r requirements.txt && pip cache purge
 RUN pip install mlem=={mlem.__version__} && pip cache purge
@@ -109,14 +115,18 @@ def test_use_wheel_installation(tmpdir):
     distr.write("wheel goes brrr")
     with use_mlem_source("whl"):
         os.environ["MLEM_DOCKER_WHEEL_PATH"] = str(os.path.basename(distr))
-        dockerfile = DockerfileGenerator().generate(env={}, packages=[])
+        dockerfile = DockerfileGenerator().generate(
+            env={}, packages=[], arg={}
+        )
         assert f"RUN pip install {MLEM_LOCAL_WHL}" in dockerfile
 
 
 def _generate_dockerfile(unix_packages=None, **kwargs):
     return _cut_empty_lines(
         DockerfileGenerator(**kwargs).generate(
-            env={}, packages=[p.package_name for p in unix_packages or []]
+            env={},
+            packages=[p.package_name for p in unix_packages or []],
+            arg={},
         )
     )