frogml-core 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

frogml_core/vector_store/utils/upsert_utils.py

@@ -0,0 +1,218 @@
+import os
+from functools import partial
+from math import ceil
+from multiprocessing import Pool, set_start_method
+from typing import Dict, Iterable, List, Optional, Tuple, TypeVar, Union
+
+from dependency_injector.wiring import Provide, inject
+from typeguard import typechecked
+
+from frogml_proto.qwak.vectors.v1.vector_pb2 import (
+    DoubleVector,
+    Property,
+    StoredVector,
+    VectorIdentifier,
+)
+from frogml_core.clients.vector_store.serving_client import VectorServingClient
+from frogml_core.exceptions import FrogmlException
+from frogml_core.inner.di_configuration import FrogmlContainer
+from frogml_core.vector_store.inference_client import VectorStoreInferenceClient
+
+_NaturalInput = TypeVar("T")
+_NaturalInputs = List[_NaturalInput]
+_Vector = List[float]
+_Properties = Dict[str, Union[str, int, bool, float]]
+
+_type_to_proto_property_mapping: Dict = {
+    str: "string_val",
+    bool: "bool_val",
+    int: "int_val",
+    float: "double_val",
+}
+
+
+def _build_property(key: str, value: Union[str, int, bool, float]):
+    type_val = _type_to_proto_property_mapping.get(type(value), None)
+    if not type_val:
+        raise FrogmlException(
+            f"Cannot upsert vector with property value type {type(value)}. "
+            f"Supported types are: {list(_type_to_proto_property_mapping.keys())}"
+        )
+
+    property_args = {"name": key, type_val: value}
+    return Property(**property_args)
+
+
+def _rewire_qwak_container(config):
+    # re-creating the container using the config from the original container.
+    # note that this runs in a fresh interpreter - at that point there's a running
+    # container because of the imports, but it won't necessarily have the same config
+    # as the container in the parent process that spawned this one.
+    # rewiring only the vector store + ecosystem and authentication - if using stuff from
+    # other modules pls feel free to add it here.
+    new_container = FrogmlContainer(config=config)
+    from frogml_core.clients import vector_store
+    from frogml_core.clients.administration import authentication, eco_system
+
+    new_container.wire(
+        packages=[
+            authentication,
+            eco_system,
+            vector_store,
+        ]
+    )
+
+
+@typechecked
+def _divide_chunks(lst: List, chunk_size: int):
+    if chunk_size <= 0:
+        raise FrogmlException("Chunk size must be a positive integer")
+
+    num_items: int = len(lst)
+    num_chunks: int = ceil(num_items / chunk_size)
+    for i in range(num_chunks):
+        yield lst[i * chunk_size : (i + 1) * chunk_size]
+
+
+def _get_vector_identifier(t: Union[str, Tuple[str, str]]):
+    if type(t) is str:
+        return VectorIdentifier(vector_id=t)
+    return VectorIdentifier(vector_id=t[0], tenant_id=t[1])
+
+
+def _upsert_vector_block(
+    vector_tuples: List[Tuple[Union[str, Tuple[str, str]], _Vector, _Properties]],
+    chunk_size: int,
+    collection_name: str,
+    edge_services_url: str,
+) -> None:
+    vector_serving_client: VectorServingClient = VectorServingClient(
+        edge_services_url=edge_services_url
+    )
+    for chunk in _divide_chunks(vector_tuples, chunk_size):
+        # chunk is a list of (id, vector, properties) tuples
+        vector_serving_client.upsert_vectors(
+            collection_name=collection_name,
+            vectors=[
+                StoredVector(
+                    vector_identifier=_get_vector_identifier(tpl[0]),
+                    vector=DoubleVector(element=tpl[1]),
+                    properties=[
+                        _build_property(key, value) for (key, value) in tpl[2].items()
+                    ],
+                )
+                for tpl in chunk
+            ],
+        )
+
+
+def _upsert_natural_input_block(
+    vector_tuples: List[Tuple[VectorIdentifier, _NaturalInput, _Properties]],
+    chunk_size: int,
+    vectorizer_name: str,
+    collection_name: str,
+    edge_services_url: str,
+) -> None:
+    vector_serving_client: VectorServingClient = VectorServingClient(
+        edge_services_url=edge_services_url
+    )
+    inference_client: VectorStoreInferenceClient = VectorStoreInferenceClient(
+        model_id=vectorizer_name
+    )
+    for chunk in _divide_chunks(vector_tuples, chunk_size):
+        # chunk is a list of (id, _NaturalInput, properties) tuples
+        vector_serving_client.upsert_vectors(
+            collection_name=collection_name,
+            vectors=[
+                StoredVector(
+                    vector_identifier=_get_vector_identifier(tpl[0]),
+                    vector=DoubleVector(
+                        element=inference_client.get_embedding(natural_input=tpl[1])
+                    ),
+                    properties=[
+                        _build_property(key, value) for (key, value) in tpl[2].items()
+                    ],
+                )
+                for tpl in chunk
+            ],
+        )
+
+
+@inject
+def _upsert_natural_input(
+    vector_tuples: List[Tuple[Union[str, Tuple[str, str]], _NaturalInput, _Properties]],
+    chunk_size: int,
+    vectorizer_name: str,
+    collection_name: str,
+    edge_services_url: str,
+    multiproc: bool = False,
+    max_processes: Optional[int] = None,
+    config=Provide[FrogmlContainer.config],
+):
+    if not multiproc:
+        _upsert_natural_input_block(
+            vector_tuples=vector_tuples,
+            chunk_size=chunk_size,
+            vectorizer_name=vectorizer_name,
+            collection_name=collection_name,
+            edge_services_url=edge_services_url,
+        )
+    else:
+        if max_processes is None:
+            max_processes = os.cpu_count()
+        effective_block_size: int = ceil(len(vector_tuples) / (max_processes * 4))
+
+        f = partial(
+            _upsert_natural_input_block,
+            chunk_size=chunk_size,
+            vectorizer_name=vectorizer_name,
+            collection_name=collection_name,
+            edge_services_url=edge_services_url,
+        )
+
+        blocks: Iterable[List[Tuple[str, _NaturalInput, _Properties]]] = _divide_chunks(
+            vector_tuples, effective_block_size
+        )
+        initializer = partial(_rewire_qwak_container, config=config)
+        set_start_method("spawn", force=True)
+
+        with Pool(processes=max_processes, initializer=initializer) as p:
+            p.map(f, blocks)
+
+
+@inject
+def _upsert_vectors(
+    vector_tuples: List[Tuple[Union[str, Tuple[str, str]], _Vector, _Properties]],
+    chunk_size: int,
+    collection_name: str,
+    edge_services_url: str,
+    multiproc: bool = False,
+    max_processes: Optional[int] = None,
+    config=Provide[FrogmlContainer.config],
+):
+    if not multiproc:
+        _upsert_vector_block(
+            vector_tuples=vector_tuples,
+            chunk_size=chunk_size,
+            collection_name=collection_name,
+            edge_services_url=edge_services_url,
+        )
+    else:
+        if max_processes is None:
+            max_processes = os.cpu_count()
+        effective_block_size: int = ceil(len(vector_tuples) / (max_processes * 4))
+
+        f = partial(
+            _upsert_vector_block,
+            chunk_size=chunk_size,
+            collection_name=collection_name,
+            edge_services_url=edge_services_url,
+        )
+        blocks: Iterable[List[Tuple[str, _Vector, _Properties]]] = list(
+            _divide_chunks(vector_tuples, effective_block_size)
+        )
+
+        set_start_method("spawn", force=True)
+        initializer = partial(_rewire_qwak_container, config=config)
+        with Pool(processes=max_processes, initializer=initializer) as p:
+            p.map(f, blocks)

