frogml-core 0.0.115__py3-none-any.whl → 0.0.117__py3-none-any.whl

frogml_services_mock/mocks/features_online_serving_api.py

@@ -1,169 +0,0 @@
-import json
-from datetime import datetime
-
-import grpc
-import pandas as pd
-
-from frogml_proto.qwak.feature_store.serving.serving_pb2 import (
-    MultiFeatureValuesResponse,
-)
-from frogml_proto.qwak.feature_store.serving.serving_pb2_grpc import (
-    ServingServiceServicer,
-)
-from frogml_core.exceptions import FrogmlException
-
-
-class NpEncoder(json.JSONEncoder):
-    def default(self, obj):
-        import numpy as np
-
-        if isinstance(obj, np.integer):
-            return int(obj)
-        if isinstance(obj, np.floating):
-            return float(obj)
-        if isinstance(obj, np.ndarray):
-            return obj.tolist()
-        if isinstance(obj, datetime):
-            return obj.isoformat()
-        return super(NpEncoder, self).default(obj)
-
-
-class ServingServiceMock(ServingServiceServicer):
-    def __init__(self):
-        self.features_by_entity_df = {}
-        self._fail_request_ledger = 0
-
-    def fail_next_request(self):
-        self._fail_request_ledger += 1
-
-    def upsert_features(self, entity_name, features_df):
-        if entity_name not in features_df.columns:
-            raise FrogmlException(
-                "Entity name does not exist in the data frame columns"
-            )
-
-        entity_feature_df = self.features_by_entity_df.get(
-            entity_name, pd.DataFrame(columns=[entity_name])
-        )
-
-        self.features_by_entity_df[entity_name] = (
-            features_df.set_index(entity_name)
-            .combine_first(entity_feature_df.set_index(entity_name))
-            .reset_index()
-            .rename({"index": entity_name})
-        )
-
-    def GetMultiFeatures(self, request, context) -> MultiFeatureValuesResponse:
-        if self._fail_request_ledger > 0:
-            self._fail_request_ledger -= 1
-            raise ValueError("Test failing feature request!")
-
-        if not self._validate_supported_features(request):
-            context.set_code(grpc.StatusCode.UNIMPLEMENTED)
-            context.set_details("Received an unimplemented one-of type for feature")
-            return MultiFeatureValuesResponse()
-
-        indexes, all_values, columns = [], [], []
-        for row_index, entity_values in enumerate(request.entity_values_matrix.rows):
-            values = []
-            for index, entity_name in enumerate(
-                request.entity_values_matrix.header.entity_names
-            ):
-                entity_feature_df = self.features_by_entity_df.get(entity_name)
-                if entity_feature_df is not None:
-                    all_features = self._filter_features(
-                        entity_feature_df, entity_name, entity_values, index
-                    )
-
-                    if len(all_features) > 1:
-                        context.set_code(grpc.StatusCode.INTERNAL)
-                        context.set_details(
-                            "Invalid state, a single entity has multiple feature rows - bug in the store :~("
-                        )
-                        return MultiFeatureValuesResponse()
-
-                    elif len(all_features) == 1:
-                        columns, values = self._populate_data(
-                            all_features, columns, entity_name, request, values
-                        )
-
-            indexes.append(row_index)
-            all_values.append(values)
-        result = {"data": all_values, "columns": columns, "index": indexes}
-
-        return MultiFeatureValuesResponse(
-            pandas_df_as_json=json.dumps(result, cls=NpEncoder)
-        )
-
-    def _populate_data(self, all_features, columns, entity_name, request, values):
-        requested_features = set()
-        self._populate_requested_features(entity_name, request, requested_features)
-        available_features = set(all_features.columns)
-        existing_features = requested_features.intersection(available_features)
-        non_existing_features = requested_features - available_features
-        values = values + (
-            [all_features.loc[0, feature_name] for feature_name in existing_features]
-        )
-
-        columns = self._populate_existing_values(columns, existing_features)
-        columns, values = self._populate_non_existing_values(
-            columns, non_existing_features, values
-        )
-        return columns, values
-
-    def _populate_existing_values(self, columns, existing_features):
-        if not any(x in columns for x in list(existing_features)):
-            columns = columns + (list(existing_features))
-        return columns
-
-    @staticmethod
-    def _populate_non_existing_values(columns, non_existing_features, values):
-        if not any(x in columns for x in list(non_existing_features)):
-            columns = columns + list(non_existing_features)
-            for _ in non_existing_features:
-                values = values + [None]
-        return columns, values
-
-    @staticmethod
-    def _populate_requested_features(entity_name, request, requested_features):
-        for entity_to_features in request.entities_to_features:
-            if entity_to_features.entity_name == entity_name:
-                for feature in entity_to_features.features:
-                    one_of = feature.WhichOneof("type")
-                    if one_of == "batch_feature":
-                        requested_features.add(feature.batch_feature.name)
-                    elif one_of == "batch_v1_feature":
-                        requested_features.add(feature.batch_v1_feature.name)
-                    elif one_of == "streaming_v1_feature":
-                        requested_features.add(feature.streaming_v1_feature.name)
-                    elif one_of == "streaming_feature":
-                        requested_features.add(feature.streaming_feature.name)
-                    elif one_of == "streaming_aggregation":
-                        requested_features.add(feature.streaming_aggregation.name)
-
-    @staticmethod
-    def _validate_supported_features(request):
-        return all(
-            [
-                (
-                    feature.WhichOneof("type") == "batch_v1_feature"
-                    or feature.WhichOneof("type") == "batch_feature"
-                    or feature.WhichOneof("type") == "streaming_feature"
-                    or feature.WhichOneof("type") == "streaming_feature_v1"
-                    or feature.WhichOneof("type") == "streaming_aggregation"
-                )
-                for entities in request.entities_to_features
-                for feature in entities.features
-            ]
-        )
-
-    @staticmethod
-    def _filter_features(entity_feature_df, entity_name, entity_values, index):
-        return (
-            entity_feature_df[
-                entity_feature_df[entity_name].astype(str)
-                == str(entity_values.entity_values[index])
-            ]
-            .drop([entity_name], axis=1)
-            .reset_index(drop=True)
-        )

