client.py

import inspect  # pylint: disable=C0302
import json
import logging
import os
import shutil
import tempfile
from io import StringIO
from typing import (
    Any,
    Callable,
    Dict,
    Iterable,
    List,
    Optional,
    Type,
    TypeVar,
    Union,
)
from zipfile import ZipFile

import cloudpickle
import requests
import sseclient
import yaml
from deprecation import deprecated
from frozendict import frozendict
from pydantic import BaseModel
from typing_extensions import Literal

from launch.api_client import ApiClient, Configuration
from launch.api_client.apis.tags.default_api import DefaultApi
from launch.api_client.model.body_upload_file_v1_files_post import (
    BodyUploadFileV1FilesPost,
)
from launch.api_client.model.callback_auth import CallbackAuth
from launch.api_client.model.clone_model_bundle_v1_request import (
    CloneModelBundleV1Request,
)
from launch.api_client.model.clone_model_bundle_v2_request import (
    CloneModelBundleV2Request,
)
from launch.api_client.model.cloudpickle_artifact_flavor import (
    CloudpickleArtifactFlavor,
)
from launch.api_client.model.completion_stream_v1_response import (
    CompletionStreamV1Response,
)
from launch.api_client.model.completion_sync_v1_request import (
    CompletionSyncV1Request,
)
from launch.api_client.model.completion_sync_v1_response import (
    CompletionSyncV1Response,
)
from launch.api_client.model.create_batch_job_v1_request import (
    CreateBatchJobV1Request,
)
from launch.api_client.model.create_docker_image_batch_job_bundle_v1_request import (
    CreateDockerImageBatchJobBundleV1Request,
)
from launch.api_client.model.create_docker_image_batch_job_v1_request import (
    CreateDockerImageBatchJobV1Request,
)
from launch.api_client.model.create_fine_tune_request import (
    CreateFineTuneRequest,
)
from launch.api_client.model.create_llm_model_endpoint_v1_request import (
    CreateLLMModelEndpointV1Request,
)
from launch.api_client.model.create_model_bundle_v1_request import (
    CreateModelBundleV1Request,
)
from launch.api_client.model.create_model_bundle_v2_request import (
    CreateModelBundleV2Request,
)
from launch.api_client.model.create_model_endpoint_v1_request import (
    CreateModelEndpointV1Request,
)
from launch.api_client.model.custom_framework import CustomFramework
from launch.api_client.model.endpoint_predict_v1_request import (
    EndpointPredictV1Request,
)
from launch.api_client.model.gpu_type import GpuType
from launch.api_client.model.llm_inference_framework import (
    LLMInferenceFramework,
)
from launch.api_client.model.llm_source import LLMSource
from launch.api_client.model.model_bundle_environment_params import (
    ModelBundleEnvironmentParams,
)
from launch.api_client.model.model_bundle_framework_type import (
    ModelBundleFrameworkType,
)
from launch.api_client.model.model_bundle_packaging_type import (
    ModelBundlePackagingType,
)
from launch.api_client.model.model_endpoint_type import ModelEndpointType
from launch.api_client.model.pytorch_framework import PytorchFramework
from launch.api_client.model.quantization import Quantization
from launch.api_client.model.runnable_image_flavor import RunnableImageFlavor
from launch.api_client.model.streaming_enhanced_runnable_image_flavor import (
    StreamingEnhancedRunnableImageFlavor,
)
from launch.api_client.model.tensorflow_framework import TensorflowFramework
from launch.api_client.model.triton_enhanced_runnable_image_flavor import (
    TritonEnhancedRunnableImageFlavor,
)
from launch.api_client.model.update_docker_image_batch_job_v1_request import (
    UpdateDockerImageBatchJobV1Request,
)
from launch.api_client.model.update_model_endpoint_v1_request import (
    UpdateModelEndpointV1Request,
)
from launch.api_client.model.zip_artifact_flavor import ZipArtifactFlavor
from launch.connection import Connection
from launch.constants import (
    BATCH_TASK_INPUT_SIGNED_URL_PATH,
    DEFAULT_SCALE_ENDPOINT,
    ENDPOINT_PATH,
    MODEL_BUNDLE_SIGNED_URL_PATH,
    SCALE_LAUNCH_V0_PATH,
    SCALE_LAUNCH_V1_PATH,
)
from launch.docker_image_batch_job_bundle import (
    CreateDockerImageBatchJobBundleResponse,
    DockerImageBatchJobBundleResponse,
    ListDockerImageBatchJobBundleResponse,
)
from launch.file import (
    DeleteFileResponse,
    GetFileContentResponse,
    GetFileResponse,
    ListFilesResponse,
    UploadFileResponse,
)
from launch.find_packages import find_packages_from_imports, get_imports
from launch.fine_tune import (
    CancelFineTuneResponse,
    CreateFineTuneResponse,
    GetFineTuneEventsResponse,
    GetFineTuneResponse,
    ListFineTunesResponse,
)
from launch.hooks import PostInferenceHooks
from launch.make_batch_file import (
    make_batch_input_dict_file,
    make_batch_input_file,
)
from launch.model import ModelDownloadResponse
from launch.model_bundle import (
    CreateModelBundleV2Response,
    ListModelBundlesV2Response,
    ModelBundle,
    ModelBundleV2Response,
)
from launch.model_endpoint import (
    AsyncEndpoint,
    Endpoint,
    ModelEndpoint,
    StreamingEndpoint,
    SyncEndpoint,
)
from launch.pydantic_schemas import get_model_definitions
from launch.request_validation import validate_task_request