frogml_core-0.0.2.dist-info/METADATA

@@ -0,0 +1,49 @@
+Metadata-Version: 2.3
+Name: frogml-core
+Version: 0.0.2
+Summary: frogml Core contains the necessary objects and communication tools for using the Jfrog ml Platform
+License: Apache-2.0
+Keywords: mlops,ml,deployment,serving,model
+Author: Jfrog
+Requires-Python: >=3.9,<3.12
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.7
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Provides-Extra: feature-store
+Requires-Dist: PyYAML
+Requires-Dist: cachetools (>=5.5.1,<6.0.0)
+Requires-Dist: chevron (>=0.14.0,<0.15.0)
+Requires-Dist: cloudpickle (==2.2.1) ; extra == "feature-store"
+Requires-Dist: dacite (==1.8.1)
+Requires-Dist: dependency-injector (>=4.0)
+Requires-Dist: filelock
+Requires-Dist: frogml-storage (>=0.9.0,<0.10.0)
+Requires-Dist: grpcio (>=1.57.0)
+Requires-Dist: joblib (>=1.3.2,<2.0.0)
+Requires-Dist: marshmallow-dataclass (>=8.5.8,<9.0.0)
+Requires-Dist: protobuf (>=3.10,<4) ; python_full_version >= "3.7.1" and python_version < "3.10"
+Requires-Dist: protobuf (>=4.21.6) ; python_version >= "3.10"
+Requires-Dist: pyarrow (>=6.0.0) ; extra == "feature-store"
+Requires-Dist: pyathena (>=2.2.0,!=2.18.0) ; extra == "feature-store"
+Requires-Dist: pyspark (==3.4.2) ; extra == "feature-store"
+Requires-Dist: python-jose
+Requires-Dist: python-json-logger (>=2.0.2)
+Requires-Dist: requests
+Requires-Dist: retrying (==1.3.4)
+Requires-Dist: typeguard (>=2,<3)
+Requires-Dist: typing-extensions (>=4.12.2,<5.0.0)
+Project-URL: Home page, https://www.jfrog.com/
+Description-Content-Type: text/markdown
+
+# Frogml Core
+
+Frogml is an end-to-end production ML platform designed to allow data scientists to build, deploy, and monitor their models in production with minimal engineering friction.
+Frogml Core contains all the objects and tools necessary to use the Frogml Platform
+