frogml_services_mock/mocks/features_operator_v3_service.py

@@ -1,29 +0,0 @@
-import uuid
-from typing import Optional
-
-from frogml_proto.qwak.features_operator.v3.features_operator_async_service_pb2 import (
-    ValidationResponse,
-)
-from frogml_proto.qwak.features_operator.v3.features_operator_async_service_pb2_grpc import (
-    FeaturesOperatorAsyncServiceServicer,
-)
-
-
-class FeaturesOperatorV3ServiceMock(FeaturesOperatorAsyncServiceServicer):
-    def __init__(self):
-        self.response: Optional[ValidationResponse] = None
-        super(FeaturesOperatorV3ServiceMock, self).__init__()
-
-    def given_next_response(self, response: ValidationResponse):
-        self.response = response
-
-    def ValidateDataSource(self, request, context) -> str:
-        request_id = str(uuid.uuid4())
-        return ValidationResponse(request_id=request_id)
-
-    def ValidateFeatureSet(self, request, context):
-        request_id = str(uuid.uuid4())
-        return ValidationResponse(request_id=request_id)
-
-    def GetValidationResult(self, request, context):
-        return self.response

frogml_services_mock/mocks/features_set_state_service_api.py

@@ -1,57 +0,0 @@
-import grpc
-
-from frogml_proto.qwak.feature_store.features.feature_set_state_pb2 import (
-    FeatureSetState,
-)
-from frogml_proto.qwak.feature_store.features.feature_set_state_service_pb2 import (
-    GetFeatureSetStateByIdV1Response,
-    GetFeatureSetStateByNameV1Response,
-)
-from frogml_proto.qwak.feature_store.features.feature_set_state_service_pb2_grpc import (
-    FeatureSetStateServiceServicer,
-)
-from frogml_services_mock.mocks.utils.exception_handlers import (
-    raise_internal_grpc_error,
-)
-
-
-class FeaturesSetStateServiceApiMock(FeatureSetStateServiceServicer):
-    def __init__(self):
-        self._feature_state_states: dict[(str, str):FeatureSetState] = {}
-        super(FeaturesSetStateServiceApiMock, self).__init__()
-
-    def given_feature_set_state(self, feature_set_state: FeatureSetState):
-        self._feature_state_states[
-            (feature_set_state.feature_set_id, feature_set_state.run_id)
-        ] = feature_set_state
-
-    def GetFeatureSetStateById(self, request, context):
-        """get feature set state by id"""
-        try:
-            return GetFeatureSetStateByIdV1Response(
-                feature_set_state=self._feature_state_states[
-                    (request.feature_set_id, request.run_id)
-                ]
-            )
-        except Exception as e:
-            raise_internal_grpc_error(context, e)
-
-    def GetFeatureSetStateByName(self, request, context):
-        """get feature set state by name"""
-        try:
-            feature_sets_by_name = [
-                self._feature_state_states[(feature_set_id, run_id)]
-                for feature_set_id, run_id in self._feature_state_states.keys()
-                if self._feature_state_states[(feature_set_id, run_id)].feature_set_name
-                == request.feature_set_name
-            ]
-            if len(feature_sets_by_name) > 0:
-                return GetFeatureSetStateByNameV1Response(
-                    feature_set_state=feature_sets_by_name.pop()
-                )
-            else:
-                context.set_code(grpc.StatusCode.NOT_FOUND)
-                context.set_details("Received a non existing feature set name")
-                return GetFeatureSetStateByNameV1Response()
-        except Exception as e:
-            raise_internal_grpc_error(context, e)

