frogml-core 0.0.1__py3-none-any.whl

frogml/clients/feature_store/offline_serving_client.py

@@ -0,0 +1,249 @@
+import logging
+from typing import List, Optional
+
+from dependency_injector.wiring import Provide
+from google.protobuf.timestamp_pb2 import Timestamp as ProtoTimestamp
+
+from _frogml_proto.qwak.offline.serving.v1.feature_values_pb2 import (
+    FeaturesetFeatures as ProtoFeaturesetFeatures,
+)
+from _frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2 import (
+    FeatureValuesRequestStatus as ProtoFeatureValuesRequestStatus,
+)
+from _frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2 import (
+    FileFormat as ProtoFileFormat,
+)
+from _frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2 import (
+    GetFeatureValuesInRangeRequest as ProtoGetFeatureValuesInRangeRequest,
+)
+from _frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2 import (
+    GetFeatureValuesInRangeResponse as ProtoGetFeatureValuesInRangeResponse,
+)
+from _frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2 import (
+    GetFeatureValuesRequest as ProtoGetFeatureValuesRequest,
+)
+from _frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2 import (
+    GetFeatureValuesResponse as ProtoGetFeatureValuesResponse,
+)
+from _frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2 import (
+    GetFeatureValuesResultRequest as ProtoGetFeatureValuesResultRequest,
+)
+from _frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2 import (
+    GetFeatureValuesResultResponse as ProtoGetFeatureValuesResultResponse,
+)
+from _frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2 import (
+    GetFileUploadUrlRequest as ProtoGetFileUploadUrlRequest,
+)
+from _frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2 import (
+    GetFileUploadUrlResponse as ProtoGetFileUploadUrlResponse,
+)
+from _frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2_grpc import (
+    FeatureStoreOfflineServingAsyncServiceStub,
+)
+from _frogml_proto.qwak.offline.serving.v1.options_pb2 import (
+    OfflineServingQueryOptions as ProtoOfflineServingQueryOptions,
+)
+from _frogml_proto.qwak.offline.serving.v1.population_pb2 import (
+    Population as ProtoPopulation,
+)
+from _frogml_proto.qwak.offline.serving.v1.population_pb2 import (
+    PopulationFileUploadUrlType as ProtoPopulationFileUploadUrlType,
+)
+from _frogml_proto.qwak.offline.serving.v1.population_pb2 import (
+    TimedPopulation as ProtoTimedPopulation,
+)
+from frogml.inner.di_configuration import FrogmlContainer
+from frogml.inner.tool.retry_utils import retry
+
+
+class FeatureValuesResultNotReadyException(Exception):
+    pass
+
+
+class FeatureValuesTimeoutException(Exception):
+    pass
+
+
+class FsOfflineServingClient:
+    """
+    Querying offline features store
+    """
+
+    def __init__(self, grpc_channel=Provide[FrogmlContainer.core_grpc_channel]):
+        self._client = FeatureStoreOfflineServingAsyncServiceStub(grpc_channel)
+
+    def get_population_file_upload_url(
+        self,
+    ) -> str:
+        """
+        Generating population file upload url
+        :return: pre-signed url to upload population data
+        """
+
+        request: ProtoGetFileUploadUrlRequest = ProtoGetFileUploadUrlRequest(
+            population=ProtoPopulationFileUploadUrlType()
+        )
+
+        response: ProtoGetFileUploadUrlResponse = self._client.GetFileUploadUrl(request)
+        return response.file_upload_url
+
+    def get_feature_values(
+        self,
+        features: List[ProtoFeaturesetFeatures],
+        population: ProtoTimedPopulation,
+        result_file_format: ProtoFileFormat,
+        options: ProtoOfflineServingQueryOptions,
+    ) -> str:
+        """
+        Getting offline feature values
+        :return: Request handle id to poll for result with
+        """
+        request: ProtoGetFeatureValuesRequest = ProtoGetFeatureValuesRequest(
+            features=features,
+            population=population,
+            result_file_format=result_file_format,
+            options=options,
+        )
+
+        response: ProtoGetFeatureValuesResponse = self._client.GetFeatureValues(request)
+        return response.request_id
+
+    def get_feature_values_in_range(
+        self,
+        features: ProtoFeaturesetFeatures,
+        lower_time_bound: ProtoTimestamp,
+        upper_time_bound: ProtoTimestamp,
+        result_file_format: ProtoFileFormat,
+        population: Optional[ProtoPopulation],
+        options: ProtoOfflineServingQueryOptions,
+    ) -> str:
+        """
+        Getting offline feature values
+        :return: Request handle id to poll for result with
+        """
+        request: ProtoGetFeatureValuesInRangeRequest = (
+            ProtoGetFeatureValuesInRangeRequest(
+                features=features,
+                lower_time_bound=lower_time_bound,
+                upper_time_bound=upper_time_bound,
+                result_file_format=result_file_format,
+                population=population,
+                options=options,
+            )
+        )
+
+        response: ProtoGetFeatureValuesInRangeResponse = (
+            self._client.GetFeatureValuesInRange(request)
+        )
+        return response.request_id
+
+    def get_result(self, request_handle: str) -> ProtoGetFeatureValuesResultResponse:
+        request: ProtoGetFeatureValuesResultRequest = (
+            ProtoGetFeatureValuesResultRequest(request_id=request_handle)
+        )
+
+        return self._client.GetFeatureValuesResult(request)
+
+    def _inner_poll(self, request_handle: str) -> ProtoGetFeatureValuesResultResponse:
+        """
+        This function receives request_handle amd polls for results.
+
+        Args:
+            request_handle: request_id for polling
+
+        Returns: if the request status is PENDING it throws FeatureValuesResultNotReadyException.
+                 else returns the Feature Values Result
+        """
+        response: ProtoGetFeatureValuesResultResponse = self.get_result(
+            request_handle=request_handle
+        )
+
+        if (
+            response.status
+            == ProtoFeatureValuesRequestStatus.FEATURE_VALUES_REQUEST_STATUS_PENDING
+        ):
+            logging.info("Feature Values query is still in progress...")
+            raise FeatureValuesResultNotReadyException()
+
+        return response
+
+    def poll_for_result(
+        self,
+        request_handle: str,
+        timeout_seconds: int = 60 * 60,
+        poll_interval_seconds: int = 10,
+    ) -> ProtoGetFeatureValuesResultResponse:
+        """
+        This function receives request_handle and polls for Feature Values Result.
+        it relies on _inner_poll() to raise FeatureValuesResultNotReadyException in case the result is not ready.
+        only once the desired amount of attempts completed it raises FeatureValuesTimeoutException
+        Args:
+            request_handle: request_id
+            timeout_seconds: total number of seconds for all attempts before timeout
+            poll_interval_seconds: sleep time between attempts
+
+        Returns: Feature Values Result or raises FeatureValuesTimeoutException if attempts exceeded timeout_seconds.
+
+        """
+        try:
+            result = retry(
+                f=self._inner_poll,
+                kwargs={"request_handle": request_handle},
+                exceptions=FeatureValuesResultNotReadyException,
+                attempts=int(timeout_seconds / poll_interval_seconds) + 1,
+                delay=poll_interval_seconds,
+            )
+        except FeatureValuesResultNotReadyException:
+            raise FeatureValuesTimeoutException(
+                f"Feature values query timed out. Frogml limits query execution time to {int(timeout_seconds / 60)} minutes"
+            )
+
+        return result
+
+    def get_feature_values_blocking(
+        self,
+        features: List[ProtoFeaturesetFeatures],
+        population: ProtoTimedPopulation,
+        result_file_format: ProtoFileFormat,
+        options: ProtoOfflineServingQueryOptions,
+        timeout_seconds: int = 60 * 60,
+        poll_interval_seconds: int = 10,
+    ) -> ProtoGetFeatureValuesResultResponse:
+        request_handle: str = self.get_feature_values(
+            features=features,
+            population=population,
+            result_file_format=result_file_format,
+            options=options,
+        )
+
+        return self.poll_for_result(
+            request_handle=request_handle,
+            timeout_seconds=timeout_seconds,
+            poll_interval_seconds=poll_interval_seconds,
+        )
+
+    def get_feature_values_in_range_blocking(
+        self,
+        features: ProtoFeaturesetFeatures,
+        lower_time_bound: ProtoTimestamp,
+        upper_time_bound: ProtoTimestamp,
+        result_file_format: ProtoFileFormat,
+        options: ProtoOfflineServingQueryOptions,
+        population: Optional[ProtoPopulation] = None,
+        timeout_seconds: int = 60 * 60,
+        poll_interval_seconds: int = 10,
+    ) -> ProtoGetFeatureValuesResultResponse:
+        request_handle: str = self.get_feature_values_in_range(
+            features=features,
+            lower_time_bound=lower_time_bound,
+            upper_time_bound=upper_time_bound,
+            population=population,
+            result_file_format=result_file_format,
+            options=options,
+        )
+
+        return self.poll_for_result(
+            request_handle=request_handle,
+            timeout_seconds=timeout_seconds,
+            poll_interval_seconds=poll_interval_seconds,
+        )

