qwak-core 0.1.37__py3-none-any.whl → 0.5.19__py3-none-any.whl

qwak_services_mock/mocks/feature_store_feature_set_manager_api.py

@@ -1,47 +1,134 @@
 import uuid
+from typing import Dict, Optional
 
 import grpc
 from _qwak_proto.qwak.feature_store.features.feature_set_pb2 import (
     DeployedModelInUseLink,
     FeatureSet,
     FeatureSetDefinition,
+    FeaturesetInfo,
+    FeatureSetList,
     FeatureSetMetadata,
     FeaturesetSchedulingState,
+    FeatureSetSpec,
     FeatureStatus,
 )
 from _qwak_proto.qwak.feature_store.features.feature_set_service_pb2 import (
     DeleteFeatureSetResponse,
+    DeleteFeaturesetVersionResponse,
+    GetEnvToFeatureSetsMappingResponse,
     GetFeatureSetByNameResponse,
     GetFeaturesetSchedulingStateResponse,
     ListFeatureSetsResponse,
+    ListFeaturesetVersionsByNameResponse,
     RegisterFeatureSetResponse,
+    SetActiveFeaturesetVersionResponse,
+    UpdateFeatureSetResponse,
 )
 from _qwak_proto.qwak.feature_store.features.feature_set_service_pb2_grpc import (
     FeatureSetServiceServicer,
 )
 
+CURRENT_ENV_NAME = "env_name"
+
 
 class FeatureSetServiceMock(FeatureSetServiceServicer):
     def __init__(self):
-        self._features_spec = {}
+        self._features_spec: Dict[str, FeatureSetSpec] = {}
+        self._fs_next_versions: Dict[str, int] = {}
+        self._fs_versions_by_name: Dict[str, Dict[int, FeatureSetDefinition]] = {}
+        self._fs_definition_by_id: Dict[str, FeatureSetDefinition] = {}
+        # this mapping is used to test the GetEnvToFeatureSetsMapping API. Note that the populated features sets may not be aligned with the self._features_spec featuresets
+        self._env_to_features_spec = (
+            {}
+        )  # env_name -> {featureset_name -> featureset_definition}
 
     def reset_features(self):
         self._features_spec.clear()
+        self._fs_next_versions.clear()
+        self._fs_versions_by_name.clear()
+        self._fs_definition_by_id.clear()
+        self._env_to_features_spec.clear()
+
+    def setup_empty_env(self, env_name):
+        self._env_to_features_spec[env_name] = {}
+
+    def register_featureset_to_env_to_features_map(
+        self, env_name, featureset_spec: FeatureSetSpec
+    ):
+        fs_id = str(uuid.uuid4())
+        self._env_to_features_spec[env_name] = self._env_to_features_spec.get(
+            env_name, {}
+        )
+        feature_set_definition = FeatureSetDefinition(
+            feature_set_id=fs_id,
+            feature_set_spec=featureset_spec,
+            status=FeatureStatus.VALID,
+        )
+        self._env_to_features_spec[env_name][
+            featureset_spec.name
+        ] = feature_set_definition
 
     def RegisterFeatureSet(self, request, context):
         fs_id = str(uuid.uuid4())
-        self._features_spec[fs_id] = request.feature_set_spec
+        curr_version = 1
+        featureset_definition: FeatureSetDefinition = FeatureSetDefinition(
+            feature_set_id=fs_id,
+            feature_set_spec=request.feature_set_spec,
+            featureset_version_number=curr_version,
+            status=FeatureStatus.VALID,
+        )
+        self._features_spec[fs_id] = featureset_definition.feature_set_spec
+        self._fs_definition_by_id[fs_id] = featureset_definition
+        self._fs_versions_by_name[featureset_definition.feature_set_spec.name] = {
+            curr_version: featureset_definition
+        }
+        self._fs_next_versions[featureset_definition.feature_set_spec.name] = (
+            curr_version + 1
+        )
+
         return RegisterFeatureSetResponse(
-            feature_set=FeatureSet(
-                feature_set_definition=FeatureSetDefinition(
-                    feature_set_id=fs_id, feature_set_spec=request.feature_set_spec
-                )
-            )
+            feature_set=FeatureSet(feature_set_definition=featureset_definition)
         )
 
