frogml-core 0.0.1__py3-none-any.whl

frogml/clients/batch_job_management/executions_config.py

@@ -0,0 +1,201 @@
+from collections import defaultdict
+from dataclasses import dataclass, field
+from typing import Dict, List, Tuple, Union
+
+from _frogml_proto.qwak.batch_job.v1.batch_job_resources_pb2 import (
+    NVIDIA_A10G,
+    NVIDIA_A100,
+    NVIDIA_K80,
+    NVIDIA_T4,
+    NVIDIA_V100,
+)
+from _frogml_proto.qwak.batch_job.v1.batch_job_service_pb2 import (
+    CSV_INPUT_FILE_TYPE,
+    CSV_OUTPUT_FILE_TYPE,
+    FEATHER_INPUT_FILE_TYPE,
+    FEATHER_OUTPUT_FILE_TYPE,
+    PARQUET_INPUT_FILE_TYPE,
+    PARQUET_OUTPUT_FILE_TYPE,
+)
+from frogml.exceptions import FrogmlException
+from frogml.inner.tool.run_config import (
+    ConfigCliMap,
+    FrogmlConfigBase,
+    YamlConfigMixin,
+    validate_float,
+    validate_int,
+    validate_list_of_strings,
+    validate_string,
+)
+
+DEFAULT_OUTPUT_FORMAT = "CSV"
+OUTPUT_FORMATTERS_MAP = {
+    DEFAULT_OUTPUT_FORMAT: CSV_OUTPUT_FILE_TYPE,
+    "PARQUET": PARQUET_OUTPUT_FILE_TYPE,
+    "FEATHER": FEATHER_OUTPUT_FILE_TYPE,
+}
+
+DEFAULT_INPUT_FORMAT = "CSV"
+INPUT_FORMATTERS_MAP = {
+    DEFAULT_INPUT_FORMAT: CSV_INPUT_FILE_TYPE,
+    "PARQUET": PARQUET_INPUT_FILE_TYPE,
+    "FEATHER": FEATHER_INPUT_FILE_TYPE,
+}
+
+GPU_TYPE_MAP = {
+    "NVIDIA_V100": NVIDIA_V100,
+    "NVIDIA_T4": NVIDIA_T4,
+    "NVIDIA_K80": NVIDIA_K80,
+    "NVIDIA_A100": NVIDIA_A100,
+    "NVIDIA_A10G": NVIDIA_A10G,
+}
+
+PURCHASE_OPTION_SET = (
+    "on-demand",
+    "spot",
+)
+
+
+CONFIG_MAPPING: List[ConfigCliMap] = [
+    ConfigCliMap("model_id", "execution.model_id", validate_string, True),
+    ConfigCliMap("build_id", "execution.build_id", validate_string, False),
+    ConfigCliMap("branch", "execution.branch", validate_string),
+    ConfigCliMap("source_bucket", "execution.source_bucket", validate_string, False),
+    ConfigCliMap("source_folder", "execution.source_folder", validate_string, False),
+    ConfigCliMap("bucket", "execution.bucket", validate_string, False),
+    ConfigCliMap(
+        "destination_bucket", "execution.destination_bucket", validate_string, False
+    ),
+    ConfigCliMap(
+        "destination_folder", "execution.destination_folder", validate_string, False
+    ),
+    ConfigCliMap(
+        "input_file_type",
+        "execution.input_file_type",
+        validate_string,
+        False,
+    ),
+    ConfigCliMap(
+        "output_file_type",
+        "execution.output_file_type",
+        validate_string,
+        False,
+    ),
+    ConfigCliMap("param_list", "execution.parameters", validate_list_of_strings, False),
+    ConfigCliMap("job_timeout", "execution.job_timeout", validate_int, False),
+    ConfigCliMap("file_timeout", "execution.file_timeout", validate_int, False),
+    ConfigCliMap(
+        "access_token_name", "execution.access_token_name", validate_string, False
+    ),
+    ConfigCliMap(
+        "access_secret_name", "execution.access_secret_name", validate_string, False
+    ),
+    ConfigCliMap(
+        "service_account_key_secret_name",
+        "execution.service_account_key_secret_name",
+        validate_string,
+        False,
+    ),
+    ConfigCliMap("pods", "resources.pods", validate_int, False),
+    ConfigCliMap("cpus", "resources.cpus", validate_float, False),
+    ConfigCliMap("memory", "resources.memory", validate_int, False),
+    ConfigCliMap("gpu_type", "resources.gpu_type", validate_string, False),
+    ConfigCliMap("gpu_amount", "resources.gpu_amount", validate_int, False),
+    ConfigCliMap("instance", "resources.instance_size", validate_string, False),
+    ConfigCliMap(
+        "purchase_option", "advanced_options.purchase_option", validate_string, False
+    ),
+]
+
+
+@dataclass
+class ExecutionConfig(YamlConfigMixin, FrogmlConfigBase):
+    def merge_bucket_arguments(self):
+        if self.execution.bucket:
+            if not self.execution.source_bucket:
+                self.execution.source_bucket = self.execution.bucket
+            if not self.execution.destination_bucket:
+                self.execution.destination_bucket = self.execution.bucket
+
+    def __post_init__(self):
+        self.merge_bucket_arguments()
+
+    def _post_merge_cli(self):
+        self.merge_bucket_arguments()
+        self.execution.parameters = dictify_params(self.execution.parameters)
+
+        if not self.execution.source_bucket:
+            raise FrogmlException(
+                "Must supply either --bucket or --source-bucket parameter"
+            )
+        if not self.execution.destination_bucket:
+            raise FrogmlException(
+                "Must supply either --bucket or --destination-bucket parameter"
+            )
+
+    @property
+    def _config_mapping(self) -> List[ConfigCliMap]:
+        return CONFIG_MAPPING
+
+    @dataclass
+    class Execution:
+        model_id: str = field(default="")
+        build_id: str = field(default="")
+        branch: str = field(default="main")
+        bucket: str = field(default="")
+        source_bucket: str = field(default="")
+        source_folder: str = field(default="")
+        destination_bucket: str = field(default="")
+        destination_folder: str = field(default="")
+        input_file_type: str = field(default=DEFAULT_INPUT_FORMAT)
+        output_file_type: str = field(default=DEFAULT_OUTPUT_FORMAT)
+        access_token_name: str = field(default="")
+        access_secret_name: str = field(default="")
+        job_timeout: int = field(default=0)
+        file_timeout: int = field(default=0)
+        parameters: dict = field(default_factory=lambda: defaultdict(str))
+        service_account_key_secret_name: str = field(default="")
+
+    @dataclass
+    class Warmup:
+        model_id: str = field(default="")
+        branch: str = field(default="main")
+        timeout: int = field(default=60 * 60)
+
+    @dataclass
+    class Resources:
+        pods: int = field(default=0)
+        cpus: float = field(default=0)
+        memory: int = field(default=0)
+        gpu_amount: int = field(default=0)
+        gpu_type: str = field(default="")
+        instance_size: str = field(default="")
+
+    @dataclass
+    class AdvancedOptions:
+        custom_iam_role_arn: str = field(default=None)
+        purchase_option: str = field(default=None)
+        service_account_key_secret_name: str = field(default=None)
+
+    warmup: Warmup = field(default_factory=Warmup)
+    execution: Execution = field(default_factory=Execution)
+    resources: Resources = field(default_factory=Resources)
+    advanced_options: AdvancedOptions = field(default_factory=AdvancedOptions)
+
+
+def dictify_params(parameters: Union[Dict[str, str], List[str], Tuple[str]]):
+    if isinstance(parameters, dict):
+        return parameters
+
+    result = dict()
+
+    if isinstance(parameters, (list, tuple)):
+        for param in parameters:
+            if "=" not in param:
+                raise FrogmlException(
+                    f'The parameter definition passed {param} is invalid. Format is "KEY=VALUE"'
+                )
+            split_param = param.split("=")
+            result[split_param[0]] = split_param[1]
+
+    return result

