indexify 0.2.21__tar.gz → 0.2.23__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (34) hide show
  1. {indexify-0.2.21 → indexify-0.2.23}/PKG-INFO +1 -1
  2. {indexify-0.2.21 → indexify-0.2.23}/indexify/__init__.py +4 -0
  3. {indexify-0.2.21 → indexify-0.2.23}/indexify/cli.py +11 -4
  4. indexify-0.2.23/indexify/error.py +8 -0
  5. {indexify-0.2.21 → indexify-0.2.23}/indexify/executor/agent.py +112 -10
  6. {indexify-0.2.21 → indexify-0.2.23}/indexify/executor/api_objects.py +1 -0
  7. {indexify-0.2.21 → indexify-0.2.23}/indexify/executor/executor_tasks.py +1 -0
  8. {indexify-0.2.21 → indexify-0.2.23}/indexify/executor/function_worker.py +42 -5
  9. indexify-0.2.23/indexify/executor/image_dependency_installer.py +64 -0
  10. {indexify-0.2.21 → indexify-0.2.23}/indexify/executor/runtime_probes.py +11 -0
  11. {indexify-0.2.21 → indexify-0.2.23}/indexify/executor/task_reporter.py +1 -1
  12. {indexify-0.2.21 → indexify-0.2.23}/indexify/functions_sdk/graph.py +13 -5
  13. {indexify-0.2.21 → indexify-0.2.23}/indexify/functions_sdk/indexify_functions.py +45 -7
  14. {indexify-0.2.21 → indexify-0.2.23}/indexify/http_client.py +25 -1
  15. {indexify-0.2.21 → indexify-0.2.23}/indexify/remote_graph.py +7 -2
  16. {indexify-0.2.21 → indexify-0.2.23}/pyproject.toml +1 -1
  17. indexify-0.2.21/indexify/error.py +0 -3
  18. {indexify-0.2.21 → indexify-0.2.23}/LICENSE.txt +0 -0
  19. {indexify-0.2.21 → indexify-0.2.23}/README.md +0 -0
  20. {indexify-0.2.21 → indexify-0.2.23}/indexify/data_loaders/__init__.py +0 -0
  21. {indexify-0.2.21 → indexify-0.2.23}/indexify/data_loaders/local_directory_loader.py +0 -0
  22. {indexify-0.2.21 → indexify-0.2.23}/indexify/data_loaders/url_loader.py +0 -0
  23. {indexify-0.2.21 → indexify-0.2.23}/indexify/executor/downloader.py +0 -0
  24. {indexify-0.2.21 → indexify-0.2.23}/indexify/executor/indexify_executor.py +0 -0
  25. {indexify-0.2.21 → indexify-0.2.23}/indexify/executor/task_store.py +0 -0
  26. {indexify-0.2.21 → indexify-0.2.23}/indexify/functions_sdk/data_objects.py +0 -0
  27. {indexify-0.2.21 → indexify-0.2.23}/indexify/functions_sdk/graph_definition.py +0 -0
  28. {indexify-0.2.21 → indexify-0.2.23}/indexify/functions_sdk/graph_validation.py +0 -0
  29. {indexify-0.2.21 → indexify-0.2.23}/indexify/functions_sdk/image.py +0 -0
  30. {indexify-0.2.21 → indexify-0.2.23}/indexify/functions_sdk/local_cache.py +0 -0
  31. {indexify-0.2.21 → indexify-0.2.23}/indexify/functions_sdk/object_serializer.py +0 -0
  32. {indexify-0.2.21 → indexify-0.2.23}/indexify/functions_sdk/pipeline.py +0 -0
  33. {indexify-0.2.21 → indexify-0.2.23}/indexify/remote_pipeline.py +0 -0
  34. {indexify-0.2.21 → indexify-0.2.23}/indexify/settings.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: indexify
3
- Version: 0.2.21
3
+ Version: 0.2.23
4
4
  Summary: Python Client for Indexify
5
5
  Home-page: https://github.com/tensorlakeai/indexify
6
6
  License: Apache 2.0
@@ -2,6 +2,8 @@ from . import data_loaders
2
2
  from .functions_sdk.graph import Graph
3
3
  from .functions_sdk.image import Image
