Add fly.io deployments

mlem/contrib/flyio/__init__.py

@@ -0,0 +1,5 @@
+"""fly.io Deployments support
+Extension type: deployment
+
+Implements MlemEnv, MlemDeployment and DeployState to work with fly.io
+"""

mlem/contrib/flyio/meta.py

@@ -0,0 +1,181 @@
+import tempfile
+from typing import ClassVar, Optional
+
+from pydantic import BaseModel
+
+from mlem import ui
+from mlem.api import build
+from mlem.config import MlemConfigBase, project_config
+from mlem.contrib.docker import DockerDirBuilder
+from mlem.contrib.flyio.utils import (
+    check_flyctl_exec,
+    get_scale,
+    get_status,
+    place_fly_toml,
+    read_fly_toml,
+    run_flyctl,
+)
+from mlem.core.objects import (
+    DeployState,
+    DeployStatus,
+    MlemDeployment,
+    MlemEnv,
+    MlemModel,
+)
+from mlem.runtime.client import Client, HTTPClient
+from mlem.runtime.server import Server
+
+FLYIO_STATE_MAPPING = {
+    "running": DeployStatus.RUNNING,
+    "deployed": DeployStatus.RUNNING,
+    "pending": DeployStatus.STARTING,
+    "suspended": DeployStatus.STOPPED,
+    "dead": DeployStatus.CRASHED,
+}
+
+
+class FlyioConfig(MlemConfigBase):
+    class Config:
+        section = "flyio"
+
+    region: str = "lax"
+
+
+class FlyioSettings(BaseModel):
+    org: Optional[str] = None
+    """Organization name"""
+    region: Optional[str] = None
+    """Region name"""
+
+
+class FlyioEnv(MlemEnv, FlyioSettings):
+    """fly.io organization/account"""
+
+    type: ClassVar = "flyio"
+
+    access_token: Optional[str] = None
+    """Access token for fly.io. Alternatively use `flyctl auth login`"""
+
+
+class FlyioAppState(DeployState):
+    """fly.io app state"""
+
+    type: ClassVar = "flyio"
+
+    fly_toml: Optional[str]
+    """Contents of fly.toml file for app"""
+    app_name: Optional[str]
+    """Application name"""
+    hostname: Optional[str]
+    """Application hostname"""
+
+
+class FlyioApp(MlemDeployment, FlyioSettings):
+    """fly.io deployment"""
+
+    type: ClassVar = "flyio"
+
+    state_type: ClassVar = FlyioAppState
+    env_type: ClassVar = FlyioEnv
+
+    image: Optional[str] = None
+    """Image name for docker image"""
+    app_name: Optional[str] = None
+    """Application name. Leave empty for auto-generated one"""
+    scale_memory: Optional[int] = None
+    """Set VM memory to a number of megabytes (256/512/1024 etc)"""
+    # TODO other scale params
+
+    server: Optional[Server] = None
+    """Server to use"""
+
+    def _get_client(self, state: FlyioAppState) -> Client:
+        return HTTPClient(host=f"https://{state.hostname}", port=443)
+
+    def _create_app(self, state: FlyioAppState):
+        with tempfile.TemporaryDirectory(
+            prefix="mlem_flyio_build_"
+        ) as tempdir:
+            args = {
+                "auto-confirm": True,
+                "region": self.region
+                or self.get_env().region
+                or project_config("", section=FlyioConfig).region,
+            }
+            if self.app_name:
+                args["name"] = self.app_name
+            else:
+                args["generate-name"] = True
+            if self.get_env().access_token:
+                args["access-token"] = self.get_env().access_token
+            run_flyctl("launch", workdir=tempdir, kwargs=args)
+            state.fly_toml = read_fly_toml(tempdir)
+
+            status = get_status(workdir=tempdir)
+            state.app_name = status.Name
+            state.hostname = status.Hostname
+            self.update_state(state)
+
+    def _scale_app(self, state: FlyioAppState):
+        if self.scale_memory is None:
+            return
+        current_scale = get_scale(app_name=state.app_name)
+
+        if current_scale.MemoryMB != self.scale_memory:
+            run_flyctl(
+                f"scale memory {self.scale_memory}",
+                kwargs={"app": state.app_name},
+            )
+            ui.echo(f"Scaled {state.app_name} memory to {self.scale_memory}MB")
+
+    def _build_in_dir(self, model: MlemModel, state: FlyioAppState):
+        with tempfile.TemporaryDirectory(
+            prefix="mlem_flyio_build_"
+        ) as tempdir:
+            assert state.fly_toml is not None
+            place_fly_toml(tempdir, state.fly_toml)
+            build(DockerDirBuilder(target=tempdir, server=self.server), model)
+
+            args = {}
+            if self.get_env().access_token:
+                args["access-token"] = self.get_env().access_token
+
+            run_flyctl("deploy", workdir=tempdir, kwargs=args)
+            state.fly_toml = read_fly_toml(tempdir)
+            state.update_model(model)
+            self.update_state(state)
+
+    def deploy(self, model: MlemModel):
+        check_flyctl_exec()
+        with self.lock_state():
+            state: FlyioAppState = self.get_state()
+
+            if state.fly_toml is None:
+                self._create_app(state)
+
+            self._scale_app(state)
+
+            if self.model_changed(model, state):
+                self._build_in_dir(model, state)
+
+    def remove(self):
+        check_flyctl_exec()
+        with self.lock_state():
+            state: FlyioAppState = self.get_state()
+            if state.app_name is not None:
+                run_flyctl(
+                    f"apps destroy {state.app_name}", kwargs={"yes": True}
+                )  # lol
+            self.purge_state()
+
+    def get_status(self, raise_on_error=True) -> "DeployStatus":
+        check_flyctl_exec()
+        with self.lock_state():
+            state: FlyioAppState = self.get_state()
+
+            if state.fly_toml is None or state.app_name is None:
+                return DeployStatus.NOT_DEPLOYED
+
+            status = get_status(app_name=state.app_name)
+
+            return FLYIO_STATE_MAPPING[status.Status]

