frogml-core 0.0.1__py3-none-any.whl

frogml/clients/integration_management/openai/openai_system_secret.py

@@ -0,0 +1,56 @@
+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.clients.administration.eco_system.eco_system_utils import EcosystemUtils
+from frogml.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")

frogml/clients/jfrog_gateway/__init__.py

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

frogml/clients/jfrog_gateway/client.py

@@ -0,0 +1,42 @@
+from typing import cast
+
+from dependency_injector.wiring import Provide
+
+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,
+)
+from frogml.exceptions import FrogmlException
+from frogml.inner.di_configuration import FrogmlContainer
+
+
+class JfrogGatewayClient:
+    """
+    Used for interacting with Feature Registry endpoints
+    """
+
+    def __init__(self, grpc_channel=Provide[FrogmlContainer.core_grpc_channel]):
+        self.__repository_service = RepositoryServiceStub(grpc_channel)
+
+    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)
+
+        try:
+            return cast(
+                GetRepositoryConfigurationResponse,
+                self.__repository_service.GetRepositoryConfiguration(request),
+            )
+        except Exception as e:
+            raise FrogmlException(
+                f"Failed to get repository configuration: {e.details()}"
+            )

frogml/clients/kube_deployment_captain/__init__.py

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

frogml/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.exceptions import FrogmlException
+from frogml.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)}"
+            )

frogml/clients/logging_client/__init__.py

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

frogml/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.clients.administration.eco_system.client import EcosystemClient
+from frogml.exceptions import FrogmlException
+from frogml.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()}]"
+            )

frogml/clients/model_management/__init__.py

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

frogml/clients/model_management/client.py

@@ -0,0 +1,128 @@
+from typing import Optional
+
+import grpc
+from dependency_injector.wiring import Provide
+
+from _frogml_proto.qwak.models.models_pb2 import (
+    CreateModelRequest,
+    DeleteModelRequest,
+    GetModelMetadataRequest,
+    GetModelRequest,
+    ListModelsMetadataRequest,
+    ListModelsRequest,
+    Model,
+    ModelSpec,
+)
+from _frogml_proto.qwak.models.models_pb2_grpc import ModelsManagementServiceStub
+from frogml.exceptions import FrogmlException
+from frogml.inner.di_configuration import FrogmlContainer
+
+
+class ModelsManagementClient:
+    """
+    Used for interacting with Feature Registry endpoints
+    """
+
+    def __init__(self, grpc_channel=Provide[FrogmlContainer.core_grpc_channel]):
+        self._models_management_service = ModelsManagementServiceStub(grpc_channel)
+
+    def get_model(self, model_id, exception_on_missing: bool = True) -> Optional[Model]:
+        try:
+            return self._models_management_service.GetModel(
+                GetModelRequest(model_id=model_id)
+            ).model
+
+        except grpc.RpcError as e:
+            if e.code() == grpc.StatusCode.NOT_FOUND and not exception_on_missing:
+                return None
+            raise FrogmlException(f"Failed to get model, error is {e.details()}")
+
+    def get_model_by_uuid(
+        self, model_uuid, exception_on_missing: bool = True
+    ) -> Optional[Model]:
+        try:
+            return self._models_management_service.GetModel(
+                GetModelRequest(model_uuid=model_uuid)
+            ).model
+
+        except grpc.RpcError as e:
+            if e.code() == grpc.StatusCode.NOT_FOUND and not exception_on_missing:
+                return None
+            raise FrogmlException(f"Failed to get model, error is {e.details()}")
+
+    def get_model_uuid(self, model_id):
+        model = self.get_model(model_id)
+        return model.uuid
+
+    def create_model(self, project_id, model_name, model_description):
+        try:
+            return self._models_management_service.CreateModel(
+                CreateModelRequest(
+                    model_spec=ModelSpec(
+                        model_named_id=model_name.lower().replace(" ", "-"),
+                        display_name=model_name,
+                        project_id=project_id,
+                        model_description=model_description,
+                    )
+                )
+            )
+
+        except grpc.RpcError as e:
+            raise FrogmlException(f"Failed to create model, error is {e}")
+
+    def delete_model(self, project_id, model_id):
+        try:
+            return self._models_management_service.DeleteModel(
+                DeleteModelRequest(model_id=model_id, project_id=project_id)
+            )
+
+        except grpc.RpcError as e:
+            raise FrogmlException(f"Failed to delete model, error is {e.details()}")
+
+    def is_model_exists(self, model_id: str) -> bool:
+        """Check if model exists in environment
+
+        Args:
+            model_id: the model id to check if exists.
+
+        Returns: if model exists.
+        """
+        try:
+            self._models_management_service.GetModel(
+                GetModelRequest(model_id=model_id)
+            ).model
+            return True
+        except grpc.RpcError as e:
+            if e.code() == grpc.StatusCode.NOT_FOUND:
+                return False
+            raise FrogmlException(
+                f"Failed to check if model {model_id} is exists, error is {e.details()}"
+            )
+
+    def list_models_metadata(self, project_id: str):
+        try:
+            return self._models_management_service.ListModelsMetadata(
+                ListModelsMetadataRequest(project_id=project_id)
+            )
+        except grpc.RpcError as e:
+            raise FrogmlException(
+                f"Failed to list models metadata, error is {e.details()}"
+            )
+
+    def list_models(self, project_id: str):
+        try:
+            return self._models_management_service.ListModels(
+                ListModelsRequest(project_id=project_id)
+            )
+        except grpc.RpcError as e:
+            raise FrogmlException(f"Failed to list models, error is {e.details()}")
+
+    def get_model_metadata(self, model_id: str):
+        try:
+            return self._models_management_service.GetModelMetadata(
+                GetModelMetadataRequest(model_id=model_id)
+            )
+        except grpc.RpcError as e:
+            raise FrogmlException(
+                f"Failed to get model metadata, error is {e.details()}"
+            )

frogml/clients/model_version_manager/__init__.py

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