frogml/clients/batch_job_management/results.py

@@ -0,0 +1,73 @@
+from dataclasses import dataclass, field
+from typing import List
+
+
+@dataclass
+class StartExecutionResult:
+    success: bool = field(default=True)
+    execution_id: str = field(default="")
+    failure_message: str = field(default="")
+
+
+@dataclass
+class StartWarmupResult:
+    success: bool = field(default=True)
+    failure_message: str = field(default="")
+
+
+@dataclass
+class ExecutionStatusResult:
+    success: bool = field(default=True)
+    failure_message: str = field(default="")
+    status: str = field(default="")
+    finished_files: int = field(default=0)
+    total_files: int = field(default=0)
+
+
+@dataclass
+class CancelExecutionResult:
+    success: bool = field(default=True)
+    failure_message: str = field(default="")
+
+
+@dataclass
+class CancelWarmupResult:
+    success: bool = field(default=True)
+    failure_message: str = field(default="")
+
+
+@dataclass
+class GetExecutionReportResult:
+    success: bool = field(default=True)
+    failure_message: str = field(default="")
+    records: List[str] = field(default_factory=list)
+    model_logs: List[str] = field(default_factory=list)
+
+
+@dataclass
+class GetBatchJobPreSignedUploadUrlResult:
+    success: bool = field(default=True)
+    failure_message: str = field(default="")
+    input_path: str = field(default="")
+    output_path: str = field(default="")
+    bucket: str = field(default="")
+    urls: List[str] = field(default_factory=list)
+
+
+@dataclass
+class GetBatchJobPreSignedDownloadUrlResult:
+    success: bool = field(default=True)
+    failure_message: str = field(default="")
+    urls: List[str] = field(default_factory=list)
+
+
+@dataclass
+class StartWarmupJobResult:
+    success: bool = field(default=True)
+    failure_message: str = field(default="")
+
+
+@dataclass
+class CancelWarmupJobResult:
+    success: bool = field(default=True)
+    failure_message: str = field(default="")