4
4
  from .functions_sdk.indexify_functions import (
5
+ IndexifyFunction,
6
+ get_ctx,
5
7
  indexify_function,
6
8
  indexify_router,
7
9
  )
@@ -19,6 +21,8 @@ __all__ = [
19
21
  "RemotePipeline",
20
22
  "Image",
21
23
  "indexify_function",
24
+ "get_ctx",
25
+ "IndexifyFunction",
22
26
  "indexify_router",
23
27
  "DEFAULT_SERVICE_URL",
24
28
  "IndexifyClient",
@@ -1,5 +1,4 @@
1
1
  import asyncio
2
- import io
3
2
  import os
4
3
  import shutil
5
4
  import signal
@@ -168,6 +167,12 @@ def executor(
168
167
  executor_cache: Optional[str] = typer.Option(
169
168
  "~/.indexify/executor_cache", help="Path to the executor cache directory"
170
169
  ),
170
+ name_alias: Optional[str] = typer.Option(
171
+ None, help="Name alias for the executor if it's spun up with the base image"
172
+ ),
173
+ image_version: Optional[int] = typer.Option(
174
+ "1", help="Requested Image Version for this executor"
175
+ ),
171
176
  ):
172
177
  id = nanoid.generate()
173
178
  console.print(
@@ -176,13 +181,14 @@ def executor(
176
181
  f"Config path: {config_path}\n"
177
182
  f"Server address: {server_addr}\n"
178
183
  f"Executor ID: {id}\n"
179
- f"Executor cache: {executor_cache}",
184
+ f"Executor cache: {executor_cache}\n"
185
+ f"Name Alias: {name_alias}"
186
+ f"Image Version: {image_version}\n",
180
187
  title="Agent Configuration",
181
188
  border_style="info",
182
189
  )
183
190
  )
184
191
 
185
- function_worker = FunctionWorker(workers=workers)
186
192
  from pathlib import Path
187
193
 
188
194
  executor_cache = Path(executor_cache).expanduser().absolute()
@@ -193,10 +199,11 @@ def executor(
193
199
  agent = ExtractorAgent(
194
200
  id,
195
201
  num_workers=workers,
196
- function_worker=function_worker,
197
202
  server_addr=server_addr,
198
203
  config_path=config_path,
199
204
  code_path=executor_cache,
205
+ name_alias=name_alias,
206
+ image_version=image_version,
200
207
  )
201
208
 
202
209
  try:
@@ -0,0 +1,8 @@
1
+ class ApiException(Exception):
2
+ def __init__(self, message: str) -> None:
3
+ super().__init__(message)
4
+
5
+
6
+ class GraphStillProcessing(Exception):
7
+ def __init__(self) -> None:
8
+ super().__init__("graph is still processing")
@@ -1,6 +1,7 @@
1
1
  import asyncio
2
2
  import json
3
3
  import ssl
4
+ import traceback
4
5
  from concurrent.futures.process import BrokenProcessPool
5
6
  from importlib.metadata import version
6
7
  from typing import Dict, List, Optional
@@ -19,7 +20,11 @@ from indexify.functions_sdk.data_objects import (
19
20
  IndexifyData,
20
21
  RouterOutput,
21
22
  )
23
+ from indexify.functions_sdk.graph_definition import ComputeGraphMetadata
24
+ from indexify.http_client import IndexifyClient
22
25
 
26
+ from ..functions_sdk.image import ImageInformation
27
+ from . import image_dependency_installer
23
28
  from .api_objects import ExecutorMetadata, Task
24
29
  from .downloader import DownloadedInputs, Downloader
25
30
  from .executor_tasks import DownloadGraphTask, DownloadInputTask, ExtractTask
@@ -55,10 +60,28 @@ class ExtractorAgent:
55
60
  executor_id: str,
56
61
  num_workers,
57
62
  code_path: str,
58
- function_worker: FunctionWorker,
59
63
  server_addr: str = "localhost:8900",
60
64
  config_path: Optional[str] = None,
65
+ name_alias: Optional[str] = None,
66
+ image_version: Optional[int] = None,
61
67
  ):
68
+ self.name_alias = name_alias
69
+ self.image_version = image_version
70
+
71
+ self._probe = RuntimeProbes()
72
+
73
+ runtime_probe: ProbeInfo = self._probe.probe()
74
+ self._require_image_bootstrap = (
75
+ True
76
+ if (runtime_probe.is_default_executor and self.name_alias is not None)
77
+ else False
78
+ )
79
+ self._executor_bootstrap_failed = False
80
+
81
+ console.print(
82
+ f"Require Bootstrap? {self._require_image_bootstrap}", style="cyan bold"
83
+ )
84
+
62
85
  self.num_workers = num_workers
63
86
  self._use_tls = False
64
87
  if config_path:
@@ -89,7 +112,12 @@ class ExtractorAgent:
89
112
 
90
113
  self._task_store: TaskStore = TaskStore()
91
114
  self._executor_id = executor_id
92
- self._function_worker = function_worker
115
+ self._function_worker = FunctionWorker(
116
+ workers=num_workers,
117
+ indexify_client=IndexifyClient(
118
+ service_url=f"{self._protocol}://{server_addr}"
119
+ ),
120
+ )
93
121
  self._has_registered = False
94
122
  self._server_addr = server_addr
95
123
  self._base_url = f"{self._protocol}://{self._server_addr}"
@@ -99,7 +127,6 @@ class ExtractorAgent:
99
127
  self._task_reporter = TaskReporter(
100
128
  base_url=self._base_url, executor_id=self._executor_id
101
129
  )
102
- self._probe = RuntimeProbes()
103
130
 
104
131
  async def task_completion_reporter(self):
105
132
  console.print(Text("Starting task completion reporter", style="bold cyan"))
@@ -145,15 +172,55 @@ class ExtractorAgent:
145
172
  async def task_launcher(self):
146
173
  async_tasks: List[asyncio.Task] = []
147
174
  fn_queue: List[FunctionInput] = []
175
+
148
176
  async_tasks.append(
149
177
  asyncio.create_task(
150
178
  self._task_store.get_runnable_tasks(), name="get_runnable_tasks"
151
179
  )
152
180
  )
181
+
153
182
  while True:
154
183
  fn: FunctionInput
155
184
  for fn in fn_queue:
156
185
  task: Task = self._task_store.get_task(fn.task_id)
186
+
187
+ if self._executor_bootstrap_failed:
188
+ completed_task = CompletedTask(
189
+ task=task,
190
+ outputs=[],
191
+ task_outcome="failure",
192
+ )
193
+ self._task_store.complete(outcome=completed_task)
194
+
195
+ continue
196
+
197
+ # Bootstrap this executor. Fail the task if we can't.
198
+ if self._require_image_bootstrap:
199
+ try:
200
+ image_info = await _get_image_info_for_compute_graph(
201
+ task, self._protocol, self._server_addr
202
+ )
203
+ image_dependency_installer.executor_image_builder(
204
+ image_info, self.name_alias, self.image_version
205
+ )
206
+ self._require_image_bootstrap = False
207
+ except Exception as e:
208
+ console.print(
209
+ Text("Failed to bootstrap the executor ", style="red bold")
210
+ + Text(f"Exception: {traceback.format_exc()}", style="red")
211
+ )
212
+
213
+ self._executor_bootstrap_failed = True
214
+
215
+ completed_task = CompletedTask(
216
+ task=task,
217
+ outputs=[],
218
+ task_outcome="failure",
219
+ )
220
+ self._task_store.complete(outcome=completed_task)
221
+
222
+ continue
223
+
157
224
  async_tasks.append(
158
225
  ExtractTask(
159
226
  function_worker=self._function_worker,
@@ -309,18 +376,26 @@ class ExtractorAgent:
309
376
 
310
377
  runtime_probe: ProbeInfo = self._probe.probe()
311
378
 
312
- # Inspect the image
313
- if runtime_probe.is_default_executor:
314
- # install dependencies
315
- # rewrite the image name
316
- pass
317
-
318
379
  executor_version = version("indexify")
380
+
381
+ image_name = (
382
+ self.name_alias
383
+ if self.name_alias is not None
384
+ else runtime_probe.image_name
385
+ )
386
+
387
+ image_version: int = (
388
+ self.image_version
389
+ if self.image_version is not None
390
+ else runtime_probe.image_version
391
+ )
392
+
319
393
  data = ExecutorMetadata(
320
394
  id=self._executor_id,
321
395
  executor_version=executor_version,
322
396
  addr="",
323
- image_name=runtime_probe.image_name,
397
+ image_name=image_name,
398
+ image_version=image_version,
324
399
  labels=runtime_probe.labels,
325
400
  ).model_dump()
326
401
 
@@ -344,6 +419,11 @@ class ExtractorAgent:
344
419
  json=data,
345
420
  headers={"Content-Type": "application/json"},
346
421
  ) as event_source:
422
+ if not event_source.response.is_success:
423
+ resp = await event_source.response.aread().decode('utf-8')
424
+ console.print(f"failed to register: {str(resp)}")
425
+ await asyncio.sleep(5)
426
+ continue
347
427
  console.print(
348
428
  Text("executor registered successfully", style="bold green")
349
429
  )
@@ -372,3 +452,25 @@ class ExtractorAgent:
372
452
  def shutdown(self, loop):
373
453
  self._function_worker.shutdown()
374
454
  loop.create_task(self._shutdown(loop))
455
+
456
+
457
+ async def _get_image_info_for_compute_graph(
458
+ task: Task, protocol, server_addr
459
+ ) -> ImageInformation:
460
+ namespace = task.namespace
461
+ graph_name: str = task.compute_graph
462
+ compute_fn_name: str = task.compute_fn
463
+
464
+ http_client = IndexifyClient(
465
+ service_url=f"{protocol}://{server_addr}", namespace=namespace
466
+ )
467
+ compute_graph: ComputeGraphMetadata = http_client.graph(graph_name)
468
+
469
+ console.print(
470
+ Text(
471
+ f"Compute_fn name {compute_fn_name}, ComputeGraph {compute_graph} \n",
472
+ style="red yellow",
473
+ )
474
+ )
475
+
476
+ return compute_graph.nodes[compute_fn_name].compute_fn.image_information
@@ -21,6 +21,7 @@ class ExecutorMetadata(BaseModel):
21
21
  executor_version: str
22
22
  addr: str
23
23
  image_name: str
24
+ image_version: int
24
25
  labels: Dict[str, Any]
25
26
 
26
27
 
@@ -66,6 +66,7 @@ class ExtractTask(asyncio.Task):
66
66
  init_value=init_value,
67
67
  code_path=code_path,
68
68
  version=task.graph_version,
69
+ invocation_id=task.invocation_id,
69
70
  ),
70
71
  **kwargs,
71
72
  )
