indexify 0.2.6__py3-none-any.whl → 0.2.7__py3-none-any.whl

indexify/cli.py

@@ -245,7 +245,7 @@ WORKDIR /app
     console.print(f"{docker_file}", style="magenta")
 
     client = docker.from_env()
-    image_name =f"{image._image_name}:{image._tag}"
+    image_name = f"{image._image_name}:{image._tag}"
     client.images.build(
         fileobj=io.BytesIO(docker_file.encode()),
         tag=image_name,

indexify/executor/function_worker.py

@@ -21,6 +21,8 @@ import concurrent.futures
 import io
 from contextlib import redirect_stderr, redirect_stdout
 
+from .runtime_probes import RuntimeProbes
+
 
 class FunctionRunException(Exception):
     def __init__(
@@ -51,7 +53,8 @@ def _load_function(
     key = f"{namespace}/{graph_name}/{version}/{fn_name}"
     if key in function_wrapper_map:
         return
-    graph = Graph.from_path(code_path)
+    image_name = RuntimeProbes().probe().image_name
+    graph = Graph.from_path(code_path, image_name)
     function_wrapper = graph.get_function(fn_name)
     function_wrapper_map[key] = function_wrapper
     graph_key = f"{namespace}/{graph_name}/{version}"
@@ -135,7 +138,9 @@ def _run_function(
     fn_output = None
     has_failed = False
     exception_msg = None
-    print(f"[bold] function_worker: [/bold] invoking function {fn_name} in graph {graph_name}")
+    print(
+        f"[bold] function_worker: [/bold] invoking function {fn_name} in graph {graph_name}"
+    )
     with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
         try:
             key = f"{namespace}/{graph_name}/{version}/{fn_name}"

indexify/functions_sdk/graph.py

@@ -132,33 +132,38 @@ class Graph:
             self.add_node(node)
             self.routers[from_node.name].append(node.name)
         return self
-    
-    
-    def _register_cloudpickle(self):
+
+    def serialize(self) -> Dict[str, bytes]:
         # Get all unique modules from nodes and edges
-        modules = set()
+        modules_by_images = defaultdict(set)
         for node in self.nodes.values():
-            modules.add(node.__module__)
+            modules_by_images[node.image._image_name].add(node.__module__)
+
+        pickled_nodes_by_image = {}
+        for image_name, modules in modules_by_images.items():
+            for module in modules:
+                print(f"registering module {module} with cloudpickle")
+                cloudpickle.register_pickle_by_value(sys.modules[module])
+            pickled_nodes_by_image[image_name] = cloudpickle.dumps(self)
+            for module in modules:
+                cloudpickle.unregister_pickle_by_value(sys.modules[module])
 
         # Register each module with cloudpickle
         for module_name in modules:
             module = sys.modules[module_name]
             print(f"registering module {module_name} with cloudpickle")
             cloudpickle.register_pickle_by_value(module)
-
-
-    def serialize(self):
-        self._register_cloudpickle()
-        return cloudpickle.dumps(self)
+        return pickled_nodes_by_image
 
     @staticmethod
-    def deserialize(graph: bytes) -> "Graph":
-        return cloudpickle.loads(graph)
+    def deserialize(serialized_code_by_images: Dict[str, bytes], image: str) -> "Graph":
+        return cloudpickle.loads(serialized_code_by_images[image])
 
     @staticmethod
-    def from_path(path: str) -> "Graph":
+    def from_path(path: str, image: str) -> "Graph":
         with open(path, "rb") as f:
-            return cloudpickle.load(f)
+            pickled_code_by_images: Dict[str, bytes] = cloudpickle.load(f)
+        return Graph.deserialize(pickled_code_by_images, image)
 
     def add_edge(
         self,

indexify/functions_sdk/graph_definition.py

@@ -27,6 +27,11 @@ class NodeMetadata(BaseModel):
     dynamic_router: Optional[RouterMetadata] = None
     compute_fn: Optional[FunctionMetadata] = None
 
+    def image_name(self):
+        if self.dynamic_router:
+            return self.dynamic_router.image_name
+        return self.compute_fn.image_name
+
 
 # RuntimeInformation is a class that holds data about the environment in which the graph should run.
 class RuntimeInformation(BaseModel):

indexify/http_client.py

@@ -99,7 +99,7 @@ class IndexifyClient:
         service_url: str = DEFAULT_SERVICE_URL_HTTPS,
         *args,
         **kwargs,
-    ) -> "RemoteClient":
+    ) -> "IndexifyClient":
         """
         Create a client with mutual TLS authentication. Also enables HTTP/2,
         which is required for mTLS.
@@ -127,7 +127,7 @@ class IndexifyClient:
 
         client_certs = (cert_path, key_path)
         verify_option = ca_bundle_path if ca_bundle_path else True
-        client = RemoteClient(
+        client = IndexifyClient(
             *args,
             **kwargs,
             service_url=service_url,
@@ -160,7 +160,7 @@ class IndexifyClient:
 
     def register_compute_graph(self, graph: Graph):
         graph_metadata = graph.definition()
-        serialized_code = graph.serialize()
+        serialized_code = cloudpickle.dumps(graph.serialize())
         response = self._post(
             f"namespaces/{self.namespace}/compute_graphs",
             files={"code": serialized_code},
@@ -177,11 +177,14 @@ class IndexifyClient:
         response = self._get(f"namespaces/{self.namespace}/compute_graphs/{name}")
         return ComputeGraphMetadata(**response.json())
 
-    def load_graph(self, name: str) -> Graph:
+    def load_graph(self, name: str, fn_name: str) -> Graph:
         response = self._get(
             f"internal/namespaces/{self.namespace}/compute_graphs/{name}/code"
         )
-        return Graph.deserialize(response.content)
+        cg_metadata: ComputeGraphMetadata = self.graph(name)
+        fn_metadata = cg_metadata.nodes[fn_name]
+        pickled_bytes_by_image = cloudpickle.loads(response.content)
+        return Graph.deserialize(pickled_bytes_by_image, fn_metadata.image_name())
 
     def namespaces(self) -> List[str]:
         response = self._get(f"namespaces")
@@ -287,7 +290,7 @@ class IndexifyClient:
         return: Union[Dict[str, List[Any]], List[Any]]: The extracted objects. If the extractor name is provided, the output is a list of extracted objects by the extractor. If the extractor name is not provided, the output is a dictionary with the extractor name as the key and the extracted objects as the value. If no objects are found, an empty list is returned.
         """
         if graph not in self._graphs:
-            self._graphs[graph] = self.load_graph(graph)
+            self._graphs[graph] = self.load_graph(graph, fn_name)
         response = self._get(
             f"namespaces/{self.namespace}/compute_graphs/{graph}/invocations/{invocation_id}/outputs",
         )

indexify/remote_graph.py

@@ -34,13 +34,19 @@ class RemoteGraph:
         )
 
     @classmethod
-    def deploy(cls, g: Graph, additional_modules=[], server_url: Optional[str] = "http://localhost:8900"):
+    def deploy(
+        cls,
+        g: Graph,
+        additional_modules=[],
+        server_url: Optional[str] = "http://localhost:8900",
+    ):
         """
         Create a new RemoteGraph from a local Graph object.
         :param g: The local Graph object.
         :param server_url: The URL of the server where the graph will be registered.
         """
         import cloudpickle
+
         for module in additional_modules:
             cloudpickle.register_pickle_by_value(module)
         client = IndexifyClient(service_url=server_url)

{indexify-0.2.6.dist-info → indexify-0.2.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: indexify
-Version: 0.2.6
+Version: 0.2.7
 Summary: Python Client for Indexify
 Home-page: https://github.com/tensorlakeai/indexify
 License: Apache 2.0

{indexify-0.2.6.dist-info → indexify-0.2.7.dist-info}/RECORD

@@ -1,5 +1,5 @@
 indexify/__init__.py,sha256=qgMBKVrM_tI-tFeWpE8ktlC5rcExk05nbWyFqxxqeEU,496
-indexify/cli.py,sha256=IK83AN8IZ0Dxgwv_azCjC7gqLYeKaGVu7b4EYs4vw5Q,7580
+indexify/cli.py,sha256=Hj-FvDzrRhgVuX1X3ReSp6QW4lurRUGHKhU9G7y7tDU,7581
 indexify/data_loaders/__init__.py,sha256=Y5NEuseTcYAICRiweYw5wBQ2m2YplbsY21I7df-rdi4,1339
 indexify/data_loaders/local_directory_loader.py,sha256=fCrgj5drnW71ZUdDDvcB1-VJjIs1w6Q8sEW0HSGSAiA,1247
 indexify/data_loaders/url_loader.py,sha256=32SERljcq1Xsi4RdLz2dgyk2TER5pQPTtXl3gUzwHbY,1533
@@ -8,24 +8,24 @@ indexify/executor/agent.py,sha256=I08CiOWeJ_mz8OHr9_iJfp07Ma1VMQirZ2MsDp8lDZw,14
 indexify/executor/api_objects.py,sha256=SysjlGYu4JtYdqfexZHHN1IW4TtaDdFUF3hYZ5mpUJU,810
 indexify/executor/downloader.py,sha256=0MPiKw0AWs3Z7ReC9l2z-3515yqq85ghPzdh485dnuw,3998
 indexify/executor/executor_tasks.py,sha256=gAZ2pvza1YwGlaR1o_tJW4SXtdCgK7sLJgp4W7rOjR0,1834
-indexify/executor/function_worker.py,sha256=NSxcB0RtrTCJnDdtYu-CIEbyZ1QgHC-nxHdH3UUQ1tg,5506
+indexify/executor/function_worker.py,sha256=qUsaZNDRcI0PxOBosL1A7x7YPdaJp_gyq3p7eTottJM,5627
 indexify/executor/indexify_executor.py,sha256=2Ut_VX-Su_lm4b4aEROyRJ3gXx-uFHA-V7EN0sWiARE,771
 indexify/executor/runtime_probes.py,sha256=tvi8KCaQTVJqcyBJ4-jzEUAnQ01ZbMmjCxV2KJ96_PI,1449
 indexify/executor/task_reporter.py,sha256=gnnse0v6rjjni8lNzeb-ZYq6iF2DgafKoT7dcGUZhQ4,3716
 indexify/executor/task_store.py,sha256=q8s2gImsFffWeXQR0mk1Xlo1Aj_2GfclNPjQ2EA_YBo,3984
 indexify/functions_sdk/data_objects.py,sha256=2LqAWJ_S2Xkp4OQTmhd3InVIrBs7juV41udnSQFMMfM,840
-indexify/functions_sdk/graph.py,sha256=ag_Hi0HuJwFn15RJG6anJaCzux2Gc8grTbEPQHReVYk,15311
-indexify/functions_sdk/graph_definition.py,sha256=EJfC0MdKEbFF1CBaU0htrveSlcAQJCH96DLSNfZ02V4,1178
+indexify/functions_sdk/graph.py,sha256=OtXZhYuPoP3KMkQLRHgr4Dx_Kj9YNAjGQxx7a_fMRSo,15945
+indexify/functions_sdk/graph_definition.py,sha256=ziiXrWXcu5x1zEfLErZyiedK75ufiSWxJBhaeV18ikQ,1329
 indexify/functions_sdk/graph_validation.py,sha256=XLHiC9PAtZungJLysU3hIUOPNDkO5TXUDZ_jiZ0H4hg,2508
 indexify/functions_sdk/image.py,sha256=l9DghfPTV1p0nlc1Fm1ZuJe21InJeBHUTULswmLmmx8,727
 indexify/functions_sdk/indexify_functions.py,sha256=xxgvnw0MQ_csIksunIdero8be0PR4mfwgoHp3UlkMZU,5851
 indexify/functions_sdk/local_cache.py,sha256=cNWF67zbhbTJe3g86hyLBy3Rqzs6dNvp2SjLazGZWvw,1348
 indexify/functions_sdk/object_serializer.py,sha256=Zz4GobW3ZamBBtFDF76QxU3TP6oJNdWnhsfKd0OUFoc,1660
-indexify/http_client.py,sha256=QwbLQAWVUZ95SZ-4XcnvtQRXyHtH2APqUj7nKh0C3DQ,12161
-indexify/remote_graph.py,sha256=zmcjuv1r6vTNueNtVnknR1aDdbDw5LPdM8y-EczXqmk,2948
+indexify/http_client.py,sha256=pe_xE0oTodlEet6XTdl434wvrWN8CWHKyv1pOxOF6Qw,12418
+indexify/remote_graph.py,sha256=gyrbgrAFxb6XSmAlPYa59TuSDa2VPscuWT6qQ4VnIWA,2988
 indexify/settings.py,sha256=LSaWZ0ADIVmUv6o6dHWRC3-Ry5uLbCw2sBSg1e_U7UM,99
-indexify-0.2.6.dist-info/LICENSE.txt,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-indexify-0.2.6.dist-info/METADATA,sha256=9tBE6KGjLKIGQ5O7KyuBc9dFEd-JXTj4cBwjMxGGRb8,6129
-indexify-0.2.6.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-indexify-0.2.6.dist-info/entry_points.txt,sha256=Pih7WV-XMpAzI5dEvROcpLr-ybVhd9Y-AtuzBKUdcDs,49
-indexify-0.2.6.dist-info/RECORD,,
+indexify-0.2.7.dist-info/LICENSE.txt,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+indexify-0.2.7.dist-info/METADATA,sha256=TxMLI-S5x3018FNbhGZBUlPRQ3V2fLEroxrbzuxfkxg,6129
+indexify-0.2.7.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+indexify-0.2.7.dist-info/entry_points.txt,sha256=Pih7WV-XMpAzI5dEvROcpLr-ybVhd9Y-AtuzBKUdcDs,49
+indexify-0.2.7.dist-info/RECORD,,