Allow to set ENV vars for build docker and deploy kubernetes (#648)

this is for
#647 (comment)

This work similarly to #645:

`mlem build docker` works the same, but I'll provide an example for k8s
deploy here:

```
$ mlem declare deployment kubernetes deployer \
  --image_name myimage --service_type loadbalancer --registry remote \
  --env docker --env.registry remote --registry.host localhost --namespace myns \
  --set_env.0 VAR1 --set_env.1 VAR2=aguschin
💾 Saving deployment to deployer.mlem
```

```yaml
# deployer.mlem
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
set_env:
- VAR1
- VAR2=aguschin
type: kubernetes
```

Then in Dockerfile for your image you'll get:
```
...
ENV VAR1=VALUE1
ENV VAR2=aguschin
```

`VALUE1` will be taken from shell/env vars with `os.getenv("VAR1")`. If
it's not set, an Exception will be raised by MLEM.

mlem/contrib/docker/base.py

@@ -27,7 +27,7 @@
 from mlem.contrib.docker.context import (
     DockerBuildArgs,
     DockerModelDirectory,
-    get_build_args,
+    parse_envs,
 )
 from mlem.contrib.docker.utils import (
     build_image_with_logs,
@@ -596,7 +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),
+                    buildargs=parse_envs(self.args.build_arg),
                 )
                 docker_image = DockerImage(**self.image.dict())
                 docker_image.image_id = image.id

mlem/contrib/docker/context.py

@@ -20,6 +20,7 @@
 
 import mlem
 from mlem.config import MlemConfigBase, project_config
+from mlem.core.errors import EnvVarNotSet
 from mlem.core.objects import MlemModel
 from mlem.core.requirements import Requirements, UnixPackageRequirement
 from mlem.runtime.server import Server
@@ -245,7 +246,9 @@ class Config:
     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"""
+    """ARG vars to use at build time https://docs.docker.com/engine/reference/commandline/build/#build-arg"""
+    set_env: List[str] = []
+    """ENV vars to set in the image https://docs.docker.com/engine/reference/builder/#env"""
 
     def get_base_image(self):
         if self.base_image is None:
@@ -272,15 +275,18 @@ 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
+def parse_envs(names_values) -> Dict[str, str]:
+    envs = {}
+    for name in names_values:
+        if "=" in name:
+            name, value = name.split("=", 1)
+            envs[name] = os.getenv(name) or value
         else:
-            args[arg] = os.getenv(arg)
-    return args
+            value = os.getenv(name)
+            if value is None:
+                raise EnvVarNotSet(name)
+            envs[name] = value
+    return envs
 
 
 class DockerModelDirectory(BaseModel):
@@ -320,6 +326,8 @@ def get_env_vars(self) -> Dict[str, str]:
         ]
         if len(used_extensions) > 0:
             envs["MLEM_EXTENSIONS"] = ",".join(used_extensions)
+
+        envs.update(parse_envs(self.docker_args.set_env))
         return envs
 
     def get_python_version(self):
@@ -384,7 +392,7 @@ def write_dockerfile(self, requirements: Requirements):
                 **self.docker_args.dict()
             ).generate(
                 env=env,
-                arg=get_build_args(self.docker_args.build_arg),
+                arg=parse_envs(self.docker_args.build_arg),
                 packages=[p.package_name for p in unix_packages or []],
             )
             df.write(dockerfile)

mlem/contrib/flyio/meta.py

@@ -118,6 +118,20 @@ def _create_app(self, state: FlyioAppState):
             if port:  # tell flyio to expose specific port
                 fly_toml = parse(state.fly_toml)
                 fly_toml["services"][0]["internal_port"] = port
+                if self.server and self.server.middlewares:
+                    # set fly to collect metrics from prometheus if exposed
+                    from mlem.contrib.prometheus import (  # noqa
+                        PrometheusFastAPIMiddleware,
+                    )
+
+                    if any(
+                        isinstance(m, PrometheusFastAPIMiddleware)
+                        for m in self.server.middlewares.__root__
+                    ):
+                        fly_toml["metrics"] = {
+                            "port": port,
+                            "path": "/metrics",
+                        }
                 state.fly_toml = fly_toml.as_string()
             status = get_status(workdir=tempdir)
             state.app_name = status.Name

mlem/contrib/kubernetes/base.py

@@ -79,7 +79,9 @@ class K8sDeployment(
     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"""
+    """ARG vars to use at build time https://docs.docker.com/engine/reference/commandline/build/#build-arg"""
+    set_env: List[str] = []
+    """ENV vars to set in the image https://docs.docker.com/engine/reference/builder/#env"""
 
     def load_kube_config(self):
         config.load_kube_config(
@@ -137,6 +139,7 @@ def deploy(self, model: MlemModel):
                     daemon=self.daemon,
                     server=self.get_server(),
                     build_arg=self.build_arg,
+                    set_env=self.set_env,
                 )
                 state.update_model(model)
                 redeploy = True

mlem/contrib/kubernetes/build.py

@@ -17,6 +17,7 @@ def build_k8s_docker(
     platform: Optional[str] = "linux/amd64",
     # runners usually do not support arm64 images built on Mac M1 devices
     build_arg: Optional[List[str]] = None,
+    set_env: Optional[List[str]] = None,
 ):
     echo(EMOJI_BUILD + f"Creating docker image {image_name}")
     with set_offset(2):
@@ -30,4 +31,5 @@ def build_k8s_docker(
             force_overwrite=True,
             platform=platform,
             build_arg=build_arg or [],
+            set_env=set_env or [],
         )

mlem/core/errors.py

@@ -178,3 +178,11 @@ def __init__(self, ext: str, reqs: List[str], extra: Optional[str]):
         super().__init__(
             f"Extension '{ext}' requires additional dependencies: {extra_install}`pip install {reqs_str}`"
         )
+
+
+class EnvVarNotSet(MlemError):
+    def __init__(self, name: str):
+        self.name = name
+        super().__init__(
+            f'"{name}" was supposed to be read from shell/env vars, but was not found there'
+        )