Sagemaker deployments (#366)

* WIP

* its alive (kinda)

* it works but it's ugly

* little less ugly

* lil fix

* fix lint

* fix lint

* fix tests

* fix tests

* fix windows bugs

* fix tests

* fix tests

mlem/api/commands.py

@@ -420,9 +420,11 @@ def deploy(
     fs: Optional[AbstractFileSystem] = None,
     external: bool = None,
     index: bool = None,
+    env_kwargs: Dict[str, Any] = None,
     **deploy_kwargs,
 ) -> MlemDeployment:
     deploy_path = None
+    update = False
     if isinstance(deploy_meta_or_path, str):
         deploy_path = deploy_meta_or_path
         try:
@@ -432,13 +434,13 @@ def deploy(
                 fs=fs,
                 force_type=MlemDeployment,
             )
+            update = True
         except MlemObjectNotFound:
             deploy_meta = None
 
     else:
         deploy_meta = deploy_meta_or_path
-        if model is not None:
-            deploy_meta.replace_model(get_model_meta(model))
+        update = True
 
     if deploy_meta is None:
         if model is None or env is None:
@@ -451,14 +453,21 @@ def deploy(
         env_meta = ensure_meta(MlemEnv, env, allow_typename=True)
         if isinstance(env_meta, type):
             env = None
+            if env_kwargs:
+                env = env_meta(**env_kwargs)
         deploy_type = env_meta.deploy_type
         deploy_meta = deploy_type(
             model_cache=model_meta,
+            model=model_meta.make_link(),
             env=env,
-            model=model,
             **deploy_kwargs,
         )
         deploy_meta.dump(deploy_path, fs, project, index, external)
+    else:
+        if model is not None:
+            deploy_meta.replace_model(get_model_meta(model, load_value=False))
+    if update:
+        pass  # todo update from deploy_args and env_args
     # ensuring links are working
     deploy_meta.get_env()
     deploy_meta.get_model()

mlem/cli/deployment.py

@@ -22,7 +22,7 @@
 from mlem.core.data_type import DataAnalyzer
 from mlem.core.errors import DeploymentError
 from mlem.core.metadata import load_meta
-from mlem.core.objects import DeployState, MlemDeployment
+from mlem.core.objects import DeployState, DeployStatus, MlemDeployment
 from mlem.ui import echo, no_echo, set_echo
 
 deployment = Typer(
@@ -67,13 +67,17 @@ def deploy_run(
     """
     from mlem.api.commands import deploy
 
+    conf = conf or []
+    env_conf = [c[len("env.") :] for c in conf if c.startswith("env.")]
+    conf = [c for c in conf if not c.startswith("env.")]
     deploy(
         path,
         model,
         env,
         project,
         external=external,
         index=index,
+        env_kwargs=parse_string_conf(env_conf),
         **parse_string_conf(conf or []),
     )
 
@@ -110,6 +114,40 @@ def deploy_status(
     echo(status)
 
 
+@mlem_command("wait", parent=deployment)
+def deploy_wait(
+    path: str = Argument(..., help="Path to deployment meta"),
+    project: Optional[str] = option_project,
+    statuses: List[DeployStatus] = Option(
+        [DeployStatus.RUNNING],
+        "-s",
+        "--status",
+        help="statuses to wait for",
+    ),
+    intermediate: List[DeployStatus] = Option(
+        None, "-i", "--intermediate", help="Possible intermediate statuses"
+    ),
+    poll_timeout: float = Option(
+        1.0, "-p", "--poll-timeout", help="Timeout between attempts"
+    ),
+    times: int = Option(
+        0, "-t", "--times", help="Number of attempts. 0 -> indefinite"
+    ),
+):
+    """Wait for status of deployed service
+
+    Examples:
+        $ mlem deployment status service_name
+    """
+    with no_echo():
+        deploy_meta = load_meta(
+            path, project=project, force_type=MlemDeployment
+        )
+        deploy_meta.wait_for_status(
+            statuses, poll_timeout, times, allowed_intermediate=intermediate
+        )
+
+
 @mlem_command("apply", parent=deployment)
 def deploy_apply(
     path: str = Argument(..., help="Path to deployment meta"),
@@ -144,7 +182,7 @@ def deploy_apply(
             raise DeploymentError(
                 f"{deploy_meta.type} deployment has no state. Either {deploy_meta.type} is not deployed yet or has been un-deployed again."
             )
-        client = state.get_client()
+        client = deploy_meta.get_client(state)
 
         result = run_apply_remote(
             client,

mlem/contrib/docker/base.py

@@ -187,10 +187,10 @@ def push(self, client, tag):
             if "error" in status:
                 error_msg = status["error"]
                 raise DeploymentError(f"Cannot push docker image: {error_msg}")
-        echo(EMOJI_OK + f"Pushed image {tag} to {self.host}")
+        echo(EMOJI_OK + f"Pushed image {tag} to {self.get_host()}")
 
     def uri(self, image: str):
-        return f"{self.host}/{image}"
+        return f"{self.get_host()}/{image}"
 
     def _get_digest(self, name, tag):
         r = requests.head(
@@ -286,9 +286,6 @@ class DockerContainerState(DeployState):
     container_name: Optional[str]
     container_id: Optional[str]
 
-    def get_client(self):
-        raise NotImplementedError
-
 
 class _DockerBuildMixin(BaseModel):
     server: Optional[Server] = None
@@ -320,6 +317,9 @@ class DockerContainer(MlemDeployment, _DockerBuildMixin):
     def ensure_image_name(self):
         return self.image_name or self.container_name
 
+    def _get_client(self, state: DockerContainerState):
+        raise NotImplementedError
+
 
 class DockerEnv(MlemEnv[DockerContainer]):
     """:class:`.MlemEnv` implementation for docker environment

mlem/contrib/docker/copy.j2

@@ -0,0 +1 @@
+COPY . ./

mlem/contrib/docker/dockerfile.j2

@@ -1,12 +1,9 @@
 FROM {{ base_image }}
 WORKDIR /app
 {% include "pre_install.j2" ignore missing %}
-{% if packages %}RUN {{ package_install_cmd }} {{ packages|join(" ")  }}{% endif %}
-COPY requirements.txt .
-RUN pip install -r requirements.txt
-{{ mlem_install }}
+{% include "install_req.j2" %}
 {% include "post_install.j2" ignore missing %}
-COPY . ./
+{% include "copy.j2" %}
 {% for name, value in env.items() %}ENV {{ name }}={{ value }}
 {% endfor %}
 {% include "post_copy.j2" ignore missing %}

mlem/contrib/docker/install_req.j2

@@ -0,0 +1,4 @@
+{% if packages %}RUN {{ package_install_cmd }} {{ packages|join(" ")  }}{% endif %}
+COPY requirements.txt .
+RUN pip install -r requirements.txt
+{{ mlem_install }}

mlem/contrib/heroku/meta.py

@@ -44,11 +44,6 @@ def ensured_app(self) -> HerokuAppMeta:
             raise ValueError("App is not created yet")
         return self.app
 
-    def get_client(self) -> Client:
-        return HTTPClient(
-            host=urlparse(self.ensured_app.web_url).netloc, port=80
-        )
-
 
 class HerokuDeployment(MlemDeployment):
     type: ClassVar = "heroku"
@@ -58,6 +53,11 @@ class HerokuDeployment(MlemDeployment):
     stack: str = "container"
     team: Optional[str] = None
 
+    def _get_client(self, state: HerokuState) -> Client:
+        return HTTPClient(
+            host=urlparse(state.ensured_app.web_url).netloc, port=80
+        )
+
 
 class HerokuEnv(MlemEnv[HerokuDeployment]):
     type: ClassVar = "heroku"

mlem/contrib/sagemaker/build.py

@@ -0,0 +1,123 @@
+import base64
+import os
+from typing import ClassVar, Optional
+
+import boto3
+import sagemaker
+from pydantic import BaseModel
+
+from ...core.objects import MlemModel
+from ...ui import EMOJI_BUILD, EMOJI_KEY, echo, set_offset
+from ..docker.base import DockerEnv, DockerImage, RemoteRegistry
+from ..docker.helpers import build_model_image
+from .runtime import SageMakerServer
+
+IMAGE_NAME = "mlem-sagemaker-runner"
+
+
+class AWSVars(BaseModel):
+    profile: str
+    bucket: str
+    region: str
+    account: str
+    role_name: str
+
+    @property
+    def role(self):
+        return f"arn:aws:iam::{self.account}:role/{self.role_name}"
+
+    def get_sagemaker_session(self):
+        return sagemaker.Session(
+            self.get_session(), default_bucket=self.bucket
+        )
+
+    def get_session(self):
+        return boto3.Session(
+            profile_name=self.profile, region_name=self.region
+        )
+
+
+def ecr_repo_check(region, repository, session: boto3.Session):
+    client = session.client("ecr", region_name=region)
+
+    repos = client.describe_repositories()["repositories"]
+
+    if repository not in {r["repositoryName"] for r in repos}:
+        echo(EMOJI_BUILD + f"Creating ECR repository {repository}")
+        client.create_repository(repositoryName=repository)
+
+
+class ECRegistry(RemoteRegistry):
+    class Config:
+        exclude = {"aws_vars"}
+
+    type: ClassVar = "ecr"
+    account: str
+    region: str
+
+    aws_vars: Optional[AWSVars] = None
+
+    def login(self, client):
+        auth_data = self.ecr_client.get_authorization_token()
+        token = auth_data["authorizationData"][0]["authorizationToken"]
+        user, token = base64.b64decode(token).decode("utf8").split(":")
+        self._login(self.get_host(), client, user, token)
+        echo(
+            EMOJI_KEY
+            + f"Logged in to remote registry at host {self.get_host()}"
+        )
+
+    def get_host(self) -> Optional[str]:
+        return f"{self.account}.dkr.ecr.{self.region}.amazonaws.com"
+
+    def image_exists(self, client, image: DockerImage):
+        images = self.ecr_client.list_images(repositoryName=image.name)[
+            "imageIds"
+        ]
+        return len(images) > 0
+
+    def delete_image(self, client, image: DockerImage, force=False, **kwargs):
+        self.ecr_client.batch_delete_image(
+            repositoryName=image.name,
+            imageIds=[{"imageTag": image.tag}],
+        )
+
+    def with_aws_vars(self, aws_vars):
+        self.aws_vars = aws_vars
+        return self
+
+    @property
+    def ecr_client(self):
+        return (
+            self.aws_vars.get_session().client("ecr")
+            if self.aws_vars
+            else boto3.client("ecr", region_name=self.region)
+        )
+
+
+def build_sagemaker_docker(
+    meta: MlemModel,
+    method: str,
+    account: str,
+    region: str,
+    image_name: str,
+    repository: str,
+    aws_vars: AWSVars,
+):
+    docker_env = DockerEnv(
+        registry=ECRegistry(account=account, region=region).with_aws_vars(
+            aws_vars
+        )
+    )
+    ecr_repo_check(region, repository, aws_vars.get_session())
+    echo(EMOJI_BUILD + "Creating docker image for sagemaker")
+    with set_offset(2):
+        return build_model_image(
+            meta,
+            name=repository,
+            tag=image_name,
+            server=SageMakerServer(method=method),
+            env=docker_env,
+            force_overwrite=True,
+            templates_dir=[os.path.dirname(__file__)],
+        )