frogml-core 0.0.1__py3-none-any.whl

frogml_services_mock/mocks/user_application_instance_service_api.py

@@ -0,0 +1,107 @@
+from _frogml_proto.qwak.admiral.user_application_instance.v0.user_application_instance_pb2 import (
+    Description,
+    Identifier,
+    Metadata,
+    SpecDescription,
+)
+from _frogml_proto.qwak.admiral.user_application_instance.v0.user_application_instance_service_pb2 import (
+    CreateUserApplicationInstanceRequest,
+    CreateUserApplicationInstanceResponse,
+    DeleteUserApplicationInstanceRequest,
+    DeleteUserApplicationInstanceResponse,
+    GetUserApplicationInstanceRequest,
+    GetUserApplicationInstanceResponse,
+    UpdateUserApplicationInstanceSpecRequest,
+    UpdateUserApplicationInstanceSpecResponse,
+)
+from _frogml_proto.qwak.admiral.user_application_instance.v0.user_application_instance_service_pb2_grpc import (
+    UserApplicationInstanceServiceServicer,
+)
+from frogml_services_mock.mocks.utils.exception_handlers import (
+    raise_internal_grpc_error,
+)
+
+
+class UserApplicationInstanceServiceApiMock(UserApplicationInstanceServiceServicer):
+    def __init__(self):
+        self.revisions = {}
+        super(UserApplicationInstanceServiceApiMock, self).__init__()
+
+    @staticmethod
+    def _make_key_from_identifier(identifier: Identifier):
+        return (
+            identifier.name
+            + "|"
+            + str(identifier.type)
+            + "|"
+            + identifier.account_id
+            + "|"
+            + identifier.environment_id
+        )
+
+    @staticmethod
+    def _get_identifier_from_key(key: str) -> Identifier:
+        splitted_key = key.split("|")
+        return Identifier(
+            name=splitted_key[0],
+            type=int(splitted_key[1]),
+            account_id=splitted_key[2],
+            environment_id=splitted_key[3],
+        )
+
+    def CreateUserApplicationInstance(
+        self, request: CreateUserApplicationInstanceRequest, context
+    ):
+        try:
+            key = self._make_key_from_identifier(request.create_options.identifier)
+            spec = self.revisions.get(key)
+            if spec is not None:
+                raise Exception("User application already exists")
+            self.revisions[key] = request.create_options.spec
+            return CreateUserApplicationInstanceResponse()
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def UpdateUserApplicationInstanceSpec(
+        self, request: UpdateUserApplicationInstanceSpecRequest, context
+    ):
+        try:
+            key = self._make_key_from_identifier(request.update_options.identifier)
+            revision = self.revisions.get(key)
+            if revision is None:
+                raise Exception("Get Update request but no revision was found")
+            self.revisions[key] = request.update_options.spec
+            return UpdateUserApplicationInstanceSpecResponse()
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def GetUserApplicationInstance(
+        self, request: GetUserApplicationInstanceRequest, context
+    ):
+        try:
+            key = self._make_key_from_identifier(request.identifier)
+            spec = self.revisions.get(key)
+            if spec is None:
+                return GetUserApplicationInstanceResponse()
+            description = Description(
+                metadata=Metadata(
+                    identifier=self._get_identifier_from_key(key=key),
+                ),
+                spec_description=(SpecDescription(spec=spec)),
+            )
+            return GetUserApplicationInstanceResponse(description=description)
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def DeleteUserApplicationInstance(
+        self, request: DeleteUserApplicationInstanceRequest, context
+    ):
+        try:
+            key = self._make_key_from_identifier(request.identifier)
+            spec = self.revisions[key]
+            if spec is None:
+                raise Exception(f"Could not find revision by the identifier key: {key}")
+            self.revisions[key] = None
+            return DeleteUserApplicationInstanceResponse()
+        except Exception as e:
+            raise_internal_grpc_error(context, e)

frogml_services_mock/mocks/utils/exception_handlers.py

@@ -0,0 +1,13 @@
+import grpc
+
+
+def raise_internal_grpc_error(context, e: Exception):
+    context.set_code(grpc.StatusCode.INTERNAL)
+    context.set_details(str(e))
+    raise e
+
+
+def raise_not_found_grpc_error(context, e: Exception):
+    context.set_code(grpc.StatusCode.NOT_FOUND)
+    context.set_details(str(e))
+    raise e

