frogml-core 0.0.116__py3-none-any.whl → 0.0.118__py3-none-any.whl

core/clients/integration_management/openai/openai_system_secret.py

@@ -0,0 +1,58 @@
+from base64 import b64decode
+from dataclasses import dataclass
+from typing import Optional
+
+from frogml._proto.qwak.admiral.secret.v0.secret_pb2 import SystemSecretValue
+from frogml._proto.qwak.integration.integration_pb2 import Integration
+from frogml._proto.qwak.integration.open_a_i_integration_pb2 import (
+    OpenAIApiKeySystemSecretDescriptor,
+    OpenAIIntegration,
+)
+from frogml.core.clients.administration.eco_system.eco_system_utils import (
+    EcosystemUtils,
+)
+from frogml.core.clients.system_secret.system_secret_client import SystemSecretClient
+
+
+@dataclass
+class OpenAIApiKeySystemSecret:
+    integration_id: str
+    integration_name: str
+    secret_name: str
+    secret_key: str
+
+    @classmethod
+    def from_proto(cls, proto: Integration) -> Optional["OpenAIApiKeySystemSecret"]:
+        if proto.WhichOneof("type") != "openai_integration":
+            raise ValueError(f"Got wrong type of integration: {proto}")
+        openai_integration: OpenAIIntegration = proto.openai_integration
+
+        if (
+            openai_integration.WhichOneof("auth")
+            != "openai_api_key_system_secret_descriptor"
+        ):
+            return None
+
+        system_secret: OpenAIApiKeySystemSecretDescriptor = (
+            openai_integration.openai_api_key_system_secret_descriptor
+        )
+
+        return cls(
+            integration_id=proto.integration_id,
+            integration_name=proto.name,
+            secret_name=system_secret.secret_name,
+            secret_key=system_secret.api_key_secret_key,
+        )
+
+    def get_api_key(self) -> str:
+        system_secret_client = SystemSecretClient()
+        default_environment_id: str = EcosystemUtils().get_default_environment_id()
+        system_secret: SystemSecretValue = system_secret_client.get_system_secret(
+            name=self.secret_name, env_id=default_environment_id
+        )
+
+        if system_secret.WhichOneof("value") != "opaque_pair":
+            raise ValueError("Got unsupported system secret value type")
+
+        encoded_api_key: str = system_secret.opaque_pair.pairs.get(key=self.secret_key)
+        return b64decode(encoded_api_key).decode("utf-8")

core/clients/jfrog_gateway/client.py

@@ -0,0 +1,60 @@
+from typing import cast
+
+from dependency_injector.wiring import Provide, inject
+
+from frogml.core.inner.di_configuration import FrogmlContainer
+from frogml.core.inner.tool.grpc.grpc_try_wrapping import grpc_try_catch_wrapper
+from frogml._proto.qwak.ecosystem.jfrog.v0.jfrog_tenant_info_service_pb2 import (
+    GetJFrogTenantInfoResponse,
+    GetJFrogTenantInfoRequest,
+)
+from frogml._proto.qwak.ecosystem.jfrog.v0.jfrog_tenant_info_service_pb2_grpc import (
+    JFrogTenantInfoServiceStub,
+)
+from frogml._proto.qwak.jfrog.gateway.v0.repository_service_pb2 import (
+    GetRepositoryConfigurationResponse,
+    GetRepositoryConfigurationRequest,
+)
+from frogml._proto.qwak.jfrog.gateway.v0.repository_service_pb2_grpc import (
+    RepositoryServiceStub,
+)
+
+
+class JfrogGatewayClient:
+    """
+    Used for interacting with Feature Registry endpoints
+    """
+
+    @inject
+    def __init__(self, grpc_channel=Provide[FrogmlContainer.core_grpc_channel]):
+        self.__repository_service = RepositoryServiceStub(grpc_channel)
+        self.__jfrog_tenant_info_service = JFrogTenantInfoServiceStub(grpc_channel)
+
+    @grpc_try_catch_wrapper("Failed to get repository configuration")
+    def get_repository_configuration(
+        self, repository_key: str
+    ) -> GetRepositoryConfigurationResponse:
+        """
+        Get repository configuration
+        :param repository_key: Repository name
+        :return: Model
+        """
+        request = GetRepositoryConfigurationRequest(repository_key=repository_key)
+
+        return cast(
+            GetRepositoryConfigurationResponse,
+            self.__repository_service.GetRepositoryConfiguration(request),
+        )
+
+    @grpc_try_catch_wrapper("Failed to get JFrog tenant info")
+    def get_jfrog_tenant_info(self) -> GetJFrogTenantInfoResponse:
+        """
+        Get the customer's JFrog tenant info
+        :return: The JFrog tenant info response
+        """
+        return cast(
+            GetJFrogTenantInfoResponse,
+            self.__jfrog_tenant_info_service.GetJFrogTenantInfo(
+                GetJFrogTenantInfoRequest()
+            ),
+        )

