qwak-sdk 0.5.102__py3-none-any.whl

qwak_sdk/commands/models/deployments/deploy/_logic/deployment.py

@@ -0,0 +1,405 @@
+from abc import ABC
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Dict, List
+
+from _qwak_proto.qwak.audience.v1.audience_pb2 import (
+    AudienceRoutesEntry as AudienceProto,
+)
+from _qwak_proto.qwak.audience.v1.audience_pb2 import Route as RouteProto
+from _qwak_proto.qwak.auto_scaling.v1.auto_scaling_pb2 import (
+    AutoScalingConfig as AutoScalingProto,
+)
+from _qwak_proto.qwak.deployment.deployment_pb2 import (
+    DeploymentDetails as DeploymentProto,
+)
+from _qwak_proto.qwak.deployment.deployment_pb2 import (
+    RealTimeConfig as RealTimeConfigProto,
+)
+from _qwak_proto.qwak.deployment.deployment_pb2 import (
+    ServingStrategy as ServingStrategyProto,
+)
+from _qwak_proto.qwak.deployment.deployment_pb2 import StreamConfig as StreamConfigProto
+
+
+@dataclass
+class TrafficSpec:
+    percentage: int = field(default=None)
+    is_shadow: bool = field(default=None)
+
+
+@dataclass
+class Variation:
+    name: str = field(default=None)
+    environment_id: str = field(default=None)
+    traffic_spec: TrafficSpec = field(default=None)
+
+
+@dataclass
+class Route:
+    variation_name: str = field(default=None)
+    weight: int = field(default=None)
+    shadow: bool = field(default=None)
+    model_id: str = field(default=None)
+    environment_id: str = field(default=None)
+    audience_id: str = field(default=None)
+
+
+@dataclass
+class AudienceRoute:
+    audience_id: str = field(default=None)
+    audience_name: str = field(default=None)
+    routes: List[Route] = field(default_factory=list)
+    order: int = field(default=None)
+
+
+@dataclass
+class EnvironmentAudienceRoute(AudienceRoute):
+    environment_id: str = field(default=None)
+
+    @staticmethod
+    def from_proto(environment_id: str, proto: AudienceProto):
+        return EnvironmentAudienceRoute(
+            audience_id=proto.audience_id,
+            audience_name=proto.audience_name,
+            routes=Deployment.parse_routes(proto.routes),
+            order=proto.order,
+            environment_id=environment_id,
+        )
+
+
+class Metric(Enum):
+    NOT_VALID = 0
+    CPU = 1
+    MEMORY = 2
+    LATENCY = 3
+    GPU = 4
+    ERROR_RATE = 5
+    THROUGHPUT = 6
+
+
+class Aggregation(Enum):
+    NOT_VALID = 0
+    MIN = 1
+    MAX = 2
+    AVERAGE = 3
+    SUM = 4
+
+
+@dataclass
+class QuerySpec:
+    metric: Metric = field(default=None)
+    aggregation: Aggregation = field(default=None)
+    time_period: int = field(default=None)
+    error_code: str = field(default=None)
+
+
+class Triggers(ABC):
+    pass
+
+
+@dataclass
+class PrometheusTrigger(Triggers):
+    query_spec: QuerySpec = field(default=None)
+    threshold: int = field(default=None)
+
+
+@dataclass
+class AutoScalingConfig:
+    min_replica_count: int = field(default=None)
+    max_replica_count: int = field(default=None)
+    pooling_interval: int = field(default=None)
+    cool_down_period: int = field(default=None)
+    triggers: List[Triggers] = field(default_factory=list)
+
+
+@dataclass
+class TrafficConfig:
+    selected_variation_name: str = field(default=None)
+    variations: List[Variation] = field(default_factory=list)
+    audiences: List[AudienceRoute] = field(default_factory=list)
+    fallback_variation: str = field(default=None)
+
+
+class ServingConfig(ABC):
+    pass
+
+
+@dataclass
+class RealTimeConfig(ServingConfig):
+    traffic_config: TrafficConfig = field(default=None)
+    autoscaling_config: AutoScalingConfig = field(default=None)
+
+
+@dataclass
+class BatchConfig(ServingConfig):
+    pass
+
+
+class AutoOffSet(Enum):
+    UNKNOWN = 0
+    EARLIEST = 1
+    LATEST = 2
+
+
+@dataclass
+class Consumer:
+    bootstrap_server: List[str] = field(default_factory=list)
+    topic: str = field(default=None)
+    group: str = field(default=None)
+    timeout: int = field(default=None)
+    auto_offset: AutoOffSet = field(default=None)
+    max_batch_size: int = field(default=None)
+    max_poll_latency: float = field(default=None)
+
+
+class Compression(Enum):
+    UNKNOWN = 0
+    UNCOMPRESSED = 1
+    GZIP = 2
+    SNAPPY = 3
+    LZ4 = 4
+    ZSTD = 5
+
+
+@dataclass
+class Producer:
+    bootstrap_server: List[str] = field(default_factory=list)
+    topic: str = field(default=None)
+    compression: Compression = field(default=None)
+
+
+@dataclass
+class KafkaConfig:
+    consumer: Consumer = field(default=None)
+    producer: Producer = field(default=None)
+    workers: int = field(default=None)
+
+
+@dataclass
+class StreamConfig(ServingConfig):
+    kafka: KafkaConfig = field(default=None)
+
+
+class BuildStatus(Enum):
+    INVALID = 0
+    IN_PROGRESS = 1
+    SUCCESSFUL = 2
+    FAILED = 3
+    REMOTE_BUILD_INITIALIZING = 4
+    REMOTE_BUILD_CANCELLED = 5
+    REMOTE_BUILD_TIMED_OUT = 6
+    REMOTE_BUILD_UNKNOWN = 7
+    SYNCING_ENVIRONMENTS = 8
+    FINISHED_SYNCING = 9
+
+
+class MemoryUnit(Enum):
+    UNKNOWN = 0
+    MIB = 1
+    GIB = 2
+
+
+class KubeDeploymentType(Enum):
+    UNDEFINED = 0
+    ONLINE = 1
+    STREAM = 2
+    BATCH = 3
+
+
+class GpuType(Enum):
+    INVALID_GPU = 0
+    ONLINE = 1
+    STREAM = 2
+    BATCH = 3
+
+
+@dataclass
+class Deployment:
+    build_id: str = field(default=None)
+    number_of_pods: int = field(default=None)
+    cpu_fraction: float = field(default=None)
+    memory_amount: int = field(default=None)
+    memory_units: MemoryUnit = field(default=MemoryUnit.UNKNOWN)
+    number_of_workers: int = field(default=None)
+    http_request_timeout_ms: int = field(default=None)
+    kube_deployment_type: KubeDeploymentType = field(
+        default=KubeDeploymentType.UNDEFINED
+    )
+    serving_config: ServingConfig = field(default=None)
+    variation: Variation = field(default=None)
+    custom_iam_role_arn: str = field(default=None)
+    max_batch_size: int = field(default=None)
+    gpu_type: GpuType = field(default=GpuType.INVALID_GPU)
+    gpu_amount: int = field(default=None)
+    environment_id: str = field(default=None)
+    available_replicas: int = field(default=None)
+    deployment_process_timeout_limit: int = field(default=None)
+    daemon_mode: bool = field(default=None)
+    purchase_option: str = field(default=None)
+    environment_variables: Dict[str, str] = field(default_factory=dict)
+    last_deployed: datetime = field(default=datetime.now())
+
+    @staticmethod
+    def parse_routes(proto: List[RouteProto]) -> List[Route]:
+        routes = []
+        for proto_route in proto:
+            route = Route(
+                variation_name=proto_route.variation_name,
+                weight=proto_route.weight,
+                shadow=proto_route.shadow,
+                model_id=proto_route.model_id,
+                environment_id=proto_route.environment_id,
+                audience_id=proto_route.audience_id,
+            )
+            routes.append(route)
+        return routes
+
+    @staticmethod
+    def parse_audience_routes(proto: List[RouteProto]) -> List[AudienceRoute]:
+        routes = []
+        for proto_route in proto:
+            route = AudienceRoute(
+                audience_id=proto_route.audience_id,
+                audience_name=proto_route.audience_name,
+                routes=Deployment.parse_routes(proto_route.routes),
+                order=proto_route.order,
+            )
+            routes.append(route)
+        return routes
+
+    @staticmethod
+    def parse_autoscaling(proto: AutoScalingProto) -> AutoScalingConfig:
+        triggers = []
+        for proto_trigger in proto.triggers.triggers:
+            trigger = PrometheusTrigger(
+                query_spec=QuerySpec(
+                    metric=Metric(
+                        proto_trigger.prometheus_trigger.query_spec.metric_type
+                    ),
+                    aggregation=Aggregation(
+                        proto_trigger.prometheus_trigger.query_spec.aggregation_type
+                    ),
+                    time_period=proto_trigger.prometheus_trigger.query_spec.time_period,
+                    error_code=proto_trigger.prometheus_trigger.query_spec.error_code,
+                ),
+                threshold=proto_trigger.prometheus_trigger.threshold,
+            )
+            triggers.append(trigger)
+        return AutoScalingConfig(
+            min_replica_count=proto.min_replica_count,
+            max_replica_count=proto.max_replica_count,
+            pooling_interval=proto.polling_interval,
+            cool_down_period=proto.cool_down_period,
+            triggers=triggers,
+        )
+
+    @staticmethod
+    def parse_realtime_config(proto: RealTimeConfigProto) -> RealTimeConfig:
+        variations = []
+        for proto_variation in proto.traffic_config.variations:
+            variation = Variation(
+                name=proto_variation.name,
+                environment_id=proto_variation.environment_id,
+                traffic_spec=TrafficSpec(
+                    percentage=proto_variation.traffic.percentage,
+                    is_shadow=proto_variation.traffic.is_shadow,
+                ),
+            )
+            variations.append(variation)
+        audiences = []
+        for proto_audience in proto.traffic_config.audience_routes_entries:
+            audience = AudienceRoute(
+                audience_id=proto_audience.audience_id,
+                audience_name=proto_audience.audience_name,
+                routes=Deployment.parse_routes(proto_audience.routes),
+                order=proto_audience.order,
+            )
+            audiences.append(audience)
+
+        autoscaling = Deployment.parse_autoscaling(proto.auto_scaling_config)
+        return RealTimeConfig(
+            traffic_config=TrafficConfig(
+                variations=variations,
+                audiences=audiences,
+                fallback_variation=proto.traffic_config.fallback_variation,
+            ),
+            autoscaling_config=autoscaling,
+        )
+
+    @staticmethod
+    def parse_streaming_config(proto_config: StreamConfigProto) -> StreamConfig:
+        return StreamConfig(
+            kafka=KafkaConfig(
+                consumer=Consumer(
+                    bootstrap_server=list(proto_config.kafka.consumer.bootstrap_server),
+                    topic=proto_config.kafka.consumer.topic,
+                    group=proto_config.kafka.consumer.group,
+                    timeout=proto_config.kafka.consumer.timeout,
+                    auto_offset=AutoOffSet(proto_config.kafka.consumer.auto_offset_type)
+                    if proto_config.kafka.consumer.auto_offset_type
+                    else None,
+                    max_batch_size=proto_config.kafka.consumer.max_batch_size,
+                    max_poll_latency=proto_config.kafka.consumer.max_poll_latency,
+                ),
+                producer=Producer(
+                    bootstrap_server=list(proto_config.kafka.producer.bootstrap_server),
+                    topic=proto_config.kafka.producer.topic,
+                    compression=Compression(
+                        proto_config.kafka.producer.compression_type
+                    )
+                    if proto_config.kafka.producer.compression_type
+                    else None,
+                ),
+                workers=proto_config.kafka.workers,
+            )
+        )
+
+    @staticmethod
+    def parse_serving_strategy(proto: ServingStrategyProto) -> ServingConfig:
+        which = proto.WhichOneof("Strategy")
+        if which == "realtime_config":
+            return Deployment.parse_realtime_config(proto.realtime_config)
+        elif which == "batch_config":
+            return BatchConfig()
+        elif which == "stream_config":
+            return Deployment.parse_streaming_config(proto.stream_config)
+        else:
+            raise ValueError(f"Unknown serving strategy: {which}")
+
+    @staticmethod
+    def from_proto(proto: DeploymentProto):
+        variation = Variation(
+            name=proto.variation.name,
+            environment_id=proto.variation.environment_id,
+            traffic_spec=TrafficSpec(
+                percentage=proto.variation.traffic.percentage,
+                is_shadow=proto.variation.traffic.is_shadow,
+            ),
+        )
+        return Deployment(
+            build_id=proto.build_id,
+            number_of_pods=proto.number_of_pods,
+            cpu_fraction=proto.cpu_fraction,
+            memory_amount=proto.memory_amount,
+            memory_units=MemoryUnit(proto.memory_units),
+            number_of_workers=proto.number_of_workers,
+            http_request_timeout_ms=proto.http_request_timeout_ms,
+            kube_deployment_type=KubeDeploymentType(proto.kube_deployment_type),
+            serving_config=Deployment.parse_serving_strategy(proto.serving_strategy),
+            variation=variation,
+            custom_iam_role_arn=proto.custom_iam_role_arn,
+            max_batch_size=proto.max_batch_size,
+            gpu_type=GpuType(proto.gpu_type),
+            gpu_amount=proto.gpu_amount,
+            environment_id=proto.environment_id,
+            available_replicas=proto.available_replicas,
+            daemon_mode=bool(proto.daemon_mode.value),
+            deployment_process_timeout_limit=proto.deployment_process_timeout_limit,
+            purchase_option=proto.purchase_option,
+            environment_variables=dict(proto.environment_variables),
+            last_deployed=datetime.fromtimestamp(
+                proto.last_deployed.seconds + proto.last_deployed.nanos / 1e9
+            ),
+        )