frogml_services_mock/mocks/feedback_service.py

@@ -1,35 +0,0 @@
-from frogml_proto.qwak.inference.feedback.feedback_pb2 import (
-    ActualValuesRequest,
-    ActualValuesResponse,
-    ConfigureFeedbackRequest,
-    ConfigureFeedbackResponse,
-)
-from frogml_proto.qwak.inference.feedback.feedback_pb2_grpc import (
-    FeedbackServiceServicer,
-)
-from frogml_services_mock.mocks.utils.exception_handlers import (
-    raise_internal_grpc_error,
-)
-
-
-class FeedbackServiceMock(FeedbackServiceServicer):
-    def __init__(self):
-        super(FeedbackServiceMock, self).__init__()
-
-    def PostFeedback(
-        self, request: ActualValuesRequest, context
-    ) -> ActualValuesResponse:
-        try:
-            return ActualValuesResponse(status=ActualValuesResponse.ValueStatus.SUCCESS)
-        except Exception as e:
-            raise_internal_grpc_error(context, e)
-
-    def ConfigureFeedback(
-        self, request: ConfigureFeedbackRequest, context
-    ) -> ConfigureFeedbackResponse:
-        try:
-            return ConfigureFeedbackResponse(
-                status=ConfigureFeedbackResponse.ValueStatus.SUCCESS
-            )
-        except Exception as e:
-            raise_internal_grpc_error(context, e)

frogml_services_mock/mocks/file_versioning_service.py