core/clients/kube_deployment_captain/client.py

@@ -0,0 +1,233 @@
+from typing import Dict, Optional
+
+import grpc
+from google.protobuf.timestamp_pb2 import Timestamp
+
+from frogml._proto.qwak.feature_store.entities.entity_pb2 import Entity
+from frogml._proto.qwak.feature_store.features.feature_set_types_pb2 import (
+    BatchFeatureSetV1,
+    StreamingAggregationFeatureSet,
+    StreamingFeatureSetV1,
+)
+from frogml._proto.qwak.kube_deployment_captain.feature_set_deployment_pb2 import (
+    BatchDetails,
+    FeatureSetDeployment,
+)
+from frogml._proto.qwak.kube_deployment_captain.kube_deployment_captain_service_pb2 import (
+    BatchFeatureSetV1DeploymentRequest,
+    DeployFeatureSetResponse,
+    DeployStreamingAggregationCompactionRequest,
+    DeployStreamingAggregationCompactionResponse,
+    StreamingFeatureSetV1DeploymentRequest,
+    UndeployBatchFeaturesetRequest,
+    UndeployFeatureSetRequest,
+    UndeployFeatureSetResponse,
+)
+from frogml._proto.qwak.kube_deployment_captain.kube_deployment_captain_service_pb2_grpc import (
+    KubeDeploymentCaptainStub,
+)
+from frogml.core.exceptions import FrogmlException
+from frogml.core.inner.tool.grpc.grpc_tools import create_grpc_channel
+
+
+class KubeDeploymentClient:
+    """
+    Used for interacting with Kube Deployment Captain
+    """
+
+    def __init__(
+        self,
+        edge_services_url: str,
+        enable_ssl: bool = True,
+        grpc_channel: Optional[grpc.Channel] = None,
+    ):
+        self._channel = None
+
+        if grpc_channel:
+            self._channel = grpc_channel
+        else:
+            self._channel = create_grpc_channel(
+                url=edge_services_url,
+                enable_ssl=enable_ssl,
+                status_for_retry=(
+                    grpc.StatusCode.UNAVAILABLE,
+                    grpc.StatusCode.DEADLINE_EXCEEDED,
+                    grpc.StatusCode.INTERNAL,
+                ),
+            )
+
+        self._kube_deployment = KubeDeploymentCaptainStub(self._channel)
+
+    def deploy_streaming_offline_feature_set(
+        self,
+        feature_set_name: str,
+        streaming_feature_set: StreamingFeatureSetV1,
+        extra_deployment_configuration: Optional[Dict[str, str]],
+        extra_env_configuration: Optional[Dict[str, str]],
+        entity: Entity,
+        feature_set_id: str,
+        environment_id: str,
+        secret_service_url: Optional[str],
+        run_id: Optional[int] = 0,
+    ) -> DeployFeatureSetResponse:
+        """
+        Deploy streaming offline feature set
+        Args: feature_set_name: Name of the feature set
+              extra_deployment_configuration: Dictionary of spark configuration
+              extra_env_configuration: Environment configuration
+              Entity: entity of the feature set
+              feature_set_id: Feature set id
+              environment_id: Environment id
+              secret_service_url: Secret service of the url
+
+        Return: DeployFeatureSetResponse
+        """
+        try:
+            streaming_featureset_request = StreamingFeatureSetV1DeploymentRequest(
+                feature_set_deployment=FeatureSetDeployment(
+                    featureset_name=feature_set_name,
+                    featureset_id=feature_set_id,
+                    streaming_feature_set=streaming_feature_set,
+                    extra_deployment_configuration=extra_deployment_configuration,
+                    extra_env_configuration=extra_env_configuration,
+                    entity=entity,
+                    secret_service_url=secret_service_url,
+                    run_id=run_id,
+                ),
+                environment_id=environment_id,
+            )
+
+            return self._kube_deployment.DeployStreamingOfflineFeatureSetV1(
+                streaming_featureset_request
+            )
+        except grpc.RpcError as e:
+            raise FrogmlException(
+                f"Failed to deploy offline streaming feature set, error is {repr(e)}"
+            )
+
+    def deploy_batch_feature_set(
+        self,
+        featureset_name: str,
+        featureset_id: str,
+        batch_featureset: BatchFeatureSetV1,
+        extra_deployment_configuration: Optional[Dict[str, str]],
+        extra_env_configuration: Optional[Dict[str, str]],
+        entity: Entity,
+        secret_service_url: str,
+        job_run_id: int,
+        environment_id: str,
+        batch_execution_date: Timestamp,
+        try_number: int,
+        batch_trigger_date: Timestamp,
+    ) -> DeployFeatureSetResponse:
+        """
+        Deploy batch featureset
+        Args: featureset_name: Name of the feature set
+              featureset_id: Id of the featureset
+              batch_featureset: batch featureset definition
+              extra_deployment_configuration: Dictionary of spark configuration
+              extra_env_configuration: Environment configuration
+              Entity: entity of the feature set
+              secret_service_url: Secret service of the url
+              job_run_id: the run id of the job
+              environment_id: Environment id
+              batch_execution_date: the current batch execution time
+              prev_batch_execution_date: previous batch exection time
+              prev_ingestion_window_date: previous ingestion window time
+              batch_trigger_date: system trigger date
+        Return: DeployFeatureSetResponse
+        """
+        try:
+            batch_featureset_request = BatchFeatureSetV1DeploymentRequest(
+                feature_set_deployment=FeatureSetDeployment(
+                    featureset_name=featureset_name,
+                    featureset_id=featureset_id,
+                    batch_feature_set=batch_featureset,
+                    extra_deployment_configuration=extra_deployment_configuration,
+                    extra_env_configuration=extra_env_configuration,
+                    secret_service_url=secret_service_url,
+                    entity=entity,
+                    run_id=job_run_id,
+                ),
+                batch_details=BatchDetails(
+                    batch_execution_date=batch_execution_date,
+                    try_number=try_number,
+                    batch_trigger_date=batch_trigger_date,
+                ),
+                environment_id=environment_id,
+            )
+            return self._kube_deployment.DeployBatchFeatureSetV1(
+                batch_featureset_request
+            )
+        except grpc.RpcError as e:
+            raise FrogmlException(
+                f"Failed to deploy batch feature set, error is {repr(e)}"
+            )
+
+    def deploy_streaming_aggregation_compaction(
+        self,
+        feature_set_id: str,
+        feature_set_name: str,
+        environment_id: str,
+        streaming_aggregation_feature_set: StreamingAggregationFeatureSet,
+        entity: Entity,
+        extra_deployment_configuration: Optional[Dict[str, str]],
+        extra_env_configuration: Optional[Dict[str, str]],
+        secret_service_url: str,
+        job_run_id: int,
+    ) -> DeployStreamingAggregationCompactionResponse:
+        try:
+            streaming_aggregation_compaction_deploy_request: (
+                DeployStreamingAggregationCompactionRequest
+            ) = DeployStreamingAggregationCompactionRequest(
+                feature_set_deployment=FeatureSetDeployment(
+                    featureset_name=feature_set_name,
+                    featureset_id=feature_set_id,
+                    streaming_aggregation_feature_set=streaming_aggregation_feature_set,
+                    extra_deployment_configuration=extra_deployment_configuration,
+                    extra_env_configuration=extra_env_configuration,
+                    entity=entity,
+                    secret_service_url=secret_service_url,
+                    run_id=job_run_id,
+                ),
+                environment_id=environment_id,
+            )
+            return self._kube_deployment.DeployStreamingAggregationCompaction(
+                streaming_aggregation_compaction_deploy_request
+            )
+        except grpc.RpcError as e:
+            raise FrogmlException(
+                f"Failed to deploy streaming aggregation feature set, error is {repr(e)}"
+            )
+
+    def undeploy_featureset_resources(
+        self, featureset_name: str
+    ) -> UndeployFeatureSetResponse:
+        """
+        Undeploy featureset by name. Delete all spark resources.
+        """
+        try:
+            undeploy_request: UndeployFeatureSetRequest = UndeployFeatureSetRequest(
+                feature_set_name=featureset_name
+            )
+            return self._kube_deployment.UndeployFeatureSet(undeploy_request)
+        except grpc.RpcError as e:
+            raise FrogmlException(
+                f"Failed to undeploy featureset: {featureset_name}, error is {repr(e)}"
+            )
+
+    def undeploy_batch_featureset(
+        self, featureset_name: str
+    ) -> UndeployFeatureSetResponse:
+        """
+        Undeploy featureset by name. Delete all spark resources.
+        """
+        try:
+            undeploy_request: UndeployBatchFeaturesetRequest = (
+                UndeployBatchFeaturesetRequest(featureset_name=featureset_name)
+            )
+            return self._kube_deployment.UndeployBatchFeatureset(undeploy_request)
+        except grpc.RpcError as e:
+            raise FrogmlException(
+                f"Failed to undeploy batch featureset: {featureset_name}, error is {repr(e)}"
+            )