qwak_sdk/commands/models/deployments/deploy/_logic/deployment_message_helpers.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from typing import Dict
+
+from _qwak_proto.qwak.deployment.deployment_pb2 import (
+    BatchConfig,
+    EnvironmentDeploymentMessage,
+    HostingService,
+    KubeDeployment,
+    KubeDeploymentType,
+    ServingStrategy,
+)
+from qwak.clients.administration.eco_system.client import EcosystemClient
+from qwak.clients.instance_template.client import InstanceTemplateManagementClient
+from qwak.exceptions import QwakException
+
+from qwak_sdk.commands.models.deployments.deploy._logic.advance_deployment_options_handler import (
+    get_advanced_deployment_options_from_deploy_config,
+)
+from qwak_sdk.commands.models.deployments.deploy._logic.deploy_config import (
+    DeployConfig,
+)
+from qwak_sdk.commands.models.deployments.deploy._logic.deployment_size_mapper import (
+    deployment_size_from_deploy_config,
+)
+from qwak_sdk.commands.models.deployments.deploy.realtime._logic.serving_strategy_mapper import (
+    create_realtime_serving_strategy_from_deploy_config,
+)
+from qwak_sdk.commands.models.deployments.deploy.streaming._logic.serving_strategy_mapper import (
+    create_streaming_serving_strategy_from_deploy_config,
+)
+from qwak_sdk.inner.tools.cli_tools import dictify_params
+
+UNKNOWN_SERVING_STRATEGY = (
+    "The deployments type doesn't have a serving strategy configured"
+)
+
+
+def get_env_to_deployment_message(
+    deploy_config: DeployConfig,
+    kube_deployment_type: KubeDeploymentType,
+    ecosystem_client: EcosystemClient,
+    instance_template_client: InstanceTemplateManagementClient,
+) -> Dict[str, EnvironmentDeploymentMessage]:
+    deployment_size = deployment_size_from_deploy_config(
+        deploy_config, instance_template_client, ecosystem_client
+    )
+    advanced_deployment_options = get_advanced_deployment_options_from_deploy_config(
+        deploy_config, kube_deployment_type
+    )
+    environment_variables = {}
+
+    if len(deploy_config.realtime.environments) != 0:
+        deploy_config.environments = deploy_config.realtime.environments
+
+    environment_name_to_config = ecosystem_client.get_environments_names_to_details(
+        deploy_config.environments
+    )
+
+    if kube_deployment_type == KubeDeploymentType.ONLINE:
+        env_to_serving_strategy = create_realtime_serving_strategy_from_deploy_config(
+            deploy_config,
+            environment_name_to_config,
+        )
+        environment_variables = dictify_params(deploy_config.realtime.env_vars)
+    elif kube_deployment_type == KubeDeploymentType.BATCH:
+        env_to_serving_strategy = {}
+        if len(environment_name_to_config) == 0:
+            env_to_serving_strategy = {
+                ecosystem_client.get_account_details().default_environment_id: ServingStrategy(
+                    batch_config=BatchConfig()
+                )
+            }
+        else:
+            for env_name, env_config in environment_name_to_config.items():
+                env_to_serving_strategy[env_config.id] = ServingStrategy(
+                    batch_config=BatchConfig()
+                )
+
+    elif kube_deployment_type == KubeDeploymentType.STREAM:
+        env_to_serving_strategy = {}
+        if len(environment_name_to_config) == 0:
+            env_to_serving_strategy[
+                ecosystem_client.get_account_details().default_environment_id
+            ] = create_streaming_serving_strategy_from_deploy_config(deploy_config)
+        else:
+            for env_name, env_config in environment_name_to_config.items():
+                env_to_serving_strategy[
+                    env_config.id
+                ] = create_streaming_serving_strategy_from_deploy_config(deploy_config)
+        environment_variables = dictify_params(deploy_config.stream.env_vars)
+    else:
+        raise QwakException(UNKNOWN_SERVING_STRATEGY)
+    env_to_hosting_services = {
+        env_id: HostingService(
+            kube_deployment=KubeDeployment(
+                deployment_size=deployment_size,
+                advanced_deployment_options=advanced_deployment_options,
+                serving_strategy=serving_strategy,
+                kube_deployment_type=kube_deployment_type,
+                environment_variables=environment_variables,
+            ),
+        )
+        for env_id, serving_strategy in env_to_serving_strategy.items()
+    }
+    env_deployment_messages = {
+        env_id: EnvironmentDeploymentMessage(
+            model_id=deploy_config.model_id,
+            build_id=deploy_config.build_id,
+            hosting_service=hosting_service,
+        )
+        for env_id, hosting_service in env_to_hosting_services.items()
+    }
+    return env_deployment_messages