frogml/clients/build_management/__init__.py

@@ -0,0 +1 @@
+from .client import BuildsManagementClient

frogml/clients/build_management/client.py

@@ -0,0 +1,115 @@
+import grpc
+from dependency_injector.wiring import Provide
+
+from _frogml_proto.qwak.builds.builds_pb2 import (
+    BuildStatus,
+    ExperimentTrackingValues,
+    GetBuildRequest,
+    GetBuildResponse,
+    ListBuildsRequest,
+    ListBuildsResponse,
+    Metric,
+    Param,
+    RegisterExperimentTrackingValuesRequest,
+    RegisterModelSchemaRequest,
+    UpdateBuildStatusRequest,
+)
+from _frogml_proto.qwak.builds.builds_pb2_grpc import BuildsManagementServiceStub
+from frogml.exceptions import FrogmlException
+from frogml.inner.di_configuration import FrogmlContainer
+
+
+class BuildsManagementClient:
+    """
+    Used for interacting with Feature Registry endpoints
+    """
+
+    def __init__(self, grpc_channel=Provide[FrogmlContainer.core_grpc_channel]):
+        self._builds_management_service = BuildsManagementServiceStub(grpc_channel)
+
+    def update_build_status(self, build_id, build_status, frogml_user_id=""):
+        try:
+            self._builds_management_service.UpdateBuildStatus(
+                UpdateBuildStatusRequest(
+                    build_id=build_id,
+                    build_status=BuildStatus.Value(build_status),
+                    qwak_calling_user_id=frogml_user_id,
+                )
+            )
+
+        except grpc.RpcError as e:
+            raise FrogmlException(
+                f"Failed to update build status, error is {e.details()}"
+            )
+
+    def get_build(self, build_id) -> GetBuildResponse:
+        try:
+            return self._builds_management_service.GetBuild(
+                GetBuildRequest(build_id=build_id)
+            )
+        except grpc.RpcError as e:
+            raise FrogmlException(f"Failed to get build, error is {e.details()}")
+
+    def is_build_exists(self, build_id: str) -> bool:
+        try:
+            self._builds_management_service.GetBuild(GetBuildRequest(build_id=build_id))
+            return True
+        except grpc.RpcError as e:
+            if e.code() == grpc.StatusCode.NOT_FOUND:
+                return False
+            raise FrogmlException(
+                f"Failed to check if build {build_id} is exists, error is {e.details()}"
+            )
+
+    def register_schema(self, build_id, model_schema):
+        try:
+            return self._builds_management_service.RegisterModelSchema(
+                RegisterModelSchemaRequest(build_id=build_id, model_schema=model_schema)
+            )
+
+        except grpc.RpcError as e:
+            raise FrogmlException(f"Failed to register schema, error is {e.details()}")
+
+    def register_experiment_tracking_values(
+        self, build_id, frogml_params=None, cli_params=None, frogml_metrics=None
+    ):
+        frogml_params = frogml_params if frogml_params else {}
+        cli_params = cli_params if cli_params else {}
+        frogml_metrics = frogml_metrics if frogml_metrics else {}
+
+        try:
+            params = cli_params.copy()
+            params.update(frogml_params)
+
+            experiment_tracking_values = ExperimentTrackingValues(
+                metrics=[
+                    Metric(key=key, value=value)
+                    for key, value in frogml_metrics.items()
+                ],
+                params=[
+                    Param(key=key, value=str(value)) for key, value in params.items()
+                ],
+            )
+
+            return self._builds_management_service.RegisterExperimentTrackingValues(
+                RegisterExperimentTrackingValuesRequest(
+                    build_id=build_id,
+                    experiment_tracking_values=experiment_tracking_values,
+                )
+            )
+
+        except grpc.RpcError as e:
+            raise FrogmlException(
+                f"Failed to register experiment tracking values, error is {e.details()}"
+            )
+
+    def list_builds(self, model_uuid="", **kwargs) -> ListBuildsResponse:
+        try:
+            _model_uuid = model_uuid if model_uuid else kwargs.get("branch_id")
+            if not _model_uuid:
+                raise FrogmlException("missing argument model uuid or branch id.")
+            return self._builds_management_service.ListBuilds(
+                ListBuildsRequest(model_uuid=_model_uuid)
+            )
+        except grpc.RpcError as e:
+            raise FrogmlException(f"Failed to list builds, error is {e.details()}")