core/clients/location_discovery/client.py

@@ -0,0 +1,73 @@
+from typing import Callable
+
+import grpc
+from dependency_injector.wiring import Provide
+from frogml.core.exceptions import FrogmlException
+from frogml.core.inner.di_configuration import FrogmlContainer
+from frogml._proto.qwak.service_discovery.service_discovery_location_pb2 import (
+    ServiceLocationDescriptor,
+)
+from frogml._proto.qwak.service_discovery.service_discovery_location_service_pb2 import (
+    GetServingUrlRequest as ProtoGetServingUrlRequest,
+    GetServingUrlRequestResponse as ProtoGetServingUrlRequestResponse,
+)
+from frogml._proto.qwak.service_discovery.service_discovery_location_service_pb2_grpc import (
+    LocationDiscoveryServiceStub,
+)
+
+
+class LocationDiscoveryClient:
+    """
+    Client for querying service locations from the LocationDiscoveryService.
+    """
+
+    def __init__(self, grpc_channel=Provide[FrogmlContainer.core_grpc_channel]):
+        self._service = LocationDiscoveryServiceStub(grpc_channel)
+
+    @staticmethod
+    def _get_location(
+        method: Callable[
+            [ProtoGetServingUrlRequest], ProtoGetServingUrlRequestResponse
+        ],
+    ) -> ServiceLocationDescriptor:
+        """
+        Calls a specific RPC and extracts the service location descriptor.
+
+        Args:
+            method: The gRPC method to call.
+
+        Returns:
+            ServiceLocationDescriptor: Contains service_url and location metadata.
+        """
+        try:
+            request = ProtoGetServingUrlRequest()
+            response = method(request)
+            return response.location
+        except grpc.RpcError as e:
+            raise FrogmlException(
+                f"Failed to retrieve service location. Error is: {e.details()}"
+            ) from e
+
+    def get_offline_serving(self) -> ServiceLocationDescriptor:
+        """Fetches the offline serving service location."""
+        return self._get_location(self._service.GetOfflineServingUrl)
+
+    def get_distribution_manager(self) -> ServiceLocationDescriptor:
+        """Fetches the distribution manager service location."""
+        return self._get_location(self._service.GetDistributionManagerUrl)
+
+    def get_analytics_engine(self) -> ServiceLocationDescriptor:
+        """Fetches the analytics engine service location."""
+        return self._get_location(self._service.GetAnalyticsEngineUrl)
+
+    def get_metrics_gateway(self) -> ServiceLocationDescriptor:
+        """Fetches the metrics gateway service location."""
+        return self._get_location(self._service.GetMetricsGatewayUrl)
+
+    def get_features_operator(self) -> ServiceLocationDescriptor:
+        """Fetches the features operator service location."""
+        return self._get_location(self._service.GetFeaturesOperatorUrl)
+
+    def get_hosting_gateway(self) -> ServiceLocationDescriptor:
+        """Fetches the hosting gateway service location."""
+        return self._get_location(self._service.GetHostingGatewayUrl)

