indexify 0.2.9__tar.gz → 0.2.11__tar.gz

{indexify-0.2.9 → indexify-0.2.11}/PKG-INFO

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

{indexify-0.2.9 → indexify-0.2.11}/indexify/cli.py

@@ -119,7 +119,9 @@ def server_dev_mode():
 
 @app.command(help="Build image for function names")
 def build_image(
-    workflow_file_path: str, func_names: List[str], python_sdk_path: Optional[str] = None
+    workflow_file_path: str,
+    func_names: List[str],
+    python_sdk_path: Optional[str] = None,
 ):
     globals_dict = {}
 

{indexify-0.2.9 → indexify-0.2.11}/indexify/executor/function_worker.py

@@ -18,6 +18,7 @@ function_wrapper_map: Dict[str, IndexifyFunctionWrapper] = {}
 
 import concurrent.futures
 
+
 class FunctionRunException(Exception):
     def __init__(
         self, exception: Exception, stdout: str, stderr: str, is_reducer: bool
@@ -131,7 +132,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}"
@@ -139,7 +142,10 @@ def _run_function(
                 _load_function(namespace, graph_name, fn_name, code_path, version)
 
             fn = function_wrapper_map[key]
-            if str(type(fn.indexify_function)) == "<class 'indexify.functions_sdk.indexify_functions.IndexifyRo'>":
+            if (
+                str(type(fn.indexify_function))
+                == "<class 'indexify.functions_sdk.indexify_functions.IndexifyRo'>"
+            ):
                 router_output = fn.invoke_router(fn_name, input)
             else:
                 fn_output = fn.invoke_fn_ser(fn_name, input, init_value)
@@ -147,6 +153,7 @@ def _run_function(
                 is_reducer = fn.indexify_function.accumulate is not None
         except Exception as e:
             import sys
+
             print(traceback.format_exc(), file=sys.stderr)
             has_failed = True
             exception_msg = str(e)

{indexify-0.2.9 → indexify-0.2.11}/indexify/functions_sdk/graph.py

@@ -117,13 +117,16 @@ class Graph:
             self.routers[from_node.name].append(node.name)
         return self
 
-    def serialize(self):
+    def serialize(self, additional_modules):
         # Get all unique modules from nodes and edges
         pickled_functions = {}
+        for module in additional_modules:
+            cloudpickle.register_pickle_by_value(module)
         for node in self.nodes.values():
             cloudpickle.register_pickle_by_value(sys.modules[node.__module__])
             pickled_functions[node.name] = cloudpickle.dumps(node)
-            cloudpickle.unregister_pickle_by_value(sys.modules[node.__module__])
+            if not sys.modules[node.__module__] in additional_modules:
+                cloudpickle.unregister_pickle_by_value(sys.modules[node.__module__])
         return pickled_functions
 
     def add_edge(

{indexify-0.2.9 → indexify-0.2.11}/indexify/http_client.py

@@ -160,9 +160,9 @@ class IndexifyClient:
     def __exit__(self, exc_type, exc_value, traceback):
         self.close()
 
-    def register_compute_graph(self, graph: Graph):
+    def register_compute_graph(self, graph: Graph, additional_modules):
         graph_metadata = graph.definition()
-        serialized_code = cloudpickle.dumps(graph.serialize())
+        serialized_code = cloudpickle.dumps(graph.serialize(additional_modules))
         response = self._post(
             f"namespaces/{self.namespace}/compute_graphs",
             files={"code": serialized_code},
@@ -197,9 +197,11 @@ class IndexifyClient:
         for item in namespaces_dict:
             namespaces.append(item["name"])
         return namespaces
-    
+
     @classmethod
-    def new_namespace(cls, namespace: str, server_addr: Optional[str] = "http://localhost:8900"):
+    def new_namespace(
+        cls, namespace: str, server_addr: Optional[str] = "http://localhost:8900"
+    ):
         # Create a new client instance with the specified server address
         client = cls(service_url=server_addr)
 
@@ -212,11 +214,10 @@ class IndexifyClient:
 
         # Set the namespace for the newly created client
         client.namespace = namespace
-        
+
         # Return the client instance with the new namespace
         return client
 
-
     def create_namespace(self, namespace: str):
         self._post("namespaces", json={"name": namespace})
 
@@ -253,8 +254,16 @@ class IndexifyClient:
                 for sse in event_source.iter_sse():
                     obj = json.loads(sse.data)
                     for k, v in obj.items():
+                        if k == "id":
+                            return v
                         if k == "InvocationFinished":
                             return v["id"]
+                        if k == "DiagnosticMessage":
+                            message = v.get("message", None)
+                            print(
+                                f"[bold red]scheduler diagnostic: [/bold red]{message}"
+                            )
+                            continue
                         event_payload = InvocationEventPayload.model_validate(v)
                         event = InvocationEvent(event_name=k, payload=event_payload)
                         if (

{indexify-0.2.9 → indexify-0.2.11}/indexify/remote_graph.py

@@ -32,7 +32,7 @@ class RemoteGraph:
         return self._client.invoke_graph_with_object(
             self._name, block_until_done, **kwargs
         )
-    
+
     def rerun(self):
         """
         Rerun the graph with the given invocation ID.
@@ -41,17 +41,19 @@ class RemoteGraph:
         self._client.rerun_graph(self._name)
 
     @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)
-        client.register_compute_graph(g)
+        client.register_compute_graph(g, additional_modules)
         return cls(name=g.name, server_url=server_url)
 
     @classmethod

{indexify-0.2.9 → indexify-0.2.11}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "indexify"
-version = "0.2.9"
+version = "0.2.11"
 description = "Python Client for Indexify"
 authors = ["Tensorlake Inc. <support@tensorlake.ai>"]
 license = "Apache 2.0"