frogml/clients/feature_store/operator_client.py

@@ -0,0 +1,164 @@
+from typing import Optional
+
+from dependency_injector.wiring import Provide
+
+from _frogml_proto.qwak.feature_store.features.feature_set_pb2 import FeatureSetSpec
+from _frogml_proto.qwak.feature_store.sources.data_source_pb2 import DataSourceSpec
+from _frogml_proto.qwak.features_operator.v3.features_operator_async_service_pb2 import (
+    DataSourceValidationOptions,
+    FeatureSetValidationOptions,
+    GetValidationResultRequest,
+    GetValidationResultResponse,
+    ValidateDataSourceRequest,
+    ValidateFeatureSetRequest,
+    ValidationResponse,
+)
+from _frogml_proto.qwak.features_operator.v3.features_operator_async_service_pb2_grpc import (
+    FeaturesOperatorAsyncServiceStub,
+)
+from _frogml_proto.qwak.features_operator.v3.features_operator_pb2 import (
+    ValidationNotReadyResponse,
+)
+from frogml.inner.di_configuration import FrogmlContainer
+from frogml.inner.tool.retry_utils import retry
+
+
+class ValidationNotReadyException(Exception):
+    pass
+
+
+class ValidationTimeoutException(Exception):
+    pass
+
+
+class FeaturesOperatorClient:
+    """
+    Validates and samples features store objects like: data sources and feature sets.
+    """
+
+    def __init__(self, grpc_channel=Provide[FrogmlContainer.core_grpc_channel]):
+        self._v3_client = FeaturesOperatorAsyncServiceStub(grpc_channel)
+
+    def validate_data_source(
+        self,
+        data_source_spec: DataSourceSpec,
+        num_samples: int = 10,
+        validation_options: Optional[DataSourceValidationOptions] = None,
+    ) -> str:
+        """
+        Validates and fetches a sample from the data source
+        :return: Request handle id to poll for result with
+        """
+        request: ValidateDataSourceRequest = ValidateDataSourceRequest(
+            data_source_spec=data_source_spec,
+            num_samples=num_samples,
+            validation_options=validation_options,
+        )
+        response: ValidationResponse = self._v3_client.ValidateDataSource(request)
+        return response.request_id
+
+    def validate_featureset(
+        self,
+        featureset_spec: FeatureSetSpec,
+        resource_path: Optional[str] = None,
+        num_samples: int = 10,
+        validation_options: Optional[FeatureSetValidationOptions] = None,
+    ) -> str:
+        """
+        Validates and fetches a sample from the featureset
+        :return: Request handle id to poll for result with
+        """
+        resource_path = str(resource_path) if resource_path is not None else None
+        request: ValidateFeatureSetRequest = ValidateFeatureSetRequest(
+            feature_set_spec=featureset_spec,
+            num_samples=num_samples,
+            zip_path=resource_path,
+            validation_options=validation_options,
+        )
+
+        response: ValidationResponse = self._v3_client.ValidateFeatureSet(request)
+        return response.request_id
+
+    def get_result(self, request_handle: str) -> GetValidationResultResponse:
+        request: GetValidationResultRequest = GetValidationResultRequest(
+            request_id=request_handle
+        )
+        response: GetValidationResultResponse = self._v3_client.GetValidationResult(
+            request
+        )
+
+        return response
+
+    def _inner_poll(self, request_handle: str) -> GetValidationResultResponse:
+        response: GetValidationResultResponse = self.get_result(
+            request_handle=request_handle
+        )
+        response_type = getattr(response, response.WhichOneof("type"))
+
+        if isinstance(response_type, ValidationNotReadyResponse):
+            raise ValidationNotReadyException()
+
+        return response
+
+    def poll_for_result(
+        self,
+        request_handle: str,
+        timeout_seconds: int = 5 * 60,
+        poll_interval_seconds: int = 3,
+    ) -> GetValidationResultResponse:
+        try:
+            result = retry(
+                f=self._inner_poll,
+                kwargs={"request_handle": request_handle},
+                exceptions=ValidationNotReadyException,
+                attempts=timeout_seconds / poll_interval_seconds,
+                delay=poll_interval_seconds,
+            )
+        except ValidationNotReadyException:
+            raise ValidationTimeoutException(
+                f"Validation timed out. Frogml limits validation execution time to {timeout_seconds} seconds"
+            )
+
+        return result
+
+    def validate_data_source_blocking(
+        self,
+        data_source_spec: DataSourceSpec,
+        num_samples: int = 10,
+        timeout_seconds: int = 5 * 60,
+        poll_interval_seconds: int = 3,
+        validation_options: Optional[DataSourceValidationOptions] = None,
+    ) -> GetValidationResultResponse:
+        request_handle: str = self.validate_data_source(
+            data_source_spec=data_source_spec,
+            num_samples=num_samples,
+            validation_options=validation_options,
+        )
+
+        return self.poll_for_result(
+            request_handle=request_handle,
+            timeout_seconds=timeout_seconds,
+            poll_interval_seconds=poll_interval_seconds,
+        )
+
+    def validate_featureset_blocking(
+        self,
+        featureset_spec: FeatureSetSpec,
+        resource_path: Optional[str] = None,
+        num_samples: int = 10,
+        timeout_seconds: int = 5 * 60,
+        poll_interval_seconds: int = 3,
+        validation_options: Optional[FeatureSetValidationOptions] = None,
+    ) -> GetValidationResultResponse:
+        request_handle: str = self.validate_featureset(
+            featureset_spec=featureset_spec,
+            resource_path=resource_path,
+            num_samples=num_samples,
+            validation_options=validation_options,
+        )
+
+        return self.poll_for_result(
+            request_handle=request_handle,
+            timeout_seconds=timeout_seconds,
+            poll_interval_seconds=poll_interval_seconds,
+        )