+    def UpdateFeatureSet(self, request, context):
+        fs_id: str = request.feature_set_id
+        if fs_id in self._fs_definition_by_id:
+            curr_version = self._fs_next_versions[request.feature_set_spec.name]
+            featureset_definition: FeatureSetDefinition = FeatureSetDefinition(
+                feature_set_id=fs_id,
+                feature_set_spec=request.feature_set_spec,
+                featureset_version_number=curr_version,
+                status=FeatureStatus.VALID,
+            )
+            self._fs_definition_by_id[fs_id] = featureset_definition
+            self._features_spec[fs_id] = featureset_definition.feature_set_spec
+            self._fs_next_versions[featureset_definition.feature_set_spec.name] = (
+                curr_version + 1
+            )
+
+            curr_versions = self._fs_versions_by_name[request.feature_set_spec.name]
+            curr_versions[curr_version] = featureset_definition
+            self._fs_versions_by_name[
+                featureset_definition.feature_set_spec.name
+            ] = curr_versions
+
+            return UpdateFeatureSetResponse(
+                feature_set=FeatureSet(feature_set_definition=featureset_definition)
+            )
+
+        context.set_details(f"Feature set ID {request.feature_set_id} doesn't exist'")
+        context.set_code(grpc.StatusCode.NOT_FOUND)
+
     def DeleteFeatureSet(self, request, context):
-        if request.feature_set_id in self._features_spec:
+        fs_definition: Optional[FeatureSetDefinition] = self._fs_definition_by_id.get(
+            request.feature_set_id, None
+        )
+        if fs_definition:
             self._features_spec.pop(request.feature_set_id)
+            self._fs_definition_by_id.pop(request.feature_set_id)
+            self._fs_versions_by_name.pop(fs_definition.feature_set_spec.name)
+            self._fs_next_versions.pop(fs_definition.feature_set_spec.name)
             return DeleteFeatureSetResponse()
 
         context.set_details(f"Feature set ID {request.feature_set_id} doesn't exist'")
@@ -49,19 +136,15 @@ class FeatureSetServiceMock(FeatureSetServiceServicer):
 
     def GetFeatureSetByName(self, request, context):
         feature_sets = [
-            (fs_id, fs_spec)
-            for (fs_id, fs_spec) in self._features_spec.items()
-            if fs_spec.name == request.feature_set_name
+            (fs_id, fs_definition)
+            for (fs_id, fs_definition) in self._fs_definition_by_id.items()
+            if fs_definition.feature_set_spec.name == request.feature_set_name
         ]
         if feature_sets:
-            fs_id, fs_spec = feature_sets[0]
+            fs_id, fs_definition = feature_sets[0]
             return GetFeatureSetByNameResponse(
                 feature_set=FeatureSet(
-                    feature_set_definition=FeatureSetDefinition(
-                        feature_set_id=fs_id,
-                        feature_set_spec=fs_spec,
-                        status=FeatureStatus.VALID,
-                    ),
+                    feature_set_definition=fs_definition,
                     metadata=FeatureSetMetadata(),
                     deployed_models_in_use_link=[DeployedModelInUseLink()],
                 )
@@ -76,13 +159,11 @@ class FeatureSetServiceMock(FeatureSetServiceServicer):
         return ListFeatureSetsResponse(
             feature_families=[
                 FeatureSet(
-                    feature_set_definition=FeatureSetDefinition(
-                        feature_set_spec=feature_spec, status=FeatureStatus.VALID
-                    ),
+                    feature_set_definition=fs_definition,
                     metadata=FeatureSetMetadata(),
                     deployed_models_in_use_link=[DeployedModelInUseLink()],
                 )
-                for feature_spec in self._features_spec.values()
+                for fs_definition in self._fs_definition_by_id.values()
             ]
         )
 
@@ -90,3 +171,98 @@ class FeatureSetServiceMock(FeatureSetServiceServicer):
         return GetFeaturesetSchedulingStateResponse(
             state=FeaturesetSchedulingState.SCHEDULING_STATE_ENABLED
         )
+
+    def GetEnvToFeatureSetsMapping(self, request, context):
+        env_to_featuresets_mapping = {}
+
+        for env_name, featureset_df_dict in self._env_to_features_spec.items():
+            env_to_featuresets_mapping[env_name] = FeatureSetList(
+                feature_sets=[
+                    FeatureSet(
+                        feature_set_definition=fs_definition,
+                        metadata=FeatureSetMetadata(),
+                        deployed_models_in_use_link=[DeployedModelInUseLink()],
+                    )
+                    for fs_name, fs_definition in featureset_df_dict.items()
+                ]
+            )
+
+        return GetEnvToFeatureSetsMappingResponse(
+            env_to_feature_set_mapping=env_to_featuresets_mapping
+        )
+
+    def DeleteFeaturesetVersion(self, request, context):
+        version_number = request.featureset_version_number
+        version_definitions: Optional[
+            Dict[int, FeatureSetDefinition]
+        ] = self._fs_versions_by_name.get(request.featureset_name, None)
+        if version_definitions:
+            version_definition = version_definitions.get(version_number, None)
+            if version_definition:
+                if (
+                    self._fs_definition_by_id[
+                        version_definition.feature_set_id
+                    ].featureset_version_number
+                    != version_number
+                ):
+                    version_definitions.pop(version_number)
+                    return DeleteFeaturesetVersionResponse()
+                else:
+                    context.set_details("Active version cant be deleted")
+                    context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
+            else:
+                context.set_details(
+                    f"Version {version_number} not found for Feature set name {request.featureset_name}"
+                )
+                context.set_code(grpc.StatusCode.NOT_FOUND)
+        else:
+            context.set_details(
+                f"Feature set name {request.featureset_name} doesn't exist'"
+            )
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+
+    def SetActiveFeaturesetVersion(self, request, context):
+        version_definitions: Optional[
+            Dict[int, FeatureSetDefinition]
+        ] = self._fs_versions_by_name.get(request.featureset_name, None)
+        if version_definitions:
+            version_definition = version_definitions.get(
+                request.featureset_version_number, None
+            )
+            if version_definition:
+                self._fs_definition_by_id[
+                    version_definition.feature_set_id
+                ] = version_definition
+                return SetActiveFeaturesetVersionResponse()
+            else:
+                context.set_details(
+                    f"Version {request.featureset_version_number} not found for Feature set name {request.featureset_name}"
+                )
+                context.set_code(grpc.StatusCode.NOT_FOUND)
+        else:
+            context.set_details(
+                f"Feature set name {request.featureset_name} doesn't exist'"
+            )
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+
+    def ListFeaturesetVersionsByName(self, request, context):
+        versions: Optional[
+            Dict[int, FeatureSetDefinition]
+        ] = self._fs_versions_by_name.get(request.featureset_name, None)
+        if versions:
+            return ListFeaturesetVersionsByNameResponse(
+                featuresets=[
+                    FeaturesetInfo(
+                        featureset_definition=feature_definition,
+                        metadata=FeatureSetMetadata(),
+                    )
+                    for feature_definition in self._fs_versions_by_name[
+                        request.featureset_name
+                    ].values()
+                ]
+            )
+
+        context.set_details(
+            f"Feature set name {request.featureset_name} doesn't exist'"
+        )
+        context.set_code(grpc.StatusCode.NOT_FOUND)

qwak_services_mock/mocks/features_online_serving_api.py

@@ -131,6 +131,12 @@ class ServingServiceMock(ServingServiceServicer):
                         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):
@@ -139,6 +145,9 @@ class ServingServiceMock(ServingServiceServicer):
                 (
                     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

qwak_services_mock/mocks/features_operator_v3_service.py

@@ -1,6 +1,7 @@
 import uuid
 from typing import Optional
 
+import grpc
 from _qwak_proto.qwak.features_operator.v3.features_operator_async_service_pb2 import (
     ValidationResponse,
 )
@@ -18,12 +19,31 @@ class FeaturesOperatorV3ServiceMock(FeaturesOperatorAsyncServiceServicer):
         self.response = response
 
     def ValidateDataSource(self, request, context) -> str:
+        if request.num_samples <= 0 or 1_000 < request.num_samples:
+            context.set_details(
+                f"[NumberOfSamples must be between 0 and 1000, got {request.num_samples}]"
+            )
+            context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
+            return
+
         request_id = str(uuid.uuid4())
         return ValidationResponse(request_id=request_id)
 
     def ValidateFeatureSet(self, request, context):
+        if request.num_samples <= 0 or 1_000 < request.num_samples:
+            context.set_details(
+                f"[NumberOfSamples must be between 0 and 1000, got {request.num_samples}]"
+            )
+            context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
+            return
+
         request_id = str(uuid.uuid4())
         return ValidationResponse(request_id=request_id)
 
     def GetValidationResult(self, request, context):
+        if not request.request_id:
+            context.set_details("Request handle id not found")
+            context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
+            return
+
         return self.response

qwak_services_mock/mocks/features_set_state_service_api.py

@@ -14,7 +14,6 @@ from qwak_services_mock.mocks.utils.exception_handlers import raise_internal_grp
 
 class FeaturesSetStateServiceApiMock(FeatureSetStateServiceServicer):
     def __init__(self):
-
         self._feature_state_states: dict[(str, str):FeatureSetState] = {}
         super(FeaturesSetStateServiceApiMock, self).__init__()
 

qwak_services_mock/mocks/file_versioning_service.py

@@ -23,7 +23,7 @@ class FileVersioningServiceMock(FileVersioningManagementServiceServicer):
         self, request: RegisterFileTagRequest, context
     ) -> RegisterFileTagResponse:
         try:
-            self.tags[request.file_tag_spec.model_id].append(request.file_tag_spec)
+            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)
@@ -32,6 +32,33 @@ class FileVersioningServiceMock(FileVersioningManagementServiceServicer):
         self, request: GetModelFileTagsRequest, context
     ) -> GetModelFileTagsResponse:
         try:
-            return GetModelFileTagsResponse(file_tags=self.tags[request.model_id])
+            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)

qwak_services_mock/mocks/fs_offline_serving_service.py

@@ -0,0 +1,53 @@
+import uuid
+from typing import Optional
+
+from _qwak_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2 import (
+    GetFeatureValuesInRangeRequest,
+    GetFeatureValuesInRangeResponse,
+    GetFeatureValuesRequest,
+    GetFeatureValuesResponse,
+    GetFeatureValuesResultResponse,
+    GetFileUploadUrlResponse,
+)
+from _qwak_proto.qwak.offline.serving.v1.offline_serving_async_service_pb2_grpc import (
+    FeatureStoreOfflineServingAsyncServiceServicer,
+)
+from _qwak_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

qwak_services_mock/mocks/instance_template_management_service.py

@@ -24,13 +24,21 @@ from qwak_services_mock.mocks.utils.exception_handlers import raise_internal_grp
 
 
 def _create_instance_template_cpu(
-    instance_id: str, name: str, order: int, cpu: float, memory_amount: int
+    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(
@@ -51,12 +59,16 @@ def _create_instance_template_gpu(
     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(
@@ -88,6 +100,28 @@ class InstanceTemplateManagementServiceMock(InstanceTemplateManagementServiceSer
         "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(

qwak_services_mock/mocks/integration_management_service.py

@@ -0,0 +1,72 @@
+import uuid
+from typing import Dict
+
+import grpc
+from _qwak_proto.qwak.integration.integration_pb2 import Integration, IntegrationSpec
+from _qwak_proto.qwak.integration.integration_service_pb2 import (
+    CreateIntegrationRequest,
+    CreateIntegrationResponse,
+    DeleteIntegrationRequest,
+    DeleteIntegrationResponse,
+    GetIntegrationRequest,
+    GetIntegrationResponse,
+    ListIntegrationRequest,
+    ListIntegrationsResponse,
+)
+from _qwak_proto.qwak.integration.integration_service_pb2_grpc import (
+    IntegrationManagementServiceServicer,
+)
+from _qwak_proto.qwak.integration.open_a_i_integration_pb2 import (
+    OpenAIApiKeySystemSecretDescriptor,
+    OpenAIIntegration,
+)
+
+
+class IntegrationManagementServiceMock(IntegrationManagementServiceServicer):
+    _integrations: Dict[str, Integration] = dict()
+
+    def clear(self):
+        self._integrations.clear()
+
+    def CreateIntegration(self, request: CreateIntegrationRequest, context):
+        """Create a new integration in account level"""
+        spec: IntegrationSpec = request.integration_spec
+        if spec.WhichOneof("spec") != "openai_integration_spec":
+            context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+            context.set_details("Method not implemented!")
+            raise NotImplementedError("Method not implemented!")
+
+        integration_id: str = str(uuid.uuid4())
+        integration_name: str = str(uuid.uuid4())
+        secret_name: str = str(uuid.uuid4())
+        api_key_secret_key: str = str(uuid.uuid4())
+
+        integration = Integration(
+            integration_id=integration_id,
+            name=integration_name,
+            openai_integration=OpenAIIntegration(
+                openai_api_key_system_secret_descriptor=OpenAIApiKeySystemSecretDescriptor(
+                    secret_name=secret_name, api_key_secret_key=api_key_secret_key
+                )
+            ),
+        )
+        self._integrations[integration_id] = integration
+        return CreateIntegrationResponse(
+            integration_id=integration_id, is_successful=True
+        )
+
+    def GetIntegration(self, request: GetIntegrationRequest, context):
+        """Get integration by id"""
+        return GetIntegrationResponse(
+            integration=self._integrations[request.integration_id]
+        )
+
+    def ListIntegrations(self, request: ListIntegrationRequest, context):
+        """List Integrations for account"""
+        return ListIntegrationsResponse(integrations=self._integrations.values())
+
+    def DeleteIntegration(self, request: DeleteIntegrationRequest, context):
+        """Delete Integration by id"""
+        if request.integration_id in self._integrations:
+            del self._integrations[request.integration_id]
+        return DeleteIntegrationResponse()