frogml_services_mock/mocks/vector_serving_api.py

@@ -0,0 +1,159 @@
+from dataclasses import dataclass
+from typing import Dict, List
+
+import grpc
+from numpy import dot
+from numpy.linalg import norm
+
+from _frogml_proto.qwak.vectors.v1.vector_pb2 import (
+    DoubleVector,
+    Property,
+    SearchResult,
+    StoredVector,
+    VectorIdentifier,
+)
+from _frogml_proto.qwak.vectors.v1.vector_service_pb2 import (
+    DeleteVectorsResponse,
+    FetchVectorRequest,
+    FetchVectorResponse,
+    SearchSimilarVectorsRequest,
+    SearchSimilarVectorsResponse,
+    UpsertVectorsRequest,
+    UpsertVectorsResponse,
+)
+from _frogml_proto.qwak.vectors.v1.vector_service_pb2_grpc import VectorServiceServicer
+
+
+@dataclass
+class VectorObject:
+    id: str
+    vector: List[float]
+    property: List[Property]
+
+
+class VectorServingServiceMock(VectorServiceServicer):
+    def __init__(self):
+        self._vector_collections: Dict[str, Dict] = dict()
+
+    def reset_vector_store(self) -> None:
+        self._vector_collections.clear()
+
+    def get_num_of_vectors(self, collection_name: str) -> int:
+        if collection_name not in self._vector_collections:
+            return 0
+
+        return len(self._vector_collections[collection_name].values())
+
+    def create_collection(self, collection_name: str) -> None:
+        self._vector_collections[collection_name] = dict()
+
+    def get_vector_by_ids(
+        self, collection_name: str, vector_ids: List[str]
+    ) -> List[dict]:
+        if collection_name not in self._vector_collections:
+            raise ValueError(f"Collection named {collection_name} doesn't exist")
+
+        collection = self._vector_collections[collection_name]
+        return [collection.get(_id) for _id in vector_ids if _id in collection]
+
+    def SearchSimilarVectors(self, request: SearchSimilarVectorsRequest, context):
+        if request.collection_name not in self._vector_collections:
+            context.set_details(f"Collection named {request.collection} doesn't exist'")
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            return
+
+        reference_vector = list(request.reference_vector.element)
+        collection_store = self._vector_collections.get(request.collection_name, dict())
+
+        # a naive impl of a "vector similarity" search - compute pairwise cosine distance on
+        # the entire set and return top results
+        result_set = sorted(
+            collection_store.values(),
+            key=lambda b: _cos_sim(reference_vector, b["vector"]),
+        )[-request.max_results :]
+
+        return SearchSimilarVectorsResponse(
+            search_results=[
+                SearchResult(
+                    id=result["id"] if request.include_id else None,
+                    properties=[
+                        p
+                        for p in result_set[0]["properties"]
+                        if p.name in request.properties
+                    ],
+                    vector=(
+                        DoubleVector(element=result["vector"])
+                        if request.include_vector
+                        else None
+                    ),
+                    distance=(
+                        _cos_sim(reference_vector, result["vector"])
+                        if request.include_distance
+                        else None
+                    ),
+                )
+                for result in result_set
+            ]
+        )
+
+    def UpsertVectors(self, request: UpsertVectorsRequest, context):
+        collection_store = self._vector_collections.get(request.collection_name, dict())
+        for stored_vector in request.vector:
+            id = stored_vector.vector_identifier.vector_id
+            collection_store[id] = {
+                "id": id,
+                "vector": list(stored_vector.vector.element),
+                "properties": stored_vector.properties,
+            }
+
+        self._vector_collections[request.collection_name] = collection_store
+        return UpsertVectorsResponse()
+
+    def DeleteVectors(self, request, context):
+        if request.collection_name not in self._vector_collections:
+            context.set_details(
+                f"Collection named {request.collection_name} doesn't exist'"
+            )
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            return
+
+        collection_store = self._vector_collections[request.collection_name]
+        ids_in_collection = [
+            vector_identifier.vector_id
+            for vector_identifier in request.vector_identifiers
+            if vector_identifier.vector_id in collection_store
+        ]
+        for id in ids_in_collection:
+            collection_store.pop(id)
+
+        self._vector_collections[request.collection_name] = collection_store
+        return DeleteVectorsResponse(num_vectors_deleted=len(ids_in_collection))
+
+    def FetchVector(self, request, context):
+        if request.collection_name not in self._vector_collections:
+            context.set_details(
+                f"Collection named {request.collection_name} doesn't exist'"
+            )
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            return
+
+        collection_store = self._vector_collections[request.collection_name]
+        vector_identifier = request.vector_identifier
+        vector_id = vector_identifier.vector_id
+        if vector_id not in collection_store:
+            context.set_details(f"Vector with id {vector_id} doesn't exist'")
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            return
+
+        stored_vector_dict = collection_store[vector_id]
+        return FetchVectorResponse(
+            vector=StoredVector(
+                vector_identifier=request.vector_identifier,
+                vector=DoubleVector(element=stored_vector_dict["vector"]),
+                properties=stored_vector_dict["properties"],
+            )
+        )
+
+
+def _cos_sim(a, b):
+    return dot(a, b) / (norm(a) * norm(b))

