Add server field to heroku deploy (#631)

closes #568

Heroku requires you to use `PORT` env to expose a port.
I changed `HerokuServer` to wrap arbitrary server (instead of
subclassing `FastAPIServer`) and swap port for value from env. For that
it needs to know which value is to swap, so I added `port_field`
classvar to servers.

---------

Co-authored-by: Alexander Guschin <[email protected]>

Author: mike0sv

mlem/contrib/fastapi.py

@@ -53,6 +53,7 @@ class FastAPIServer(Server, LibRequirementsMixin):
 
     libraries: ClassVar[List[ModuleType]] = [uvicorn, fastapi]
     type: ClassVar[str] = "fastapi"
+    port_field: ClassVar = "port"
 
     host: str = "0.0.0.0"
     """Network interface to use"""

mlem/contrib/heroku/build.py

@@ -3,8 +3,9 @@
 
 from mlem.core.objects import MlemModel
 
+from ...runtime.server import Server
 from ...ui import EMOJI_BUILD, echo, set_offset
-from ..docker.base import DockerEnv, DockerImage, RemoteRegistry
+from ..docker.base import DockerImage, RemoteRegistry
 from ..docker.helpers import build_model_image
 from .server import HerokuServer
 
@@ -41,23 +42,21 @@ def login(self, client):
 
 def build_heroku_docker(
     meta: MlemModel,
+    server: Server,
     app_name: str,
     process_type: str = "web",
     api_key: str = None,
     push: bool = True,
 ) -> DockerImage:
-    docker_env = DockerEnv(
-        registry=HerokuRemoteRegistry(
-            host="registry.heroku.com", api_key=api_key
-        )
-    )
     echo(EMOJI_BUILD + "Creating docker image for heroku")
     with set_offset(2):
         return build_model_image(
             meta,
             process_type,
-            server=HerokuServer(),
-            env=docker_env,
+            server=HerokuServer(server=server),
+            registry=HerokuRemoteRegistry(
+                host="registry.heroku.com", api_key=api_key
+            ),
             repository=app_name,
             force_overwrite=True,
             # heroku does not support arm64 images built on Mac M1 devices

mlem/contrib/heroku/meta.py

@@ -12,7 +12,9 @@
 )
 from mlem.runtime.client import Client, HTTPClient
 
+from ...config import project_config
 from ...core.errors import DeploymentError
+from ...runtime.server import Server
 from ...ui import EMOJI_OK, echo
 from ..docker.base import DockerImage
 from .build import build_heroku_docker
@@ -77,6 +79,8 @@ class HerokuDeployment(MlemDeployment[HerokuState, HerokuEnv]):
     """Stack to use"""
     team: Optional[str] = None
     """Heroku team"""
+    server: Optional[Server] = None
+    """Server to use"""
 
     def _get_client(self, state: HerokuState) -> Client:
         return HTTPClient(
@@ -95,7 +99,13 @@ def deploy(self, model: MlemModel):
             redeploy = False
             if state.image is None or self.model_changed(model):
                 state.image = build_heroku_docker(
-                    model, state.app.name, api_key=self.get_env().api_key
+                    model,
+                    self.server
+                    or project_config(
+                        self.loc.project if self.is_saved else None
+                    ).server,
+                    state.app.name,
+                    api_key=self.get_env().api_key,
                 )
                 state.update_model(model)
                 self.update_state(state)

mlem/contrib/heroku/server.py

@@ -1,19 +1,47 @@
 import logging
 import os
-from typing import ClassVar
+from typing import ClassVar, Dict
 
-from mlem.contrib.fastapi import FastAPIServer
+from pydantic import validator
+
+from mlem.core.requirements import Requirements
 from mlem.runtime import Interface
+from mlem.runtime.server import Server
 
 logger = logging.getLogger(__name__)
 
 
-class HerokuServer(FastAPIServer):
+class HerokuServer(Server):
     """Special FastAPI server to pickup port from env PORT"""
 
     type: ClassVar = "_heroku"
+    server: Server
+
+    @validator("server")
+    @classmethod
+    def server_validator(cls, value: Server):
+        if value.port_field is None:
+            raise ValueError(
+                f"{value} does not have port field and can not be exposed on heroku"
+            )
+        return value
 
     def serve(self, interface: Interface):
-        self.port = int(os.environ["PORT"])
-        logger.info("Switching port to %s", self.port)
-        return super().serve(interface)
+        assert self.server.port_field is not None  # ensured by validator
+        setattr(self.server, self.server.port_field, int(os.environ["PORT"]))
+        logger.info(
+            "Switching port to %s",
+            getattr(self.server, self.server.port_field),
+        )
+        return self.server.serve(interface)
+
+    def get_requirements(self) -> Requirements:
+        return self.server.get_requirements()
+
+    def get_env_vars(self) -> Dict[str, str]:
+        env_vars = super().get_env_vars()
+        env_vars.update(self.server.get_env_vars())
+        return env_vars
+
+    def get_sources(self) -> Dict[str, bytes]:
+        return self.server.get_sources()

mlem/contrib/streamlit/server.py

@@ -5,7 +5,7 @@
 import tempfile
 from threading import Thread
 from time import sleep
-from typing import ClassVar, Dict, Optional
+from typing import ClassVar, Dict, List, Optional
 
 import streamlit
 import streamlit_pydantic
@@ -52,6 +52,7 @@ class StreamlitServer(Server, StreamlitScript, LibRequirementsMixin):
 
     type: ClassVar = "streamlit"
     libraries: ClassVar = (streamlit, streamlit_pydantic)
+    port_field: ClassVar = "ui_port"
 
     run_server: bool = True
     """Whether to run backend server or use existing one"""
@@ -159,3 +160,6 @@ def get_sources(self) -> Dict[str, bytes]:
                 sources[template] = f.read()
             self.template = template
         return sources
+
+    def get_ports(self) -> List[int]:
+        return [self.ui_port, self.server_port]

mlem/runtime/server.py

@@ -118,6 +118,7 @@ class Config:
     abs_name: ClassVar[str] = "server"
     env_vars: ClassVar[Optional[Dict[str, str]]] = None
     additional_source_files: ClassVar[Optional[List[str]]] = None
+    port_field: ClassVar[Optional[str]] = None
 
     # @validator("interface")
     # @classmethod
@@ -158,6 +159,11 @@ def get_requirements(self) -> Requirements:
             [self.request_serializer, self.response_serializer, self.methods]
         )
 
+    def get_ports(self) -> List[int]:
+        if self.port_field is not None:
+            return [getattr(self, self.port_field)]
+        return []
+
 
 class ServerInterface(Interface):
     type: ClassVar = "_server"

tests/contrib/test_heroku.py

@@ -4,11 +4,14 @@
 
 import pytest
 import requests
+import testcontainers.core.container
 from docker import DockerClient
 from sklearn.datasets import load_iris
 from sklearn.ensemble import RandomForestClassifier
+from testcontainers.core.waiting_utils import wait_for_logs
 
 from mlem.api.commands import deploy
+from mlem.contrib.fastapi import FastAPIServer
 from mlem.contrib.heroku.build import (
     DEFAULT_HEROKU_REGISTRY,
     build_heroku_docker,
@@ -28,6 +31,7 @@
 )
 from mlem.core.errors import DeploymentError
 from mlem.core.objects import DeployStatus, MlemModel
+from mlem.runtime.client import HTTPClient
 from tests.conftest import flaky, long
 
 heroku = pytest.mark.skipif(
@@ -97,12 +101,21 @@ def test_create_app(heroku_app_name, heroku_env, model):
 
 @long
 def test_build_heroku_docker(model: MlemModel, uses_docker_build):
-    image_meta = build_heroku_docker(model, "test_build", push=False)
+    image_meta = build_heroku_docker(
+        model, FastAPIServer(), "test_build", push=False
+    )
     client = DockerClient.from_env()
     image = client.images.get(image_meta.image_id)
     assert image is not None
     assert f"{DEFAULT_HEROKU_REGISTRY}/test_build/web:latest" in image.tags
 
+    with testcontainers.core.container.DockerContainer(
+        image_meta.uri
+    ).with_env("PORT", "4567").with_bind_ports(4567, 80) as c:
+        wait_for_logs(c, ".*Uvicorn running on.*")
+        client = HTTPClient(port=80)
+        assert client.interface is not None
+
 
 def test_state_ensured_app():
     state = HerokuState(declaration=HerokuDeployment(app_name=""))