DEFAULT_NETWORK_TIMEOUT_SEC = 120
DEFAULT_LLM_COMPLETIONS_TIMEOUT = 300

logger = logging.getLogger(__name__)
logging.basicConfig()

LaunchModel_T = TypeVar("LaunchModel_T")


def _model_bundle_to_name(model_bundle: Union[ModelBundle, str]) -> str:
    if isinstance(model_bundle, ModelBundle):
        return model_bundle.name
    elif isinstance(model_bundle, str):
        return model_bundle
    else:
        raise TypeError("model_bundle should be type ModelBundle or str")


def _model_bundle_to_id(model_bundle: Union[ModelBundle, str]) -> str:
    if isinstance(model_bundle, ModelBundle):
        if model_bundle.id is None:
            raise ValueError(
                "You need to pass in a ModelBundle that has an id, "
                "i.e. one that has already been registered on the server"
            )
        return model_bundle.id
    elif isinstance(model_bundle, str):
        return model_bundle
    else:
        raise TypeError("model_bundle should be type ModelBundle or str")


def _model_endpoint_to_name(model_endpoint: Union[ModelEndpoint, str]) -> str:
    if isinstance(model_endpoint, ModelEndpoint):
        return model_endpoint.name
    elif isinstance(model_endpoint, str):
        return model_endpoint
    else:
        raise TypeError("model_endpoint should be type ModelEndpoint or str")


def _add_app_config_to_bundle_create_payload(payload: Dict[str, Any], app_config: Optional[Union[Dict[str, Any], str]]):
    """
    Edits a request payload (for creating a bundle) to include a (not serialized) app_config if it's
    not None
    """
    if isinstance(app_config, Dict):
        payload["app_config"] = app_config
    elif isinstance(app_config, str):
        with open(app_config, "r") as f:  # pylint: disable=unspecified-encoding
            app_config_dict = yaml.safe_load(f)
            payload["app_config"] = app_config_dict


def _get_model_bundle_framework(
    pytorch_image_tag: Optional[str] = None,
    tensorflow_version: Optional[str] = None,
    custom_base_image_repository: Optional[str] = None,
    custom_base_image_tag: Optional[str] = None,
):
    if pytorch_image_tag is not None:
        return PytorchFramework(
            pytorch_image_tag=pytorch_image_tag,
            framework_type=ModelBundleFrameworkType.PYTORCH,
        )
    elif tensorflow_version is not None:
        return TensorflowFramework(
            tensorflow_version=tensorflow_version,
            framework_type=ModelBundleFrameworkType.TENSORFLOW,
        )
    elif custom_base_image_repository is not None and custom_base_image_tag is not None:
        return CustomFramework(
            image_repository=custom_base_image_repository,
            image_tag=custom_base_image_tag,
            framework_type=ModelBundleFrameworkType.CUSTOM_BASE_IMAGE,
        )
    else:
        raise ValueError(
            "You must specify one of pytorch_image_tag, tensorflow_version, or "
            "custom_base_image_repository and custom_base_image_tag"
        )


def dict_not_none(**kwargs):
    return {k: v for k, v in kwargs.items() if v is not None}