frogml/clients/file_versioning/client.py

@@ -0,0 +1,77 @@
+from pathlib import Path
+from typing import Optional
+
+import grpc
+from dependency_injector.wiring import Provide
+
+from _frogml_proto.qwak.file_versioning.file_versioning_pb2 import (
+    FileTagFilter,
+    FileTagSpec,
+)
+from _frogml_proto.qwak.file_versioning.file_versioning_service_pb2 import (
+    GetModelFileTagsRequest,
+    RegisterFileTagRequest,
+)
+from _frogml_proto.qwak.file_versioning.file_versioning_service_pb2_grpc import (
+    FileVersioningManagementServiceStub,
+)
+from frogml.exceptions import FrogmlException
+from frogml.inner.di_configuration import FrogmlContainer
+
+
+class FileVersioningManagementClient:
+    def __init__(self, grpc_channel=Provide[FrogmlContainer.core_grpc_channel]):
+        self._file_management_service = FileVersioningManagementServiceStub(
+            grpc_channel
+        )
+
+    def register_file_tag(
+        self, model_id: str, tag: str, file_path: str, build_id: Optional[str] = ""
+    ) -> None:
+        """
+        Register file tag to service
+        Args:
+            model_id: model id to save tag by.
+            file_path: The file path
+            tag: tag to save the file under.
+            build_id: build id to save tag by.
+
+        Returns:
+
+        """
+        try:
+            file_extension = Path(file_path).suffix.replace(".", "")
+
+            self._file_management_service.RegisterFileTag(
+                RegisterFileTagRequest(
+                    file_tag_spec=(
+                        FileTagSpec(
+                            build_id=build_id,
+                            model_id=model_id,
+                            tag=tag,
+                            extension_type=file_extension,
+                        )
+                    )
+                )
+            )
+        except grpc.RpcError as e:
+            if e.args[0].code != grpc.StatusCode.ALREADY_EXISTS:
+                raise FrogmlException(
+                    f"Failed to register file tag, error is {e.details()}"
+                )
+            else:
+                raise e
+
+    def get_model_file_tags(
+        self, model_id: str, build_id: str, file_tag_filter: FileTagFilter = None
+    ):
+        try:
+            return self._file_management_service.GetModelFileTags(
+                GetModelFileTagsRequest(
+                    model_id=model_id, build_id=build_id, filter=file_tag_filter
+                )
+            )
+        except grpc.RpcError as e:
+            raise FrogmlException(
+                f"Failed to list model file tags, error is {e.details()}"
+            )