@@ -1,66 +0,0 @@
-from collections import defaultdict
-from typing import Dict, List
-
-from frogml_proto.qwak.file_versioning.file_versioning_pb2 import FileTagSpec
-from frogml_proto.qwak.file_versioning.file_versioning_service_pb2 import (
-    GetModelFileTagsRequest,
-    GetModelFileTagsResponse,
-    RegisterFileTagRequest,
-    RegisterFileTagResponse,
-)
-from frogml_proto.qwak.file_versioning.file_versioning_service_pb2_grpc import (
-    FileVersioningManagementServiceServicer,
-)
-from frogml_services_mock.mocks.utils.exception_handlers import (
-    raise_internal_grpc_error,
-)
-
-
-class FileVersioningServiceMock(FileVersioningManagementServiceServicer):
-    def __init__(self):
-        super(FileVersioningServiceMock, self).__init__()
-        self.tags: Dict[str : List[FileTagSpec]] = defaultdict(list)
-
-    def RegisterFileTag(
-        self, request: RegisterFileTagRequest, context
-    ) -> RegisterFileTagResponse:
-        try:
-            self.tags[request.file_tag_spec.build_id].append(request.file_tag_spec)
-            return RegisterFileTagResponse()
-        except Exception as e:
-            raise_internal_grpc_error(context, e)
-
-    def GetModelFileTags(
-        self, request: GetModelFileTagsRequest, context
-    ) -> GetModelFileTagsResponse:
-        try:
-            file_tags_list = []
-
-            if not request.build_id:
-                for all_file_tags in self.tags.values():
-                    for file_tag in all_file_tags:
-                        file_tags_list.append(file_tag)
-            else:
-                for file_tag_by_build_id in self.tags[request.build_id]:
-                    if file_tag_by_build_id.model_id == request.model_id:
-                        file_tags_list.append(file_tag_by_build_id)
-
-            if request.filter:
-                file_tags_list_filtered = list(file_tags_list)
-                filter_type: str = request.filter.WhichOneof("filter")
-
-                if filter_type == "tag_contains":
-                    for file_tag in file_tags_list:
-                        if request.filter.tag_contains not in file_tag.tag:
-                            file_tags_list_filtered.remove(file_tag)
-
-                if filter_type == "tag_prefix":
-                    for file_tag in file_tags_list:
-                        if not file_tag.tag.startswith(request.filter.tag_prefix):
-                            file_tags_list_filtered.remove(file_tag)
-
-                file_tags_list = file_tags_list_filtered
-
-            return GetModelFileTagsResponse(file_tags=file_tags_list)
-        except Exception as e:
-            raise_internal_grpc_error(context, e)

frogml_services_mock/mocks/frogml_mocks.py