core/clients/logging_client/client.py

@@ -0,0 +1,148 @@
+from typing import Optional
+
+from grpc import RpcError
+
+from frogml._proto.qwak.logging.log_filter_pb2 import LogText, SearchFilter
+from frogml._proto.qwak.logging.log_reader_service_pb2 import (
+    ReadLogsRequest,
+    ReadLogsResponse,
+)
+from frogml._proto.qwak.logging.log_reader_service_pb2_grpc import LogReaderServiceStub
+from frogml._proto.qwak.logging.log_source_pb2 import (
+    InferenceExecutionSource,
+    LogSource,
+    ModelRuntimeSource,
+    RemoteBuildSource,
+)
+from frogml.core.clients.administration.eco_system.client import EcosystemClient
+from frogml.core.exceptions import FrogmlException
+from frogml.core.inner.tool.grpc.grpc_tools import create_grpc_channel
+
+
+class LoggingClient:
+    """
+    Used for interacting with Logging endpoint
+    """
+
+    def __init__(
+        self,
+        endpoint_url: Optional[str] = None,
+        enable_ssl: bool = True,
+        environment_id: Optional[str] = None,
+    ):
+        if endpoint_url is None:
+            user_context = EcosystemClient().get_authenticated_user_context().user
+            if environment_id is None:
+                environment_id = user_context.account_details.default_environment_id
+
+            if environment_id not in user_context.account_details.environment_by_id:
+                raise FrogmlException(
+                    f"Configuration for environment [{environment_id}] was not found"
+                )
+
+            endpoint_url = user_context.account_details.environment_by_id[
+                environment_id
+            ].configuration.edge_services_url
+
+        self._channel = create_grpc_channel(url=endpoint_url, enable_ssl=enable_ssl)
+
+        self._logging_service = LogReaderServiceStub(self._channel)
+
+    def read_build_logs(
+        self,
+        build_id=None,
+        before_offset=None,
+        after_offset=None,
+        max_number_of_results=None,
+        log_text_filter=None,
+    ):
+        try:
+            response = self.read_logs(
+                source=LogSource(remote_build=RemoteBuildSource(build_id=build_id)),
+                before_offset=before_offset,
+                after_offset=after_offset,
+                log_text_filter=log_text_filter,
+                max_number_of_results=max_number_of_results,
+            )
+
+            return response
+        except FrogmlException as e:
+            raise FrogmlException(f"Failed to fetch build logs, error is [{e}]")
+
+    def read_model_runtime_logs(
+        self,
+        build_id=None,
+        deployment_id=None,
+        before_offset=None,
+        after_offset=None,
+        max_number_of_results=None,
+        log_text_filter=None,
+    ) -> ReadLogsResponse:
+        try:
+            response = self.read_logs(
+                source=LogSource(
+                    model_runtime=ModelRuntimeSource(
+                        build_id=build_id, deployment_id=deployment_id
+                    )
+                ),
+                before_offset=before_offset,
+                after_offset=after_offset,
+                log_text_filter=log_text_filter,
+                max_number_of_results=max_number_of_results,
+            )
+
+            return response
+        except FrogmlException as e:
+            raise FrogmlException(f"Failed to fetch runtime logs, error is [{e}]")
+
+    def read_execution_models_logs(
+        self,
+        execution_id,
+        before_offset=None,
+        after_offset=None,
+        max_number_of_results=None,
+        log_text_filter=None,
+    ) -> ReadLogsResponse:
+        try:
+            response = self.read_logs(
+                source=LogSource(
+                    inference_execution=InferenceExecutionSource(
+                        inference_job_id=execution_id
+                    )
+                ),
+                before_offset=before_offset,
+                after_offset=after_offset,
+                log_text_filter=log_text_filter,
+                max_number_of_results=max_number_of_results,
+            )
+
+            return response
+        except FrogmlException as e:
+            raise FrogmlException(f"Failed to fetch execution logs, error is [{e}]")
+
+    def read_logs(
+        self,
+        source,
+        before_offset,
+        after_offset,
+        max_number_of_results,
+        log_text_filter,
+    ):
+        try:
+            response = self._logging_service.ReadLogs(
+                ReadLogsRequest(
+                    source=source,
+                    before_offset=before_offset,
+                    after_offset=after_offset,
+                    search_filter=SearchFilter(
+                        log_text_filter=LogText(contains=log_text_filter)
+                    ),
+                    max_number_of_results=max_number_of_results,
+                )
+            )
+            return response
+        except RpcError as e:
+            raise FrogmlException(
+                f"Failed grpc read logs request, grpc error is "
+                f"[{e.details() if e.details() else e.code()}]"
+            )