frogml/clients/file_versioning/file_tag_filter.py

@@ -0,0 +1,32 @@
+from dataclasses import dataclass, field
+from enum import Enum
+
+from _frogml_proto.qwak.file_versioning.file_versioning_pb2 import (
+    FileTagFilter as FileTagFilterProto,
+)
+
+
+class FileTagFilerType(Enum):
+    FILE_TAG_FILTER_TYPE_INVALID = 0
+    FILE_TAG_FILTER_TYPE_CONTAINS = 1
+    FILE_TAG_FILTER_TYPE_PREFIX = 2
+
+
+@dataclass
+class FileTagFilter:
+    value: str = field(default=None)
+    type: FileTagFilerType = field(
+        default=FileTagFilerType.FILE_TAG_FILTER_TYPE_INVALID
+    )
+
+    def to_proto(self):
+        if self.type is FileTagFilerType.FILE_TAG_FILTER_TYPE_CONTAINS:
+            return FileTagFilterProto(
+                tag_contains=self.value,
+            )
+        elif self.type is FileTagFilerType.FILE_TAG_FILTER_TYPE_PREFIX:
+            return FileTagFilterProto(
+                tag_prefix=self.value,
+            )
+        else:
+            return None

frogml/clients/instance_template/client.py

@@ -0,0 +1,67 @@
+from typing import List
+
+import grpc
+from dependency_injector.wiring import Provide
+
+from _frogml_proto.qwak.instance_template.instance_template_pb2 import (
+    InstanceFilter,
+    InstanceTemplateSpec,
+    InstanceTypeFilter,
+)
+from _frogml_proto.qwak.instance_template.instance_template_service_pb2 import (
+    GetInstanceTemplateRequest,
+    GetInstanceTemplateResponse,
+    ListInstanceTemplatesRequest,
+    ListInstanceTemplatesResponse,
+)
+from _frogml_proto.qwak.instance_template.instance_template_service_pb2_grpc import (
+    InstanceTemplateManagementServiceStub,
+)
+from frogml.exceptions import FrogmlException
+from frogml.inner.di_configuration import FrogmlContainer
+
+GET_INSTANCE_TEMPLATE_ERROR_FORMAT = (
+    "An error occurred while trying to get instance template: {e}"
+)
+
+
+class InstanceTemplateManagementClient:
+    def __init__(self, grpc_channel=Provide[FrogmlContainer.core_grpc_channel]):
+        self._instance_template_service = InstanceTemplateManagementServiceStub(
+            grpc_channel
+        )
+
+    def get_instance_template(self, instance_template_id: str) -> InstanceTemplateSpec:
+        try:
+            result: GetInstanceTemplateResponse = (
+                self._instance_template_service.GetInstanceTemplate(
+                    GetInstanceTemplateRequest(id=instance_template_id)
+                )
+            )
+            return result.instance_template
+        except grpc.RpcError as e:
+            raise FrogmlException(
+                GET_INSTANCE_TEMPLATE_ERROR_FORMAT.format(
+                    e=e.details(), error_code=e.code()
+                )
+            )
+        except Exception as e:
+            raise FrogmlException(GET_INSTANCE_TEMPLATE_ERROR_FORMAT.format(e=e))
+
+    def list_instance_templates(self) -> List[InstanceTemplateSpec]:
+        try:
+            result: ListInstanceTemplatesResponse = self._instance_template_service.ListInstanceTemplates(
+                ListInstanceTemplatesRequest(
+                    optional_instance_filter=InstanceFilter(
+                        instance_type_filter=InstanceTypeFilter.INSTANCE_TYPE_FILTER_ALL
+                    )
+                )
+            )
+            return list(result.instance_template_list)
+        except grpc.RpcError as e:
+            raise FrogmlException(
+                GET_INSTANCE_TEMPLATE_ERROR_FORMAT.format(e=e.details()),
+                status_code=e.code(),
+            )
+        except Exception as e:
+            raise FrogmlException(GET_INSTANCE_TEMPLATE_ERROR_FORMAT.format(e=e))