frogml_services_mock/mocks/vectors_management_api.py

@@ -0,0 +1,97 @@
+import uuid
+from datetime import datetime
+
+import grpc
+from google.protobuf.timestamp_pb2 import Timestamp
+
+from _frogml_proto.qwak.vectors.v1.collection.collection_pb2 import (
+    COLLECTION_STATUS_CREATE_REQUESTED,
+    QwakMetadata,
+    VectorCollection,
+    VectorCollectionDefinition,
+)
+from _frogml_proto.qwak.vectors.v1.collection.collection_service_pb2 import (
+    CreateCollectionResponse,
+    DeleteCollectionByIdResponse,
+    DeleteCollectionByNameResponse,
+    GetCollectionByIdResponse,
+    GetCollectionByNameResponse,
+    ListCollectionsResponse,
+)
+from _frogml_proto.qwak.vectors.v1.collection.collection_service_pb2_grpc import (
+    VectorCollectionServiceServicer,
+)
+
+
+class VectorCollectionManagementServiceMock(VectorCollectionServiceServicer):
+    def __init__(self):
+        self._collections_spec_by_ids = {}
+        self._collections_spec_by_name = {}
+
+    def reset_collections(self):
+        self._collections_spec_by_ids.clear()
+        self._collections_spec_by_name.clear()
+
+    def CreateCollection(self, request, context):
+        timestamp = Timestamp()
+        timestamp.FromDatetime(datetime.now())
+
+        collection_id = str(uuid.uuid4())
+        vector_collection = VectorCollection(
+            metadata=QwakMetadata(
+                created_at=timestamp,
+                created_by="it@frogml.com",
+                last_modified_at=timestamp,
+                last_modified_by="it@frogml.com",
+            ),
+            definition=VectorCollectionDefinition(
+                id=collection_id,
+                collection_spec=request.collection_spec,
+            ),
+            status=COLLECTION_STATUS_CREATE_REQUESTED,
+        )
+
+        self._collections_spec_by_ids[collection_id] = vector_collection
+        self._collections_spec_by_name[request.collection_spec.name] = vector_collection
+        return CreateCollectionResponse(vector_collection=vector_collection)
+
+    def GetCollectionById(self, request, context):
+        if request.id in self._collections_spec_by_ids:
+            return GetCollectionByIdResponse(
+                vector_collection=self._collections_spec_by_ids[request.id]
+            )
+
+        context.set_details(f"Collection ID {request.id} doesn't exist'")
+        context.set_code(grpc.StatusCode.NOT_FOUND)
+
+    def GetCollectionByName(self, request, context):
+        if request.name in self._collections_spec_by_name:
+            return GetCollectionByNameResponse(
+                vector_collection=self._collections_spec_by_name[request.name]
+            )
+
+        context.set_details(f"Collection name {request.name} doesn't exist'")
+        context.set_code(grpc.StatusCode.NOT_FOUND)
+
+    def DeleteCollectionById(self, request, context):
+        if request.id in self._collections_spec_by_ids:
+            self._collections_spec_by_ids.pop(request.id)
+            return DeleteCollectionByIdResponse()
+
+        context.set_details(f"Collection ID {request.id} doesn't exist'")
+        context.set_code(grpc.StatusCode.NOT_FOUND)
+
+    def DeleteCollectionByName(self, request, context):
+        if request.name in self._collections_spec_by_name:
+            self._collections_spec_by_name.pop(request.name)
+            return DeleteCollectionByNameResponse()
+
+        context.set_details(f"Collection name {request.name} doesn't exist'")
+        context.set_code(grpc.StatusCode.NOT_FOUND)
+
+    def ListCollections(self, request, context):
+        return ListCollectionsResponse(
+            vector_collections=[
+                collection for collection in self._collections_spec_by_ids.values()
+            ]
+        )