@@ -1,128 +0,0 @@
-import dataclasses
-
-from frogml_services_mock.mocks.alert_manager_service_api import (
-    AlertManagerServiceApiMock,
-)
-from frogml_services_mock.mocks.alert_registry_service_api import (
-    AlertsRegistryServiceApiMock,
-)
-from frogml_services_mock.mocks.analytics_api import AnalyticsApiMock
-from frogml_services_mock.mocks.audience_service_api import AudienceServiceApiMock
-from frogml_services_mock.mocks.authentication_service import AuthenticationServiceMock
-from frogml_services_mock.mocks.automation_management_service import (
-    AutomationManagementServiceMock,
-)
-from frogml_services_mock.mocks.autoscaling_service_api import AutoscalingServiceApiMock
-from frogml_services_mock.mocks.batch_job_manager_service import BatchJobManagerService
-from frogml_services_mock.mocks.build_management import BuildsManagementServiceMock
-from frogml_services_mock.mocks.build_orchestrator_build_api import (
-    BuildOrchestratorBuildApiMock,
-)
-from frogml_services_mock.mocks.build_orchestrator_build_settings_api import (
-    BuildOrchestratorBuildSettingsApiMock,
-)
-from frogml_services_mock.mocks.build_orchestrator_service_api import (
-    BuildOrchestratorServiceApiMock,
-)
-from frogml_services_mock.mocks.data_versioning_service import DataVersioningServiceMock
-from frogml_services_mock.mocks.deployment_management_service import (
-    DeploymentManagementServiceMock,
-)
-from frogml_services_mock.mocks.ecosystem_service_api import EcoSystemServiceMock
-from frogml_services_mock.mocks.execution_management_service import (
-    ExecutionManagementServiceMock,
-)
-from frogml_services_mock.mocks.feature_store_data_sources_manager_api import (
-    DataSourceServiceMock,
-)
-from frogml_services_mock.mocks.feature_store_entities_manager_api import (
-    EntityServiceMock,
-)
-from frogml_services_mock.mocks.feature_store_feature_set_manager_api import (
-    FeatureSetServiceMock,
-)
-from frogml_services_mock.mocks.features_operator_v3_service import (
-    FeaturesOperatorV3ServiceMock,
-)
-from frogml_services_mock.mocks.file_versioning_service import FileVersioningServiceMock
-from frogml_services_mock.mocks.fs_offline_serving_service import (
-    FsOfflineServingServiceMock,
-)
-from frogml_services_mock.mocks.instance_template_management_service import (
-    InstanceTemplateManagementServiceMock,
-)
-from frogml_services_mock.mocks.integration_management_service import (
-    IntegrationManagementServiceMock,
-)
-from frogml_services_mock.mocks.internal_build_orchestrator_service import (
-    InternalBuildOrchestratorServiceMock,
-)
-from frogml_services_mock.mocks.jfrog_tenant_info_service_mock import JFrogTenantInfoServiceMock
-from frogml_services_mock.mocks.job_registry_service_api import (
-    JobRegistryServiceApiMock,
-)
-from frogml_services_mock.mocks.kube_captain_service_api import (
-    KubeCaptainServiceApiMock,
-)
-from frogml_services_mock.mocks.location_discovery_service_api import (
-    LocationDiscoveryServiceApiMock,
-)
-from frogml_services_mock.mocks.logging_service import LoggingServiceApiMock
-from frogml_services_mock.mocks.model_management_service import (
-    ModelsManagementServiceMock,
-)
-from frogml_services_mock.mocks.model_version_manager_service import (
-    ModelVersionManagerServiceMock,
-)
-from frogml_services_mock.mocks.project_manager_service import ProjectManagerServiceMock
-from frogml_services_mock.mocks.repository_service_mock import RepositoryServiceMock
-from frogml_services_mock.mocks.secret_service import SecretServiceMock
-from frogml_services_mock.mocks.self_service_user_service import (
-    SelfServiceUserServiceMock,
-)
-from frogml_services_mock.mocks.system_secret_service import SystemSecretServiceMock
-from frogml_services_mock.mocks.user_application_instance_service_api import (
-    UserApplicationInstanceServiceApiMock,
-)
-
-@dataclasses.dataclass
-class FrogmlMocks:
-    port: int
-    build_orchestrator_build_api: BuildOrchestratorBuildApiMock
-    build_orchestrator_service_api: BuildOrchestratorServiceApiMock
-    build_orchestrator_build_settings_api: BuildOrchestratorBuildSettingsApiMock
-    internal_build_orchestrator_service: InternalBuildOrchestratorServiceMock
-    build_management_service_mock: BuildsManagementServiceMock
-    alert_manager_service_mock: AlertManagerServiceApiMock
-    automation_management_service_mock: AutomationManagementServiceMock
-    job_registry_service_mock: JobRegistryServiceApiMock
-    project_manager_service_mock: ProjectManagerServiceMock
-    kube_captain_service_mock: KubeCaptainServiceApiMock
-    file_versioning_service_mock: FileVersioningServiceMock
-    data_versioning_service_mock: DataVersioningServiceMock
-    model_management_service_mock: ModelsManagementServiceMock
-    logging_service_mock: LoggingServiceApiMock
-    autoscaling_service_mock: AutoscalingServiceApiMock
-    audience_api_mock: AudienceServiceApiMock
-    self_service_user_service_mock: SelfServiceUserServiceMock
-    analytics_api_mock: AnalyticsApiMock
-    deployment_management_service_mock: DeploymentManagementServiceMock
-    batch_job_manager_service: BatchJobManagerService
-    ecosystem_client_mock: EcoSystemServiceMock
-    user_application_instance_service_mock: UserApplicationInstanceServiceApiMock
-    secret_service_mock: SecretServiceMock
-    authentication_service_mock: AuthenticationServiceMock
-    features_operator_service: FeaturesOperatorV3ServiceMock
-    fs_offline_serving_service: FsOfflineServingServiceMock
-    fs_data_sources_service: DataSourceServiceMock
-    fs_entities_service: EntityServiceMock
-    fs_feature_sets_service: FeatureSetServiceMock
-    instance_templates_service: InstanceTemplateManagementServiceMock
-    alerts_registry_service: AlertsRegistryServiceApiMock
-    execution_management_service: ExecutionManagementServiceMock
-    system_secret_service: SystemSecretServiceMock
-    integration_management_service: IntegrationManagementServiceMock
-    model_version_manager_service: ModelVersionManagerServiceMock
-    repository_service: RepositoryServiceMock
-    location_discovery_service: LocationDiscoveryServiceApiMock
-    jfrog_tenant_info_service: JFrogTenantInfoServiceMock