frogml/clients/integration_management/integration_manager_client.py

@@ -0,0 +1,35 @@
+from typing import List
+
+from dependency_injector.wiring import Provide, inject
+
+from _frogml_proto.qwak.integration.integration_pb2 import Integration
+from _frogml_proto.qwak.integration.integration_service_pb2 import (
+    GetIntegrationRequest,
+    GetIntegrationResponse,
+    ListIntegrationRequest,
+    ListIntegrationsResponse,
+)
+from _frogml_proto.qwak.integration.integration_service_pb2_grpc import (
+    IntegrationManagementServiceStub,
+)
+from frogml.inner.di_configuration import FrogmlContainer
+
+
+class IntegrationManagerClient:
+    @inject
+    def __init__(self, grpc_channel=Provide[FrogmlContainer.core_grpc_channel]):
+        self._grpc_client: IntegrationManagementServiceStub = (
+            IntegrationManagementServiceStub(grpc_channel)
+        )
+
+    def list_integrations(self) -> List[Integration]:
+        response: ListIntegrationsResponse = self._grpc_client.ListIntegrations(
+            ListIntegrationRequest()
+        )
+        return response.integrations
+
+    def get_by_id(self, integration_id: str) -> Integration:
+        response: GetIntegrationResponse = self._grpc_client.GetIntegration(
+            GetIntegrationRequest(integration_id=integration_id)
+        )
+        return response.integration

frogml/clients/integration_management/integration_utils.py

@@ -0,0 +1,27 @@
+from typing import List
+
+from _frogml_proto.qwak.integration.integration_pb2 import Integration
+from frogml.clients.integration_management.integration_manager_client import (
+    IntegrationManagerClient,
+)
+from frogml.clients.integration_management.openai.openai_system_secret import (
+    OpenAIApiKeySystemSecret,
+)
+
+
+class IntegrationUtils:
+    def __init__(self):
+        self._client: IntegrationManagerClient = IntegrationManagerClient()
+
+    def get_openai_api_keys(self) -> List[OpenAIApiKeySystemSecret]:
+        all_integrations: List[Integration] = self._client.list_integrations()
+
+        openai_integration: List[Integration] = [
+            i for i in all_integrations if i.WhichOneof("type") == "openai_integration"
+        ]
+
+        openai_secrets: List[OpenAIApiKeySystemSecret] = [
+            OpenAIApiKeySystemSecret.from_proto(proto=o) for o in openai_integration
+        ]
+
+        return [s for s in openai_secrets if s is not None]