qwak_sdk/commands/models/deployments/deploy/_logic/deployment_response_handler.py

@@ -0,0 +1,156 @@
+from typing import Dict
+
+from _qwak_proto.qwak.deployment.deployment_pb2 import ModelDeploymentStatus
+from _qwak_proto.qwak.deployment.deployment_service_pb2 import DeployModelResponse
+from _qwak_proto.qwak.ecosystem.v0.ecosystem_pb2 import UserContextEnvironmentDetails
+from qwak.clients.administration.eco_system.client import EcosystemClient
+from qwak.exceptions import QwakException
+from qwak.tools.logger.logger import get_qwak_logger
+from tabulate import tabulate
+
+from qwak_sdk.commands.models.deployments.deploy._logic.base_deploy_executor import (
+    FAILED_STATUS,
+    BaseDeployExecutor,
+)
+from qwak_sdk.tools.utils import qwak_spinner
+
+logger = get_qwak_logger()
+
+
+def __client_deployment_remote(
+    deploy: BaseDeployExecutor, sync: bool, output_multiple: bool = False
+):
+    try:
+        deploy_response: DeployModelResponse = deploy.deploy()
+        deploy_status, deploy_id = (
+            ModelDeploymentStatus.Name(deploy_response.status),
+            deploy_response.deployment_named_id,
+        )
+
+        if deploy_response.environment_to_deployment_result:
+            account_environments = EcosystemClient().get_environments()
+            output_deployments_response(account_environments, deploy, deploy_response)
+            __check_failed_deployments(account_environments, deploy_response)
+
+            if sync:
+                deployment_id_to_env = {}
+                for (
+                    env_id,
+                    env_deploy_response,
+                ) in deploy_response.environment_to_deployment_result.items():
+                    deployment_id_to_env[
+                        env_deploy_response.deployment_named_id
+                    ] = account_environments.get(env_id).name
+
+                __sync_deployments(deploy, deployment_id_to_env)
+
+            if output_multiple:
+                return deploy_response.environment_to_deployment_result
+        else:
+            if "FAILED" in deploy_status:
+                raise QwakException(deploy_response.info)
+
+            if sync:
+                __sync_deployment(deploy, deploy_id)
+
+        if not sync:
+            logger.info(
+                "Deployment initiated successfully, Use --sync to wait for deployments to be ready."
+            )
+            return deploy_id, deploy_status
+    except Exception as e:
+        logger.error(f"Deployment failed, Error: {e}")
+        raise e
+
+
+def __check_failed_deployments(
+    account_environments: Dict[str, UserContextEnvironmentDetails],
+    deploy_response: DeployModelResponse,
+):
+    failed_deployments = {}
+    for (
+        env_id,
+        env_deploy_response,
+    ) in deploy_response.environment_to_deployment_result.items():
+        deploy_status = ModelDeploymentStatus.Name(env_deploy_response.status)
+        if "FAILED" in deploy_status:
+            env_config = account_environments.get(env_id)
+            failed_deployments[
+                env_config.name if env_config else env_id
+            ] = env_deploy_response.info
+    if failed_deployments:
+        failed_deployments_message = "\n".join(
+            [
+                f"Environment {env_name}: {reason}"
+                for env_name, reason in failed_deployments.items()
+            ]
+        )
+        raise QwakException(
+            f"The following deployments failed:\n {failed_deployments_message}"
+        )
+
+
+def output_deployments_response(
+    account_environments: Dict[str, UserContextEnvironmentDetails],
+    deploy: BaseDeployExecutor,
+    deploy_response: DeployModelResponse,
+):
+    for (
+        env_id,
+        env_deploy_response,
+    ) in deploy_response.environment_to_deployment_result.items():
+        env_config = account_environments.get(env_id)
+        if env_deploy_response.status == ModelDeploymentStatus.INITIATING_DEPLOYMENT:
+            logger.info(
+                tabulate(
+                    tabular_data=[
+                        ["Environment", env_config.name],
+                        ["Model ID", deploy.config.model_id],
+                        ["Build ID", deploy.config.build_id],
+                        ["Deployment ID", env_deploy_response.deployment_named_id],
+                    ],
+                    tablefmt="fancy_grid",
+                )
+            )
+
+
+def __sync_deployments(
+    deploy: BaseDeployExecutor, deployment_id_to_env: Dict[str, str]
+):
+    with qwak_spinner(
+        begin_text=f"Deploy - model: {deploy.config.model_id}, build: {deploy.config.build_id} "
+        f"Environments: {', '.join(deploy.config.realtime.environments)}",
+        end_text="Successful deployments",
+        print_callback=print,
+    ):
+        deployment_id_to_status = deploy.poll_until_complete_multiple(
+            list(deployment_id_to_env.keys()), check_every_n_seconds=5
+        )
+        failed_deployments = []
+        for deployment_id, status in deployment_id_to_status.items():
+            if status in FAILED_STATUS:
+                failed_deployments.append(
+                    f"Environment {deployment_id_to_env.get(deployment_id)} "
+                    f"failed with status: {status}"
+                )
+        if failed_deployments:
+            raise QwakException("\n".join(failed_deployments))
+
+
+def __sync_deployment(deploy: BaseDeployExecutor, deploy_name: str):
+    with qwak_spinner(
+        begin_text=f"Deploy - model: {deploy.config.model_id}, build: {deploy.config.build_id}",
+        end_text="Successful deployments",
+        print_callback=print,
+    ):
+        deployment_status = deploy.poll_until_complete(
+            deploy_name, check_every_n_seconds=5
+        )
+        if deployment_status in FAILED_STATUS:
+            raise QwakException(deployment_status)
+
+
+def client_deployment(
+    deploy: BaseDeployExecutor, sync: bool, output_multiple: bool = False
+):
+    return __client_deployment_remote(deploy, sync, output_multiple)