mlem/contrib/flyio/utils.py

@@ -0,0 +1,88 @@
+import json
+import os.path
+import subprocess
+from typing import Any, Dict
+
+from pydantic import BaseModel, parse_obj_as
+
+from mlem.core.errors import DeploymentError
+
+FLY_TOML = "fly.toml"
+
+
+def check_flyctl_exec():
+    try:
+        run_flyctl("version", wrap_error=False)
+    except subprocess.SubprocessError as e:
+        raise DeploymentError(
+            "flyctl executable is not available. Please install it using <https://fly.io/docs/hands-on/install-flyctl/>"
+        ) from e
+
+
+def run_flyctl(
+    command: str,
+    workdir: str = None,
+    kwargs: Dict[str, Any] = None,
+    wrap_error=True,
+):
+    kwargs = kwargs or {}
+    cmd = (
+        ["flyctl"]
+        + command.split(" ")
+        + " ".join(
+            [
+                f"--{k} {v}" if v is not True else f"--{k}"
+                for k, v in kwargs.items()
+            ]
+        ).split()
+    )
+    try:
+        return subprocess.check_output(cmd, cwd=workdir)
+    except subprocess.SubprocessError as e:
+        if wrap_error:
+            raise DeploymentError(e) from e
+        raise
+
+
+def read_fly_toml(workdir: str):
+    with open(os.path.join(workdir, FLY_TOML), encoding="utf8") as f:
+        return f.read()
+
+
+def place_fly_toml(workdir: str, fly_toml: str):
+    with open(os.path.join(workdir, FLY_TOML), "w", encoding="utf8") as f:
+        f.write(fly_toml)
+
+
+class FlyioStatusModel(BaseModel):
+    Name: str
+    Status: str
+    Hostname: str
+
+
+def get_status(workdir: str = None, app_name: str = None) -> FlyioStatusModel:
+    args: Dict[str, Any] = {"json": True}
+    if app_name is not None:
+        args["app"] = app_name
+    status = run_flyctl("status", kwargs=args, workdir=workdir)
+    return parse_obj_as(FlyioStatusModel, json.loads(status))
+
+
+class FlyioScaleModel(BaseModel):
+    Name: str
+    CPUCores: int
+    CPUClass: str
+    MemoryGB: float
+    MemoryMB: int
+    PriceMonth: float
+    PriceSecond: float
+    Count: str
+    MaxPerRegion: str
+
+
+def get_scale(workdir: str = None, app_name: str = None) -> FlyioScaleModel:
+    args: Dict[str, Any] = {"json": True}
+    if app_name is not None:
+        args["app"] = app_name
+    status = run_flyctl("scale show", kwargs=args, workdir=workdir)
+    return parse_obj_as(FlyioScaleModel, json.loads(status))

mlem/ext.py

@@ -129,6 +129,9 @@ class ExtensionLoader:
             False,
         ),
         Extension("mlem.contrib.git", ["pygit2"], True),
+        Extension(
+            "mlem.contrib.flyio", ["docker", "fastapi", "uvicorn"], False
+        ),
         Extension("mlem.contrib.torchvision", ["torchvision"], False),
     )
 

setup.py

@@ -91,6 +91,7 @@
     "kubernetes": ["docker", "kubernetes"],
     "dvc": ["dvc~=2.0"],
     "git": ["pygit2"],
+    "flyio": ["docker", "fastapi", "uvicorn"],
     "torchvision": ["torchvision"],
 }
 
@@ -166,6 +167,9 @@
             "artifact.dvc = mlem.contrib.dvc:DVCArtifact",
             "storage.dvc = mlem.contrib.dvc:DVCStorage",
             "server.fastapi = mlem.contrib.fastapi:FastAPIServer",
+            "deployment.flyio = mlem.contrib.flyio.meta:FlyioApp",
+            "deploy_state.flyio = mlem.contrib.flyio.meta:FlyioAppState",
+            "env.flyio = mlem.contrib.flyio.meta:FlyioEnv",
             "resolver.local_git = mlem.contrib.git:LocalGitResolver",
             "resolver.github = mlem.contrib.github:GithubResolver",
             "resolver.gitlab = mlem.contrib.gitlabfs:GitlabResolver",
@@ -245,6 +249,7 @@
             "core = mlem.config:MlemConfig",
             "bitbucket = mlem.contrib.bitbucketfs:BitbucketConfig",
             "docker = mlem.contrib.docker.context:DockerConfig",
+            "flyio = mlem.contrib.flyio.meta:FlyioConfig",
             "heroku = mlem.contrib.heroku.config:HerokuConfig",
             "pandas = mlem.contrib.pandas:PandasConfig",
             "aws = mlem.contrib.sagemaker.config:AWSConfig",