class LaunchClient:
    """Scale Launch Python Client."""

    def __init__(
        self,
        api_key: str,
        endpoint: Optional[str] = None,
        self_hosted: bool = False,
        use_path_with_custom_endpoint: bool = False,
    ):
        """
        Initializes a Scale Launch Client.

        Parameters:
            api_key: Your Scale API key
            endpoint: The Scale Launch Endpoint (this should not need to be changed)
            self_hosted: True iff you are connecting to a self-hosted Scale Launch
            use_path_with_custom_endpoint: True iff you are not using the default Scale Launch endpoint
                but your endpoint has path routing (to SCALE_LAUNCH_VX_PATH) set up
        """
        self.endpoint = endpoint or DEFAULT_SCALE_ENDPOINT
        self.connection = Connection(api_key, self.endpoint + SCALE_LAUNCH_V0_PATH)
        self.self_hosted = self_hosted
        self.upload_bundle_fn: Optional[Callable[[str, str], None]] = None
        self.upload_batch_csv_fn: Optional[Callable[[str, str], None]] = None
        self.bundle_location_fn: Optional[Callable[[], str]] = None
        self.batch_csv_location_fn: Optional[Callable[[], str]] = None
        host = self.endpoint + SCALE_LAUNCH_V1_PATH if endpoint is None else self.endpoint
        if use_path_with_custom_endpoint:
            host = self.endpoint + SCALE_LAUNCH_V1_PATH
        self.configuration = Configuration(
            host=host,
            discard_unknown_keys=True,
            username=api_key,
            password="",
        )

    def __repr__(self):
        return f"LaunchClient(connection='{self.connection}')"

    def __eq__(self, other):
        return self.connection == other.connection

    def register_upload_bundle_fn(self, upload_bundle_fn: Callable[[str, str], None]):
        """
        For self-hosted mode only. Registers a function that handles model bundle upload. This
        function is called as

            upload_bundle_fn(serialized_bundle, bundle_url)

        This function should directly write the contents of ``serialized_bundle`` as a
        binary string into ``bundle_url``.

        See ``register_bundle_location_fn`` for more notes on the signature of ``upload_bundle_fn``

        Parameters:
            upload_bundle_fn: Function that takes in a serialized bundle (bytes type),
                and uploads that bundle to an appropriate location. Only needed for self-hosted mode.
        """
        self.upload_bundle_fn = upload_bundle_fn

    def register_upload_batch_csv_fn(self, upload_batch_csv_fn: Callable[[str, str], None]):
        """
        For self-hosted mode only. Registers a function that handles batch text upload. This
        function is called as

            upload_batch_csv_fn(csv_text, csv_url)

        This function should directly write the contents of ``csv_text`` as a text string into
        ``csv_url``.

        Parameters:
            upload_batch_csv_fn: Function that takes in a csv text (string type),
                and uploads that bundle to an appropriate location. Only needed for self-hosted mode.
        """
        self.upload_batch_csv_fn = upload_batch_csv_fn

    def register_bundle_location_fn(self, bundle_location_fn: Callable[[], str]):
        """
        For self-hosted mode only. Registers a function that gives a location for a model bundle.
        Should give different locations each time. This function is called as
        ``bundle_location_fn()``, and should return a ``bundle_url`` that
        ``register_upload_bundle_fn`` can take.

        Strictly, ``bundle_location_fn()`` does not need to return a ``str``. The only
        requirement is that if ``bundle_location_fn`` returns a value of type ``T``,
        then ``upload_bundle_fn()`` takes in an object of type T as its second argument (i.e.
        bundle_url).

        Parameters:
            bundle_location_fn: Function that generates bundle_urls for upload_bundle_fn.
        """
        self.bundle_location_fn = bundle_location_fn

    def register_batch_csv_location_fn(self, batch_csv_location_fn: Callable[[], str]):
        """
        For self-hosted mode only. Registers a function that gives a location for batch CSV
        inputs. Should give different locations each time. This function is called as
        batch_csv_location_fn(), and should return a batch_csv_url that upload_batch_csv_fn can
        take.

        Strictly, batch_csv_location_fn() does not need to return a str. The only requirement is
        that if batch_csv_location_fn returns a value of type T, then upload_batch_csv_fn() takes
        in an object of type T as its second argument (i.e. batch_csv_url).

        Parameters:
            batch_csv_location_fn: Function that generates batch_csv_urls for upload_batch_csv_fn.
        """
        self.batch_csv_location_fn = batch_csv_location_fn

    def _upload_data(self, data: bytes) -> str:
        if self.self_hosted:
            if self.upload_bundle_fn is None:
                raise ValueError("Upload_bundle_fn should be registered")
            if self.bundle_location_fn is None:
                raise ValueError("Need either bundle_location_fn to know where to upload bundles")
            raw_bundle_url = self.bundle_location_fn()  # type: ignore
            self.upload_bundle_fn(data, raw_bundle_url)  # type: ignore
        else:
            model_bundle_url = self.connection.post({}, MODEL_BUNDLE_SIGNED_URL_PATH)
            s3_path = model_bundle_url["signedUrl"]
            raw_bundle_url = f"s3://{model_bundle_url['bucket']}/{model_bundle_url['key']}"
            requests.put(s3_path, data=data)
        return raw_bundle_url

    def _get_bundle_url_from_base_paths(self, base_paths: List[str]) -> str:
        tmpdir = tempfile.mkdtemp()
        try:
            zip_path = os.path.join(tmpdir, "bundle.zip")
            _zip_directories(zip_path, base_paths)
            with open(zip_path, "rb") as zip_f:
                data = zip_f.read()
        finally:
            shutil.rmtree(tmpdir)

        raw_bundle_url = self._upload_data(data)
        return raw_bundle_url

    def _upload_model_bundle(
        self,
        load_model_fn: Callable,
        load_predict_fn: Callable,
    ):
        bundle = dict(load_model_fn=load_model_fn, load_predict_fn=load_predict_fn)
        serialized_bundle = cloudpickle.dumps(bundle)
        bundle_location = self._upload_data(data=serialized_bundle)
        return bundle_location

    def _upload_schemas(self, request_schema: Type[BaseModel], response_schema: Type[BaseModel]) -> str:
        model_definitions = get_model_definitions(
            request_schema=request_schema,
            response_schema=response_schema,
        )
        model_definitions_encoded = json.dumps(model_definitions).encode()
        return self._upload_data(model_definitions_encoded)

    def create_model_bundle_from_callable_v2(
        self,
        *,
        model_bundle_name: str,
        load_predict_fn: Callable[[LaunchModel_T], Callable[[Any], Any]],
        load_model_fn: Callable[[], LaunchModel_T],
        request_schema: Type[BaseModel],
        response_schema: Type[BaseModel],
        requirements: Optional[List[str]] = None,
        pytorch_image_tag: Optional[str] = None,
        tensorflow_version: Optional[str] = None,
        custom_base_image_repository: Optional[str] = None,
        custom_base_image_tag: Optional[str] = None,
        app_config: Optional[Union[Dict[str, Any], str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> CreateModelBundleV2Response:
        """
        Uploads and registers a model bundle to Scale Launch.

        Parameters:
            model_bundle_name: Name of the model bundle.

            load_predict_fn: Function that takes in a model and returns a predict function.
                When your model bundle is deployed, this predict function will be called as follows:
                ```
                input = {"input": "some input"} # or whatever your request schema is.

                def load_model_fn():
                    # load model
                    return model

                def load_predict_fn(model, app_config=None):
                    def predict_fn(input):
                        # do pre-processing
                        output = model(input)
                        # do post-processing
                        return output
                    return predict_fn

                predict_fn = load_predict_fn(load_model_fn(), app_config=optional_app_config)
                response = predict_fn(input)
                ```

            load_model_fn: A function that, when run, loads a model.

            request_schema: A pydantic model that represents the request schema for the model
                bundle. This is used to validate the request body for the model bundle's endpoint.

            response_schema: A pydantic model that represents the request schema for the model
                bundle. This is used to validate the response for the model bundle's endpoint.

            requirements: List of pip requirements.

            pytorch_image_tag: The image tag for the PyTorch image that will be used to run the
                bundle. Exactly one of ``pytorch_image_tag``, ``tensorflow_version``, or
                ``custom_base_image_repository`` must be specified.

            tensorflow_version: The version of TensorFlow that will be used to run the bundle.
                If not specified, the default version will be used. Exactly one of
                ``pytorch_image_tag``, ``tensorflow_version``, or ``custom_base_image_repository``
                must be specified.

            custom_base_image_repository: The repository for a custom base image that will be
                used to run the bundle. If not specified, the default base image will be used.
                Exactly one of ``pytorch_image_tag``, ``tensorflow_version``, or
                ``custom_base_image_repository`` must be specified.

            custom_base_image_tag: The tag for a custom base image that will be used to run the
                bundle. Must be specified if ``custom_base_image_repository`` is specified.

            app_config: An optional dictionary of configuration values that will be passed to the
                bundle when it is run. These values can be accessed by the bundle via the
                ``app_config`` global variable.

            metadata: Metadata to record with the bundle.

        Returns:
            An object containing the following keys:

                - ``model_bundle_id``: The ID of the created model bundle.
        """
        nonnull_requirements = requirements or []
        bundle_location = self._upload_model_bundle(load_model_fn, load_predict_fn)
        schema_location = self._upload_schemas(request_schema=request_schema, response_schema=response_schema)
        framework = _get_model_bundle_framework(
            pytorch_image_tag=pytorch_image_tag,
            tensorflow_version=tensorflow_version,
            custom_base_image_repository=custom_base_image_repository,
            custom_base_image_tag=custom_base_image_tag,
        )
        flavor = CloudpickleArtifactFlavor(
            **dict_not_none(
                flavor="cloudpickle_artifact",
                load_predict_fn=inspect.getsource(load_predict_fn),
                load_model_fn=inspect.getsource(load_model_fn),
                framework=framework,
                requirements=nonnull_requirements,
                app_config=app_config,
                location=bundle_location,
            )
        )
        create_model_bundle_request = CreateModelBundleV2Request(
            **dict_not_none(
                name=model_bundle_name,
                schema_location=schema_location,
                flavor=flavor,
                metadata=metadata,
            )
        )
        with ApiClient(self.configuration) as api_client:
            api_instance = DefaultApi(api_client)
            response = api_instance.create_model_bundle_v2_model_bundles_post(
                body=create_model_bundle_request,
                skip_deserialization=True,
            )
            resp = CreateModelBundleV2Response.parse_raw(response.response.data)

        return resp

    def create_model_bundle_from_dirs_v2(
        self,
        *,
        model_bundle_name: str,
        base_paths: List[str],
        load_predict_fn_module_path: str,
        load_model_fn_module_path: str,
        request_schema: Type[BaseModel],
        response_schema: Type[BaseModel],
        requirements_path: Optional[str] = None,
        pytorch_image_tag: Optional[str] = None,
        tensorflow_version: Optional[str] = None,
        custom_base_image_repository: Optional[str] = None,
        custom_base_image_tag: Optional[str] = None,
        app_config: Optional[Dict[str, Any]] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> CreateModelBundleV2Response:
        """
        Packages up code from one or more local filesystem folders and uploads them as a bundle
        to Scale Launch. In this mode, a bundle is just local code instead of a serialized object.

        For example, if you have a directory structure like so, and your current working
        directory is ``my_root``:

        ```text
           my_root/
               my_module1/
                   __init__.py
                   ...files and directories
                   my_inference_file.py
               my_module2/
                   __init__.py
                   ...files and directories
        ```

        then calling ``create_model_bundle_from_dirs_v2`` with ``base_paths=["my_module1",
        "my_module2"]`` essentially creates a zip file without the root directory, e.g.:

        ```text
           my_module1/
               __init__.py
               ...files and directories
               my_inference_file.py
           my_module2/
               __init__.py
               ...files and directories
        ```

        and these contents will be unzipped relative to the server side application root. Bear
        these points in mind when referencing Python module paths for this bundle. For instance,
        if ``my_inference_file.py`` has ``def f(...)`` as the desired inference loading function,
        then the `load_predict_fn_module_path` argument should be `my_module1.my_inference_file.f`.

        Parameters:
            model_bundle_name: The name of the model bundle you want to create.

            base_paths: A list of paths to directories that will be zipped up and uploaded
                as a bundle. Each path must be relative to the current working directory.

            load_predict_fn_module_path: The Python module path to the function that will be
                used to load the model for inference. This function should take in a path to a
                model directory, and return a model object. The model object should be pickleable.

            load_model_fn_module_path: The Python module path to the function that will be
                used to load the model for training. This function should take in a path to a
                model directory, and return a model object. The model object should be pickleable.

            request_schema: A Pydantic model that defines the request schema for the bundle.

            response_schema: A Pydantic model that defines the response schema for the bundle.

            requirements_path: Path to a requirements.txt file that will be used to install
                dependencies for the bundle. This file must be relative to the current working
                directory.

            pytorch_image_tag: The image tag for the PyTorch image that will be used to run the
                bundle. Exactly one of ``pytorch_image_tag``, ``tensorflow_version``, or
                ``custom_base_image_repository`` must be specified.

            tensorflow_version: The version of TensorFlow that will be used to run the bundle.
                If not specified, the default version will be used. Exactly one of
                ``pytorch_image_tag``, ``tensorflow_version``, or ``custom_base_image_repository``
                must be specified.

            custom_base_image_repository: The repository for a custom base image that will be
                used to run the bundle. If not specified, the default base image will be used.
                Exactly one of ``pytorch_image_tag``, ``tensorflow_version``, or
                ``custom_base_image_repository`` must be specified.

            custom_base_image_tag: The tag for a custom base image that will be used to run the
                bundle. Must be specified if ``custom_base_image_repository`` is specified.

            app_config: An optional dictionary of configuration values that will be passed to the
                bundle when it is run. These values can be accessed by the bundle via the
                ``app_config`` global variable.

            metadata: Metadata to record with the bundle.

        Returns:
            An object containing the following keys:

                - ``model_bundle_id``: The ID of the created model bundle.
        """
        requirements = []
        if requirements_path is not None:
            with open(requirements_path, "r", encoding="utf-8") as req_f:
                requirements = req_f.read().splitlines()
        bundle_location = self._get_bundle_url_from_base_paths(base_paths)
        schema_location = self._upload_schemas(request_schema=request_schema, response_schema=response_schema)
        framework = _get_model_bundle_framework(
            pytorch_image_tag=pytorch_image_tag,
            tensorflow_version=tensorflow_version,
            custom_base_image_repository=custom_base_image_repository,
            custom_base_image_tag=custom_base_image_tag,
        )
        flavor = ZipArtifactFlavor(
            **dict_not_none(
                flavor="zip_artifact",
                load_predict_fn_module_path=load_predict_fn_module_path,
                load_model_fn_module_path=load_model_fn_module_path,
                framework=framework,
                requirements=requirements,
                app_config=app_config,
                location=bundle_location,
            )
        )
        create_model_bundle_request = CreateModelBundleV2Request(
            **dict_not_none(
                name=model_bundle_name,
                schema_location=schema_location,
                flavor=flavor,
                metadata=metadata,
            )
        )
        with ApiClient(self.configuration) as api_client:
            api_instance = DefaultApi(api_client)
            response = api_instance.create_model_bundle_v2_model_bundles_post(
                body=create_model_bundle_request,
                skip_deserialization=True,
            )
            resp = CreateModelBundleV2Response.parse_raw(response.response.data)

        return resp

    def create_model_bundle_from_runnable_image_v2(
        self,
        *,
        model_bundle_name: str,
        request_schema: Type[BaseModel],
        response_schema: Type[BaseModel],
        repository: str,
        tag: str,
        command: List[str],
        healthcheck_route: Optional[str] = None,
        predict_route: Optional[str] = None,
        env: Dict[str, str],
        readiness_initial_delay_seconds: int,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> CreateModelBundleV2Response:
        """
        Create a model bundle from a runnable image. The specified ``command`` must start a process
        that will listen for requests on port 5005 using HTTP.

        Inference requests must be served at the `POST /predict` route while the `GET /readyz` route is a healthcheck.

        Parameters:
            model_bundle_name: The name of the model bundle you want to create.

            request_schema: A Pydantic model that defines the request schema for the bundle.

            response_schema: A Pydantic model that defines the response schema for the bundle.

            repository: The name of the Docker repository for the runnable image.

            tag: The tag for the runnable image.

            command: The command that will be used to start the process that listens for requests.

            predict_route: The endpoint route on the runnable image that will be called.

            healthcheck_route: The healthcheck endpoint route on the runnable image.

            env: A dictionary of environment variables that will be passed to the bundle when it
                is run.

            readiness_initial_delay_seconds: The number of seconds to wait for the HTTP server to become ready and
                successfully respond on its healthcheck.

            metadata: Metadata to record with the bundle.

        Returns:
            An object containing the following keys:

                - ``model_bundle_id``: The ID of the created model bundle.
        """
        schema_location = self._upload_schemas(request_schema=request_schema, response_schema=response_schema)
        flavor = RunnableImageFlavor(
            **dict_not_none(
                flavor="runnable_image",
                repository=repository,
                tag=tag,
                command=command,
                healthcheck_route=healthcheck_route,
                predict_route=predict_route,
                env=env,
                protocol="http",
                readiness_initial_delay_seconds=readiness_initial_delay_seconds,
            )
        )
        create_model_bundle_request = CreateModelBundleV2Request(
            **dict_not_none(
                name=model_bundle_name,
                schema_location=schema_location,
                flavor=flavor,
                metadata=metadata,
            )
        )

        with ApiClient(self.configuration) as api_client:
            api_instance = DefaultApi(api_client)
            response = api_instance.create_model_bundle_v2_model_bundles_post(
                body=create_model_bundle_request,
                skip_deserialization=True,
            )
            resp = CreateModelBundleV2Response.parse_raw(response.response.data)

        return resp

    def create_model_bundle_from_streaming_enhanced_runnable_image_v2(
        self,
        *,
        model_bundle_name: str,
        request_schema: Type[BaseModel],
        response_schema: Type[BaseModel],
        repository: str,
        tag: str,
        command: Optional[List[str]] = None,
        healthcheck_route: Optional[str] = None,
        predict_route: Optional[str] = None,
        streaming_command: List[str],
        streaming_predict_route: Optional[str] = None,
        env: Dict[str, str],
        readiness_initial_delay_seconds: int,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> CreateModelBundleV2Response:
        """
        Create a model bundle from a runnable image. The specified ``command`` must start a process
        that will listen for requests on port 5005 using HTTP.

        Inference requests must be served at the `POST /predict` route while the `GET /readyz` route is a healthcheck.

        Parameters:
            model_bundle_name: The name of the model bundle you want to create.

            request_schema: A Pydantic model that defines the request schema for the bundle.

            response_schema: A Pydantic model that defines the response schema for the bundle.

            repository: The name of the Docker repository for the runnable image.

            tag: The tag for the runnable image.

            command: The command that will be used to start the process that listens for requests if
                this bundle is used as a SYNC or ASYNC endpoint.

            healthcheck_route: The healthcheck endpoint route on the runnable image.

            predict_route: The endpoint route on the runnable image that will be called if this bundle is used as a SYNC
                or ASYNC endpoint.

            streaming_command: The command that will be used to start the process that listens for
                requests if this bundle is used as a STREAMING endpoint.

            streaming_predict_route: The endpoint route on the runnable image that will be called if this bundle is used
                as a STREAMING endpoint.

            env: A dictionary of environment variables that will be passed to the bundle when it
                is run.

            readiness_initial_delay_seconds: The number of seconds to wait for the HTTP server to become ready and
                successfully respond on its healthcheck.

            metadata: Metadata to record with the bundle.

        Returns:
            An object containing the following keys:

                - ``model_bundle_id``: The ID of the created model bundle.
        """
        schema_location = self._upload_schemas(request_schema=request_schema, response_schema=response_schema)
        flavor = StreamingEnhancedRunnableImageFlavor(
            **dict_not_none(
                flavor="streaming_enhanced_runnable_image",
                repository=repository,
                tag=tag,
                command=command,
                healthcheck_route=healthcheck_route,
                predict_route=predict_route,
                streaming_command=streaming_command,
                streaming_predict_route=streaming_predict_route,
                env=env,
                protocol="http",
                readiness_initial_delay_seconds=readiness_initial_delay_seconds,
            )
        )
        create_model_bundle_request = CreateModelBundleV2Request(
            **dict_not_none(
                name=model_bundle_name,
                schema_location=schema_location,
                flavor=flavor,
                metadata=metadata,
            )
        )

        with ApiClient(self.configuration) as api_client:
            api_instance = DefaultApi(api_client)
            response = api_instance.create_model_bundle_v2_model_bundles_post(
                body=create_model_bundle_request,
                skip_deserialization=True,
            )
            resp = CreateModelBundleV2Response.parse_raw(response.response.data)

        return resp

    def create_model_bundle_from_triton_enhanced_runnable_image_v2(
        self,
        *,
        model_bundle_name: str,
        request_schema: Type[BaseModel],
        response_schema: Type[BaseModel],
        repository: str,
        tag: str,
        command: List[str],
        healthcheck_route: Optional[str] = None,
        predict_route: Optional[str] = None,
        env: Dict[str, str],
        readiness_initial_delay_seconds: int,
        triton_model_repository: str,
        triton_model_replicas: Optional[Dict[str, str]] = None,
        triton_num_cpu: float,
        triton_commit_tag: str,
        triton_storage: Optional[str] = None,
        triton_memory: Optional[str] = None,
        triton_readiness_initial_delay_seconds: int,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> CreateModelBundleV2Response:
        """
        Create a model bundle from a runnable image and a tritonserver image.

        Same requirements as :param:`create_model_bundle_from_runnable_image_v2` with additional constraints necessary
        for configuring tritonserver's execution.

        Parameters:
            model_bundle_name: The name of the model bundle you want to create.

            request_schema: A Pydantic model that defines the request schema for the bundle.

            response_schema: A Pydantic model that defines the response schema for the bundle.

            repository: The name of the Docker repository for the runnable image.

            tag: The tag for the runnable image.

            command: The command that will be used to start the process that listens for requests.

            predict_route: The endpoint route on the runnable image that will be called.

            healthcheck_route: The healthcheck endpoint route on the runnable image.

            env: A dictionary of environment variables that will be passed to the bundle when it
                is run.

            readiness_initial_delay_seconds: The number of seconds to wait for the HTTP server to
                become ready and successfully respond on its healthcheck.

            triton_model_repository: The S3 prefix that contains the contents of the model
                repository, formatted according to
                https://github.com/triton-inference-server/server/blob/main/docs/user_guide/model_repository.md

            triton_model_replicas: If supplied, the name and number of replicas to make for each
                model.

            triton_num_cpu: Number of CPUs, fractional, to allocate to tritonserver.

            triton_commit_tag: The image tag of the specific trionserver version.

            triton_storage: Amount of storage space to allocate for the tritonserver container.

            triton_memory: Amount of memory to allocate for the tritonserver container.

            triton_readiness_initial_delay_seconds: Like readiness_initial_delay_seconds, but for
                tritonserver's own healthcheck.

            metadata: Metadata to record with the bundle.

        Returns:
            An object containing the following keys:

                - ``model_bundle_id``: The ID of the created model bundle.
        """
        schema_location = self._upload_schemas(request_schema=request_schema, response_schema=response_schema)
        flavor = TritonEnhancedRunnableImageFlavor(
            **dict_not_none(
                flavor="triton_enhanced_runnable_image",
                repository=repository,
                tag=tag,
                command=command,
                healthcheck_route=healthcheck_route,
                predict_route=predict_route,
                env=env,
                protocol="http",
                readiness_initial_delay_seconds=readiness_initial_delay_seconds,
                triton_model_repository=triton_model_repository,
                triton_model_replicas=triton_model_replicas,
                triton_num_cpu=triton_num_cpu,
                triton_commit_tag=triton_commit_tag,
                triton_storage=triton_storage,
                triton_memory=triton_memory,
                triton_readiness_initial_delay_seconds=triton_readiness_initial_delay_seconds,
            )
        )
        create_model_bundle_request = CreateModelBundleV2Request(
            **dict_not_none(
                name=model_bundle_name,
                schema_location=schema_location,
                flavor=flavor,
                metadata=metadata,
            )
        )

        with ApiClient(self.configuration) as api_client:
            api_instance = DefaultApi(api_client)
            response = api_instance.create_model_bundle_v2_model_bundles_post(
                body=create_model_bundle_request,
                skip_deserialization=True,
            )
            resp = CreateModelBundleV2Response.parse_raw(response.response.data)

        return resp

    def get_model_bundle_v2(self, model_bundle_id: str) -> ModelBundleV2Response:
        """
        Get a model bundle.

        Parameters:
            model_bundle_id: The ID of the model bundle you want to get.

        Returns:
            An object containing the following fields:

                - ``id``: The ID of the model bundle.
                - ``name``: The name of the model bundle.
                - ``flavor``: The flavor of the model bundle. Either `RunnableImage`,
                    `CloudpickleArtifact`, `ZipArtifact`, or `TritonEnhancedRunnableImageFlavor`.
                - ``created_at``: The time the model bundle was created.
                - ``metadata``: A dictionary of metadata associated with the model bundle.
                - ``model_artifact_ids``: A list of IDs of model artifacts associated with the
                    bundle.
        """
        with ApiClient(self.configuration) as api_client:
            api_instance = DefaultApi(api_client)
            path_params = frozendict({"model_bundle_id": model_bundle_id})
            response = api_instance.get_model_bundle_v2_model_bundles_model_bundle_id_get(  # type: ignore
                path_params=path_params,
                skip_deserialization=True,
            )
            resp = ModelBundleV2Response.parse_raw(response.response.data)

        return resp

    def get_latest_model_bundle_v2(self, model_bundle_name: str) -> ModelBundleV2Response:
        """
        Get the latest version of a model bundle.