frogml/clients/build_orchestrator/__init__.py

@@ -0,0 +1,2 @@
+from .client import BuildOrchestratorClient
+from .internal_client import InternalBuildOrchestratorClient

frogml/clients/build_orchestrator/build_model_request_getter.py

@@ -0,0 +1,138 @@
+import json
+import re
+from typing import List, Optional
+
+import yaml
+from yaml import Loader
+
+from _frogml_proto.qwak.builds.build_pb2 import (
+    BuildEnv,
+    BuildProperties,
+    BuildPropertiesV1,
+    DockerEnv,
+    MemoryUnit,
+    ModelUriSpec,
+    PythonEnv,
+    RemoteBuildSpec,
+)
+from _frogml_proto.qwak.fitness_service.fitness_pb2 import PurchaseOption
+from _frogml_proto.qwak.user_application.common.v0.resources_pb2 import (
+    CpuResources,
+    GpuResources,
+    PodComputeResourceTemplateSpec,
+)
+from frogml.exceptions import FrogmlException
+
+
+def _get_build_model_spec(
+    build_conf,
+    verbose: int = 3,
+    git_commit_id: str = "",
+    resolved_model_url: str = "",
+    build_code_path: str = "",
+    build_v1_flag: bool = False,
+    build_config_url: str = "",
+    qwak_sdk_wheel_url: str = "",
+    qwak_sdk_version_url: str = "",
+    build_steps: Optional[List[str]] = None,
+    sdk_version: str = "",
+) -> RemoteBuildSpec:
+    build_spec = RemoteBuildSpec(
+        build_properties=BuildProperties(
+            build_id=build_conf.build_properties.build_id,
+            build_name=build_conf.build_properties.build_name,
+            model_id=build_conf.build_properties.model_id,
+            branch=build_conf.build_properties.branch,
+            tags=build_conf.build_properties.tags,
+            gpu_compatible=build_conf.build_properties.gpu_compatible,
+            model_uri=ModelUriSpec(
+                uri=resolved_model_url or build_conf.build_properties.model_uri.uri,
+                git_credentials=build_conf.build_properties.model_uri.git_credentials,
+                git_credentials_secret=build_conf.build_properties.model_uri.git_credentials_secret,
+                git_branch=build_conf.build_properties.model_uri.git_branch,
+                commit_id=git_commit_id,
+                main_dir=build_conf.build_properties.model_uri.main_dir,
+            ),
+        ),
+        build_env=BuildEnv(
+            docker_env=DockerEnv(
+                base_image=build_conf.build_env.docker.base_image,
+                assumed_iam_role_arn=build_conf.build_env.docker.assumed_iam_role_arn,
+                service_account_key_secret_name=build_conf.build_env.docker.service_account_key_secret_name,
+                no_cache=not build_conf.build_env.docker.cache,
+                env_vars=build_conf.build_env.docker.env_vars,
+            ),
+            python_env=PythonEnv(
+                git_credentials=build_conf.build_properties.model_uri.git_credentials,
+                git_credentials_secret=build_conf.build_properties.model_uri.git_credentials_secret,
+                qwak_sdk_version=sdk_version,
+            ),
+        ),
+        verbose=verbose,
+        build_code_path=build_code_path,
+        build_config=json.dumps(
+            yaml.load(build_conf.to_yaml(), Loader=Loader)  # nosec B506
+        ),
+        build_properties_v1=BuildPropertiesV1(
+            build_config_url=build_config_url,
+            qwak_sdk_wheel_url=qwak_sdk_wheel_url,
+            qwak_sdk_version_url=qwak_sdk_version_url,
+        ),
+        build_v1_flag=build_v1_flag,
+        build_steps=build_steps,
+        purchase_option=_purchase_option_to_enum(build_conf.purchase_option),
+        provision_instance_timeout=build_conf.provision_instance_timeout,
+    )
+
+    if build_conf.build_env.remote.resources.instance:
+        build_spec.client_pod_compute_resources.template_spec.CopyFrom(
+            PodComputeResourceTemplateSpec(
+                template_id=build_conf.build_env.remote.resources.instance
+            )
+        )
+    elif build_conf.build_env.remote.resources.gpu_type:
+        build_spec.client_pod_compute_resources.gpu_resources.CopyFrom(
+            GpuResources(
+                gpu_type=build_conf.build_env.remote.resources.gpu_type,
+                gpu_amount=build_conf.build_env.remote.resources.gpu_amount,
+            )
+        )
+    else:
+        build_spec.client_pod_compute_resources.cpu_resources.CopyFrom(
+            CpuResources(
+                cpu=build_conf.build_env.remote.resources.cpus,
+                memory_amount=int(
+                    re.sub(
+                        r"\D",
+                        "",
+                        build_conf.build_env.remote.resources.memory,
+                    )
+                ),
+                memory_units=map_memory_units(
+                    build_conf.build_env.remote.resources.memory
+                ),
+            )
+        )
+
+    return build_spec
+
+
+def map_memory_units(memory):
+    memory_unit = re.sub(r"\d+", "", memory)
+    if memory_unit == "Gi":
+        return MemoryUnit.GIB
+    elif memory_unit == "Mib":
+        return MemoryUnit.MIB
+    else:
+        return MemoryUnit.UNKNOWN_MEMORY_UNIT
+
+
+def _purchase_option_to_enum(purchase_option: Optional[str]) -> PurchaseOption:
+    if purchase_option == "" or purchase_option is None:
+        return PurchaseOption.INVALID_PURCHASE_OPTION
+    elif purchase_option == "ondemand":
+        return PurchaseOption.ONDEMAND_PURCHASE_OPTION
+    elif purchase_option == "spot":
+        return PurchaseOption.SPOT_PURCHASE_OPTION
+    else:
+        raise FrogmlException("Purchase option must be either 'ondemand' or 'spot'")