frogml_services_mock/mocks/workspace_manager_service_mock.py

@@ -0,0 +1,202 @@
+import logging
+import uuid
+from dataclasses import dataclass
+
+import grpc
+
+from _frogml_proto.qwak.user_application.common.v0.resources_pb2 import (
+    ClientPodComputeResources,
+    PodComputeResourceTemplateSpec,
+)
+from _frogml_proto.qwak.workspace.workspace_pb2 import (
+    DefaultWorkspaceDetails,
+    Workspace,
+    WorkspaceImage,
+)
+from _frogml_proto.qwak.workspace.workspace_service_pb2 import (
+    CountWorkspacesRequest,
+    CountWorkspacesResponse,
+    CreateWorkspaceResponse,
+    DeleteWorkspaceResponse,
+    GetDefaultWorkspaceValuesRequest,
+    GetDefaultWorkspaceValuesResponse,
+    GetWorkspaceByIdResponse,
+    ListWorkspaceImagesResponse,
+    ListWorkspacesResponse,
+    UpdateWorkspaceResponse,
+)
+from _frogml_proto.qwak.workspace.workspace_service_pb2_grpc import (
+    WorkspaceManagementServiceServicer,
+)
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_CPU_IMAGE = "1"
+DEFAULT_GPU_IMAGE = "3"
+DEFAULT_CPU_INSTANCE = "small"
+DEFAULT_GPU_INSTANCE = "g5_xlarge"
+
+
+@dataclass
+class WorkspaceData:
+    id: str
+    workspace_name: str
+    image_id: str
+    template_id: str
+    status: str
+
+
+class WorkspaceManagerServiceMock(WorkspaceManagementServiceServicer):
+    def __init__(self):
+        super(WorkspaceManagerServiceMock, self).__init__()
+        self.workspaces = dict()
+        self.image_ids_to_names = {
+            "1": "cpu_3.7",
+            "2": "cpu_3.8",
+            "3": "gpu_3.7",
+            "4": "gpu_3.8",
+        }
+
+    def CreateWorkspace(self, request, context):
+        logger.info(f"Creating workspace {request}")
+        workspace_id = self._get_workspace_id()
+        self.workspaces[workspace_id] = WorkspaceData(
+            id=workspace_id,
+            workspace_name=request.workspace_spec.workspace_name,
+            image_id=request.workspace_spec.image_id,
+            template_id=request.workspace_spec.client_pod_compute_resources.template_spec.template_id,
+            status="CREATED",
+        )
+        return CreateWorkspaceResponse(workspace_id=workspace_id)
+
+    def UpdateWorkspace(self, request, context):
+        logger.info(f"Updating workspace {request}")
+        if request.workspace_id not in self.workspaces:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            context.set_details("Received a non existing workspace id")
+            return UpdateWorkspaceResponse()
+        self.workspaces[request.workspace_id] = WorkspaceData(
+            id=request.workspace_id,
+            workspace_name=request.workspace_spec.workspace_name,
+            image_id=request.workspace_spec.image_id,
+            template_id=request.workspace_spec.client_pod_compute_resources.template_spec.template_id,
+            status="CREATED",
+        )
+        return UpdateWorkspaceResponse()
+
+    def DeleteWorkspace(self, request, context):
+        logger.info(f"delete workspace request: {request}")
+        if request.workspace_id not in self.workspaces:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            context.set_details("Received a non existing workspace id")
+            return UpdateWorkspaceResponse()
+        del self.workspaces[request.workspace_id]
+        return DeleteWorkspaceResponse()
+
+    def GetWorkspaceById(self, request, context):
+        logger.info(f"get workspace by id request: {request}")
+        if request.workspace_id not in self.workspaces:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            context.set_details("Received a non existing workspace id")
+            return UpdateWorkspaceResponse()
+
+        return GetWorkspaceByIdResponse(
+            workspace=Workspace(
+                workspace_id=request.workspace_id,
+                workspace_name=self.workspaces[request.workspace_id].workspace_name,
+                image_id=self.workspaces[request.workspace_id].image_id,
+                client_pod_compute_resources=ClientPodComputeResources(
+                    template_spec=PodComputeResourceTemplateSpec(
+                        template_id=self.workspaces[request.workspace_id].template_id
+                    )
+                ),
+            )
+        )
+
+    def ListWorkspaces(self, request, context):
+        logger.info(f"list workspaces: {request}")
+        return ListWorkspacesResponse(
+            workspaces=[
+                Workspace(
+                    workspace_id=workspace_id,
+                    workspace_name=workspace.workspace_name,
+                    image_id=workspace.image_id,
+                    client_pod_compute_resources=ClientPodComputeResources(
+                        template_spec=PodComputeResourceTemplateSpec(
+                            template_id=workspace.template_id
+                        )
+                    ),
+                )
+                for workspace_id, workspace in self.workspaces.items()
+            ]
+        )
+
+    def DeployWorkspace(self, request, context):
+        logger.info(f"Deploying workspace {request}")
+        if request.workspace_id not in self.workspaces:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            context.set_details("Received a non existing workspace id")
+            return UpdateWorkspaceResponse()
+
+        current_workspace = self.workspaces[request.workspace_id]
+        self.workspaces[request.workspace_id] = WorkspaceData(
+            id=request.workspace_id,
+            workspace_name=current_workspace.workspace_name,
+            image_id=current_workspace.image_id,
+            template_id=current_workspace.template_id,
+            status="DEPLOYING",
+        )
+        return UpdateWorkspaceResponse()
+
+    def UndeployWorkspace(self, request, context):
+        logger.info(f"Undeploying workspace {request}")
+        if request.workspace_id not in self.workspaces:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            context.set_details("Received a non existing workspace id")
+            return UpdateWorkspaceResponse()
+
+        current_workspace = self.workspaces[request.workspace_id]
+        self.workspaces[request.workspace_id] = WorkspaceData(
+            id=request.workspace_id,
+            workspace_name=current_workspace.workspace_name,
+            image_id=current_workspace.image_id,
+            template_id=current_workspace.template_id,
+            status="UNDEPLOYING",
+        )
+        return UpdateWorkspaceResponse()
+
+    def ListWorkspaceImages(self, request, context):
+        logger.info("List workspace images")
+        workspace_images = []
+        for k, v in self.image_ids_to_names.items():
+            workspace_images.append(WorkspaceImage(name=v, id=k))
+        return ListWorkspaceImagesResponse(workspace_images=workspace_images)
+
+    def CountWorkspaces(
+        self, request: CountWorkspacesRequest, context
+    ) -> CountWorkspacesResponse:
+        return CountWorkspacesResponse(count=len(self.workspaces))
+
+    def GetDefaultWorkspaceValues(
+        self, request: GetDefaultWorkspaceValuesRequest, context
+    ) -> GetDefaultWorkspaceValuesResponse:
+        return GetDefaultWorkspaceValuesResponse(
+            default_workspace_details=DefaultWorkspaceDetails(
+                cpu_image_id=DEFAULT_CPU_IMAGE,
+                gpu_image_id=DEFAULT_GPU_IMAGE,
+                cpu_compute_resources=ClientPodComputeResources(
+                    template_spec=PodComputeResourceTemplateSpec(
+                        template_id=DEFAULT_CPU_INSTANCE
+                    )
+                ),
+                gpu_compute_resources=ClientPodComputeResources(
+                    template_spec=PodComputeResourceTemplateSpec(
+                        template_id=DEFAULT_GPU_INSTANCE
+                    )
+                ),
+            )
+        )
+
+    @staticmethod
+    def _get_workspace_id():
+        return str(uuid.uuid4())