frogml_services_mock/mocks/fs_offline_serving_service.py

@@ -1,53 +0,0 @@
-import uuid
-from typing import Optional
-
-from frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2 import (
-    GetFeatureValuesInRangeRequest,
-    GetFeatureValuesInRangeResponse,
-    GetFeatureValuesRequest,
-    GetFeatureValuesResponse,
-    GetFeatureValuesResultResponse,
-    GetFileUploadUrlResponse,
-)
-from frogml_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2_grpc import (
-    FeatureStoreOfflineServingAsyncServiceServicer,
-)
-from frogml_proto.qwak.offline.serving.v1.options_pb2 import (
-    OfflineServingQueryOptions as ProtoOfflineServingQueryOptions,
-)
-
-
-class FsOfflineServingServiceMock(FeatureStoreOfflineServingAsyncServiceServicer):
-    def __init__(self):
-        self.file_upload_url_response: Optional[GetFileUploadUrlResponse] = None
-        self.response: Optional[GetFeatureValuesResultResponse] = None
-        self.latest_query_options: Optional[ProtoOfflineServingQueryOptions] = None
-        super(FsOfflineServingServiceMock, self).__init__()
-
-    def given_next_file_upload_url(
-        self, file_upload_url_response: GetFileUploadUrlResponse
-    ):
-        self.file_upload_url_response = file_upload_url_response
-
-    def given_next_response(self, response: GetFeatureValuesResultResponse):
-        self.response = response
-
-    def GetFileUploadUrl(self, request, context) -> Optional[GetFileUploadUrlResponse]:
-        return self.file_upload_url_response
-
-    def GetFeatureValues(
-        self, request: GetFeatureValuesRequest, context
-    ) -> GetFeatureValuesResponse:
-        request_id = str(uuid.uuid4())
-        self.latest_query_options = request.options
-        return GetFeatureValuesResponse(request_id=request_id)
-
-    def GetFeatureValuesInRange(self, request: GetFeatureValuesInRangeRequest, context):
-        request_id = str(uuid.uuid4())
-        self.latest_query_options = request.options
-        return GetFeatureValuesInRangeResponse(request_id=request_id)
-
-    def GetFeatureValuesResult(
-        self, request, context
-    ) -> Optional[GetFeatureValuesResultResponse]:
-        return self.response

frogml_services_mock/mocks/instance_template_management_service.py