@@ -6,6 +6,7 @@ import cloudpickle
6
6
  from pydantic import BaseModel
7
7
  from rich import print
8
8
 
9
+ from indexify import IndexifyClient
9
10
  from indexify.functions_sdk.data_objects import (
10
11
  FunctionWorkerOutput,
11
12
  IndexifyData,
@@ -13,6 +14,7 @@ from indexify.functions_sdk.data_objects import (
13
14
  )
14
15
  from indexify.functions_sdk.indexify_functions import (
15
16
  FunctionCallResult,
17
+ GraphInvocationContext,
16
18
  IndexifyFunctionWrapper,
17
19
  RouterCallResult,
18
20
  )
@@ -44,7 +46,13 @@ class FunctionOutput(BaseModel):
44
46
 
45
47
 
46
48
  def _load_function(
47
- namespace: str, graph_name: str, fn_name: str, code_path: str, version: int
49
+ namespace: str,
50
+ graph_name: str,
51
+ fn_name: str,
52
+ code_path: str,
53
+ version: int,
54
+ invocation_id: str,
55
+ indexify_client: IndexifyClient,
48
56
  ):
49
57
  """Load an extractor to the memory: extractor_wrapper_map."""
50
58
  global function_wrapper_map
@@ -54,18 +62,28 @@ def _load_function(
54
62
  with open(code_path, "rb") as f:
55
63
  code = f.read()
56
64
  pickled_functions = cloudpickle.loads(code)
65
+ context = GraphInvocationContext(
66
+ invocation_id=invocation_id,
67
+ graph_name=graph_name,
68
+ graph_version=str(version),
69
+ indexify_client=indexify_client,
70
+ )
57
71
  function_wrapper = IndexifyFunctionWrapper(
58
- cloudpickle.loads(pickled_functions[fn_name])
72
+ cloudpickle.loads(pickled_functions[fn_name]),
73
+ context,
59
74
  )
60
75
  function_wrapper_map[key] = function_wrapper
61
76
 
62
77
 
63
78
  class FunctionWorker:
64
- def __init__(self, workers: int = 1) -> None:
79
+ def __init__(
80
+ self, workers: int = 1, indexify_client: IndexifyClient = None
81
+ ) -> None:
65
82
  self._executor: concurrent.futures.ProcessPoolExecutor = (
66
83
  concurrent.futures.ProcessPoolExecutor(max_workers=workers)
67
84
  )
68
85
  self._workers = workers
86
+ self._indexify_client = indexify_client
69
87
 
70
88
  async def async_submit(
71
89
  self,
@@ -76,10 +94,19 @@ class FunctionWorker:
76
94
  code_path: str,
77
95
  version: int,
78
96
  init_value: Optional[IndexifyData] = None,
97
+ invocation_id: Optional[str] = None,
79
98
  ) -> FunctionWorkerOutput:
80
99
  try:
81
100
  result = _run_function(
82
- namespace, graph_name, fn_name, input, code_path, version, init_value
101
+ namespace,
102
+ graph_name,
103
+ fn_name,
104
+ input,
105
+ code_path,
106
+ version,
107
+ init_value,
108
+ invocation_id,
109
+ self._indexify_client,
83
110
  )
84
111
  # TODO - bring back running in a separate process
85
112
  except Exception as e:
@@ -113,6 +140,8 @@ def _run_function(
113
140
  code_path: str,
114
141
  version: int,
115
142
  init_value: Optional[IndexifyData] = None,
143
+ invocation_id: Optional[str] = None,
144
+ indexify_client: Optional[IndexifyClient] = None,
116
145
  ) -> FunctionOutput:
117
146
  import io
118
147
  from contextlib import redirect_stderr, redirect_stdout
@@ -131,7 +160,15 @@ def _run_function(
131
160
  try:
132
161
  key = f"{namespace}/{graph_name}/{version}/{fn_name}"
133
162
  if key not in function_wrapper_map:
134
- _load_function(namespace, graph_name, fn_name, code_path, version)
163
+ _load_function(
164
+ namespace,
165
+ graph_name,
166
+ fn_name,
167
+ code_path,
168
+ version,
169
+ invocation_id,
170
+ indexify_client,
171
+ )
135
172
 
136
173
  fn = function_wrapper_map[key]
137
174
  if (
@@ -0,0 +1,64 @@
1
+ import os
2
+ import subprocess
3
+
4
+ from rich.console import Console
5
+ from rich.text import Text
6
+ from rich.theme import Theme
7
+
8
+ from indexify.functions_sdk.image import ImageInformation
9
+
10
+ custom_theme = Theme(
11
+ {
12
+ "info": "cyan",
13
+ "warning": "yellow",
14
+ "error": "red",
15
+ "success": "green",
16
+ }
17
+ )
18
+
19
+ console = Console(theme=custom_theme)
20
+
21
+
22
+ def _record_image_name(name: str, version: int):
23
+ dir_path = os.path.expanduser("~/.indexify/")
24
+
25
+ file_path = os.path.expanduser("~/.indexify/image_name")
26
+ os.makedirs(dir_path, exist_ok=True)
27
+ with open(file_path, "w") as file:
28
+ file.write(name)
29
+
30
+ file_path = os.path.expanduser("~/.indexify/image_version")
31
+ os.makedirs(dir_path, exist_ok=True)
32
+ with open(file_path, "w") as file:
33
+ file.write(str(version))
34
+
35
+
36
+ def _install_dependencies(run_str: str):
37
+ # Throw error to the caller if these subprocesses fail.
38
+ proc = subprocess.run(run_str.split())
39
+ if proc.returncode != 0:
40
+ raise Exception(f"Unable to install dep `{run_str}`")
41
+
42
+
43
+ def executor_image_builder(
44
+ image_info: ImageInformation, name_alias: str, image_version: int
45
+ ):
46
+ console.print(Text("Attempting Executor Bootstrap.", style="red bold"))
47
+
48
+ run_strs = image_info.run_strs
49
+ console.print(Text("Attempting to install dependencies.", style="red bold"))
50
+
51
+ for run_str in run_strs:
52
+ console.print(Text(f"Attempting {run_str}", style="red bold"))
53
+ _install_dependencies(run_str)
54
+
55
+ console.print(Text("Install dependencies done.", style="red bold"))
56
+
57
+ console.print(
58
+ Text(
59
+ f"Recording image name {name_alias} and version {image_version}",
60
+ style="red bold",
61
+ )
62
+ )
63
+
64
+ _record_image_name(name_alias, image_version)
@@ -6,10 +6,12 @@ from typing import Any, Dict, Tuple
6
6
  from pydantic import BaseModel
7
7
 
8
8
  DEFAULT_EXECUTOR = "tensorlake/indexify-executor-default"
9
+ DEFAULT_VERSION = 1
9
10
 
10
11
 
11
12
  class ProbeInfo(BaseModel):
12
13
  image_name: str
14
+ image_version: int
13
15
  python_major_version: int
14
16
  labels: Dict[str, Any] = {}
15
17
  is_default_executor: bool
@@ -18,6 +20,7 @@ class ProbeInfo(BaseModel):
18
20
  class RuntimeProbes:
19
21
  def __init__(self) -> None:
20
22
  self._image_name = self._read_image_name()
23
+ self._image_version = self._read_image_version()
21
24
  self._os_name = platform.system()
22
25
  self._architecture = platform.machine()
23
26
  (
@@ -32,6 +35,13 @@ class RuntimeProbes:
32
35
  return file.read().strip()
33
36
  return DEFAULT_EXECUTOR
34
37
 
38
+ def _read_image_version(self) -> int:
39
+ file_path = os.path.expanduser("~/.indexify/image_version")
40
+ if os.path.exists(file_path):
41
+ with open(file_path, "r") as file:
42
+ return int(file.read().strip())
43
+ return DEFAULT_VERSION
44
+
35
45
  def _get_python_version(self) -> Tuple[int, int]:
36
46
  version_info = sys.version_info
37
47
  return version_info.major, version_info.minor
@@ -50,6 +60,7 @@ class RuntimeProbes:
50
60
 
51
61
  return ProbeInfo(
52
62
  image_name=self._image_name,
63
+ image_version=self._image_version,
53
64
  python_major_version=self._python_version_major,
54
65
  labels=labels,
55
66
  is_default_executor=self._is_default_executor(),
@@ -32,7 +32,7 @@ class TaskReporter:
32
32
  print(
33
33
  f"[bold]task-reporter[/bold] uploading output of size: {len(output.payload)} bytes"
34
34
  )
35
- output_bytes = MsgPackSerializer.serialize(output)
35
+ output_bytes = MsgPackSerializer.serialize(output)
36
36
  fn_outputs.append(
37
37
  ("node_outputs", (nanoid.generate(), io.BytesIO(output_bytes)))
38
38
  )
@@ -33,6 +33,7 @@ from .indexify_functions import (
33
33
  IndexifyFunction,
34
34
  IndexifyFunctionWrapper,
35
35
  IndexifyRouter,
36
+ GraphInvocationContext,
36
37
  )
37
38
  from .local_cache import CacheAwareFunctionWrapper
38
39
  from .object_serializer import get_serializer
@@ -96,9 +97,9 @@ class Graph:
96
97
  return self
97
98
 
98
99
  if issubclass(indexify_fn, IndexifyFunction) and indexify_fn.accumulate:
99
- self.accumulator_zero_values[indexify_fn.name] = (
100
- indexify_fn.accumulate().model_dump()
101
- )
100
+ self.accumulator_zero_values[
101
+ indexify_fn.name
102
+ ] = indexify_fn.accumulate().model_dump()
102
103
 
103
104
  self.nodes[indexify_fn.name] = indexify_fn
104
105
  return self
@@ -214,7 +215,13 @@ class Graph:
214
215
  }
215
216
  self._results[input.id] = outputs
216
217
  enable_cache = kwargs.get("enable_cache", True)
217
- self._run(input, outputs, enable_cache)
218
+ ctx = GraphInvocationContext(
219
+ invocation_id=input.id,
220
+ graph_name=self.name,
221
+ graph_version="1",
222
+ indexify_client=None,
223
+ )
224
+ self._run(input, outputs, enable_cache, ctx)
218
225
  return input.id
219
226
 
220
227
  def _run(
@@ -222,6 +229,7 @@ class Graph:
222
229
  initial_input: IndexifyData,
223
230
  outputs: Dict[str, List[bytes]],
224
231
  enable_cache: bool,
232
+ ctx: GraphInvocationContext,
225
233
  ):
226
234
  accumulator_values = self._accumulator_values[initial_input.id]
227
235
  queue = deque([(self._start_node, initial_input)])
@@ -229,7 +237,7 @@ class Graph:
229
237
  node_name, input = queue.popleft()
230
238
  node = self.nodes[node_name]
231
239
  function_outputs: FunctionCallResult = IndexifyFunctionWrapper(
232
- node
240
+ node, context=ctx
233
241
  ).invoke_fn_ser(node_name, input, accumulator_values.get(node_name, None))
234
242
  if function_outputs.traceback_msg is not None:
235
243
  print(function_outputs.traceback_msg)
@@ -2,7 +2,6 @@ import inspect
2
2
  import re
3
3
  import sys
4
4
  import traceback
5
- from abc import ABC, abstractmethod
6
5
  from functools import update_wrapper
7
6
  from typing import (
8
7
  Any,
@@ -18,14 +17,37 @@ from typing import (
18
17
  )
19
18
 
20
19
  import msgpack
21
- from pydantic import BaseModel
20
+ from pydantic import BaseModel, Field, PrivateAttr, model_validator
22
21
  from typing_extensions import get_type_hints
23
22
 
24
- from .data_objects import IndexifyData, RouterOutput
23
+ from .data_objects import IndexifyData
25
24
  from .image import DEFAULT_IMAGE_3_10, Image
26
25
  from .object_serializer import CloudPickleSerializer, get_serializer
27
26
 
28
27
 
28
+ class GraphInvocationContext(BaseModel):
29
+ invocation_id: str
30
+ graph_name: str
31
+ graph_version: str
32
+ indexify_client: Optional[Any] = Field(default=None) # avoids circular import
33
+ _local_state: Dict[str, Any] = PrivateAttr(default_factory=dict)
34
+
35
+ def set_state_key(self, key: str, value: Any) -> None:
36
+ if self.indexify_client is None:
37
+ self._local_state[key] = value
38
+ return
39
+ self.indexify_client.set_state_key(
40
+ self.graph_name, self.invocation_id, key, value
41
+ )
42
+
43
+ def get_state_key(self, key: str) -> Any:
44
+ if self.indexify_client is None:
45
+ return self._local_state.get(key)
46
+ return self.indexify_client.get_state_key(
47
+ self.graph_name, self.invocation_id, key
48
+ )
49
+
50
+
29
51
  def format_filtered_traceback(exc_info=None):
30
52
  """
31
53
  Format a traceback excluding indexify_functions.py lines.
@@ -205,10 +227,15 @@ class RouterCallResult(BaseModel):
205
227
 
206
228
 
207
229
  class IndexifyFunctionWrapper:
208
- def __init__(self, indexify_function: Union[IndexifyFunction, IndexifyRouter]):
209
- self.indexify_function: Union[IndexifyFunction, IndexifyRouter] = (
210
- indexify_function()
211
- )
230
+ def __init__(
231
+ self,
232
+ indexify_function: Union[IndexifyFunction, IndexifyRouter],
233
+ context: GraphInvocationContext,
234
+ ):
235
+ self.indexify_function: Union[
236
+ IndexifyFunction, IndexifyRouter
237
+ ] = indexify_function()
238
+ self.indexify_function._ctx = context
212
239
 
213
240
  def get_output_model(self) -> Any:
214
241
  if not isinstance(self.indexify_function, IndexifyFunction):
@@ -322,3 +349,14 @@ class IndexifyFunctionWrapper:
322
349
  payload = list(payload.values())[0]
323
350
  return arg_type.model_validate(payload)
324
351
  return payload
352
+
353
+
354
+ def get_ctx() -> GraphInvocationContext:
355
+ frame = inspect.currentframe()
356
+ caller_frame = frame.f_back.f_back
357
+ function_instance = caller_frame.f_locals["self"]
358
+ del frame
359
+ del caller_frame
360
+ if isinstance(function_instance, IndexifyFunctionWrapper):
361
+ return function_instance.indexify_function._ctx
362
+ return function_instance._ctx
@@ -10,7 +10,7 @@ from httpx_sse import connect_sse
10
10
  from pydantic import BaseModel, Json
11
11
  from rich import print
12
12
 
13
- from indexify.error import ApiException
13
+ from indexify.error import ApiException, GraphStillProcessing
14
14
  from indexify.functions_sdk.data_objects import IndexifyData
15
15
  from indexify.functions_sdk.graph import ComputeGraphMetadata, Graph
16
16
  from indexify.functions_sdk.indexify_functions import IndexifyFunction
@@ -36,7 +36,9 @@ class GraphOutputMetadata(BaseModel):
36
36
 
37
37
 
38
38
  class GraphOutputs(BaseModel):
39
+ status: str
39
40
  outputs: List[GraphOutputMetadata]
41
+ cursor: Optional[str] = None
40
42
 
41
43
 
42
44
  class IndexifyClient:
@@ -196,6 +198,23 @@ class IndexifyClient:
196
198
  namespaces.append(item["name"])
197
199
  return namespaces
198
200
 
201
+ def set_state_key(
202
+ self, compute_graph: str, invocation_id: str, key: str, value: Json
203
+ ) -> None:
204
+ response = self._post(
205
+ f"internal/namespaces/{self.namespace}/compute_graphs/{compute_graph}/invocations/{invocation_id}/ctx",
206
+ json={"key": key, "value": value},
207
+ )
208
+ response.raise_for_status()
209
+
210
+ def get_state_key(self, compute_graph: str, invocation_id: str, key: str) -> Json:
211
+ response = self._get(
212
+ f"internal/namespaces/{self.namespace}/compute_graphs/{compute_graph}/invocations/{invocation_id}/ctx",
213
+ json={"key": key},
214
+ )
215
+ response.raise_for_status()
216
+ return response.json().get("value")
217
+
199
218
  @classmethod
200
219
  def new_namespace(
201
220
  cls, namespace: str, server_addr: Optional[str] = "http://localhost:8900"
@@ -249,6 +268,9 @@ class IndexifyClient:
249
268
  data=ser_input,
250
269
  params=params,
251
270
  ) as event_source:
271
+ if not event_source.response.is_success:
272
+ resp = event_source.response.read().decode("utf-8")
273
+ raise Exception(f"failed to invoke graph: {resp}")
252
274
  for sse in event_source.iter_sse():
253
275
  obj = json.loads(sse.data)
254
276
  for k, v in obj.items():
@@ -329,6 +351,8 @@ class IndexifyClient:
329
351
  )
330
352
  response.raise_for_status()
331
353
  graph_outputs = GraphOutputs(**response.json())
354
+ if graph_outputs.status == "pending":
355
+ raise GraphStillProcessing()
332
356
  outputs = []
333
357
  for output in graph_outputs.outputs:
334
358
  if output.compute_fn == fn_name:
@@ -84,7 +84,12 @@ class RemoteGraph:
84
84
  return cls(name=g.name, server_url=server_url, client=client)
85
85
 
86
86
  @classmethod
87
- def by_name(cls, name: str, server_url: Optional[str] = DEFAULT_SERVICE_URL, client: Optional[IndexifyClient] = None):
87
+ def by_name(
88
+ cls,
89
+ name: str,
90
+ server_url: Optional[str] = DEFAULT_SERVICE_URL,
91
+ client: Optional[IndexifyClient] = None,
92
+ ):
88
93
  """
89
94
  Create a handle to call a RemoteGraph by name.
90
95
 
@@ -104,7 +109,7 @@ class RemoteGraph:
104
109
  ) -> List[Any]:
105
110
  """
106
111
  Returns the extracted objects by a graph for an ingested object.
107
-
112
+
108
113
  - If the extractor name is provided, only the objects extracted by that extractor are returned.
109
114
  - If the extractor name is not provided, all the extracted objects are returned for the input object.
110
115
 
@@ -1,6 +1,6 @@
1
1
  [tool.poetry]
2
2
  name = "indexify"
3
- version = "0.2.21"
3
+ version = "0.2.23"
4
4
  description = "Python Client for Indexify"
5
5
  authors = ["Tensorlake Inc. <support@tensorlake.ai>"]
6
6
  license = "Apache 2.0"
@@ -1,3 +0,0 @@
1
- class ApiException(Exception):
2
- def __init__(self, message: str) -> None:
3
- super().__init__(message)
File without changes
File without changes