@@ -1,150 +0,0 @@
-from frogml_proto.qwak.instance_template.instance_template_pb2 import (
-    InstanceTemplateSpec,
-    InstanceType,
-)
-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 (
-    InstanceTemplateManagementServiceServicer,
-)
-from frogml_proto.qwak.user_application.common.v0.resources_pb2 import (
-    CpuResources,
-    GpuResources,
-    GpuType,
-    MemoryUnit,
-    NodeOptimizationType,
-    PodComputeResources,
-)
-from frogml_core.exceptions import FrogmlException
-from frogml_services_mock.mocks.utils.exception_handlers import (
-    raise_internal_grpc_error,
-)
-
-
-def _create_instance_template_cpu(
-    instance_id: str,
-    name: str,
-    order: int,
-    cpu: float,
-    memory_amount: int,
-    aws_supported: bool = True,
-    gcp_supported: bool = True,
-) -> InstanceTemplateSpec:
-    return InstanceTemplateSpec(
-        id=instance_id,
-        display_name=name,
-        order=order,
-        enabled=True,
-        aws_supported=aws_supported,
-        gcp_supported=gcp_supported,
-        pod_compute_resources=PodComputeResources(
-            optimization_type=NodeOptimizationType.NODE_OPTIMIZATION_NONE,
-            cpu_resources=CpuResources(
-                cpu=cpu,
-                memory_amount=memory_amount,
-                memory_units=MemoryUnit.GIB,
-            ),
-        ),
-        instance_type=InstanceType.INSTANCE_TYPE_CPU,
-    )
-
-
-def _create_instance_template_gpu(
-    instance_id: str,
-    name: str,
-    order: int,
-    gpu: int,
-    gpu_type: GpuType,
-    cpu: float,
-    memory_amount: int,
-    aws_supported: bool = True,
-    gcp_supported: bool = True,
-) -> InstanceTemplateSpec:
-    return InstanceTemplateSpec(
-        id=instance_id,
-        display_name=name,
-        order=order,
-        enabled=True,
-        aws_supported=aws_supported,
-        gcp_supported=gcp_supported,
-        pod_compute_resources=PodComputeResources(
-            optimization_type=NodeOptimizationType.NODE_OPTIMIZATION_NONE,
-            gpu_resources=GpuResources(
-                gpu_amount=gpu,
-                gpu_type=gpu_type,
-            ),
-            cpu_resources=CpuResources(
-                cpu=cpu,
-                memory_amount=memory_amount,
-                memory_units=MemoryUnit.GIB,
-            ),
-        ),
-        instance_type=InstanceType.INSTANCE_TYPE_GPU,
-    )
-
-
-class InstanceTemplateManagementServiceMock(InstanceTemplateManagementServiceServicer):
-    INSTANCES = {
-        "small": _create_instance_template_cpu("small", "Small", 1, 1, 8),
-        "medium": _create_instance_template_cpu("medium", "Medium", 2, 2, 16),
-        "large": _create_instance_template_cpu("large", "Large", 3, 4, 32),
-        "xlarge": _create_instance_template_cpu("xlarge", "XLarge", 4, 8, 64),
-        "g5_xlarge": _create_instance_template_gpu(
-            "g5_xlarge", "G5 XLarge", 1, 1, GpuType.NVIDIA_A10G, 1, 8
-        ),
-        "g5_2xlarge": _create_instance_template_gpu(
-            "g5_2xlarge", "G5 2XLarge", 2, 2, GpuType.NVIDIA_A10G, 2, 16
-        ),
-        "g5_4xlarge": _create_instance_template_gpu(
-            "g5_4xlarge", "G5 4XLarge", 2, 4, GpuType.NVIDIA_A10G, 4, 32
-        ),
-        "aws_only": _create_instance_template_gpu(
-            "aws_only",
-            "Aws Only",
-            2,
-            4,
-            GpuType.NVIDIA_A10G,
-            4,
-            32,
-            aws_supported=True,
-            gcp_supported=False,
-        ),
-        "gcp_only": _create_instance_template_gpu(
-            "gcp_only",
-            "GCP Only",
-            2,
-            4,
-            GpuType.NVIDIA_A10G,
-            4,
-            32,
-            aws_supported=False,
-            gcp_supported=True,
-        ),
-    }
-
-    def ListInstanceTemplates(
-        self, request: ListInstanceTemplatesRequest, context
-    ) -> ListInstanceTemplatesResponse:
-        try:
-            return ListInstanceTemplatesResponse(
-                instance_template_list=list(self.INSTANCES.values())
-            )
-        except Exception as e:
-            raise_internal_grpc_error(context, e)
-
-    def GetInstanceTemplate(
-        self, request: GetInstanceTemplateRequest, context
-    ) -> GetInstanceTemplateResponse:
-        try:
-            if request.id in self.INSTANCES:
-                return GetInstanceTemplateResponse(
-                    instance_template=self.INSTANCES[request.id]
-                )
-            else:
-                raise FrogmlException("Instance template not found")
-        except Exception as e:
-            raise_internal_grpc_error(context, e)