class DocumentCaster:
"""Cast source documents to :class:`GraphContainer` via ingestion resources."""
# Concurrent sources may interleave batches of different resources; keeping a
# small set of pools alive avoids tearing one down (and re-importing graflo in
# every worker) each time the resource changes.
_MAX_LIVE_POOLS = 4
def __init__(self, ingestion_model: IngestionModel) -> None:
self.ingestion_model = ingestion_model
# Worker startup is dominated by importing graflo and rebuilding the runtime,
# which costs far more than casting a batch. Pools are therefore kept alive
# across batches, one per (resource, workers), LRU-bounded.
self._pools: OrderedDict[tuple[str, int], ProcessPoolExecutor] = OrderedDict()
# Serializing a resource + its schema is not free, and it is identical for
# every batch of that resource.
self._spec_cache: dict[str, CastSpec | None] = {}
self._cache_lock = threading.Lock()
def close(self) -> None:
"""Shut down all cast worker pools, if any were started."""
with self._cache_lock:
pools = list(self._pools.values())
self._pools.clear()
for pool in pools:
pool.shutdown(wait=True)
def _get_pool(
self, spec: CastSpec, workers: int, *, resource_name: str
) -> ProcessPoolExecutor:
key = (resource_name, workers)
evicted: ProcessPoolExecutor | None = None
with self._cache_lock:
pool = self._pools.get(key)
if pool is not None:
self._pools.move_to_end(key)
return pool
pool = ProcessPoolExecutor(
max_workers=workers,
mp_context=_MP_CONTEXT,
initializer=init_worker,
initargs=(spec,),
)
self._pools[key] = pool
if len(self._pools) > self._MAX_LIVE_POOLS:
_, evicted = self._pools.popitem(last=False)
if evicted is not None:
evicted.shutdown(wait=True)
return pool
async def cast_batch(
self,
data: Iterable[Any],
resource_name: str | None,
*,
params: IngestionParams,
allowed_vertex_names: set[str] | None = None,
) -> CastBatchResult:
runtime = self.ingestion_model.fetch_resource(resource_name)
resolved_name = runtime.name
vertex_filter = resolve_effective_vertex_names(
runtime.collect_vertex_names(),
allowed_vertex_names=allowed_vertex_names,
)
docs = [_coerce_doc(doc) for doc in data]
mode = self._resolve_cast_mode(params, n_docs=len(docs))
workers = max(1, params.n_cores)
if mode == "process":
spec = self._cast_spec(runtime, params=params)
if spec is None:
mode = "thread"
else:
try:
return await self._cast_batch_in_processes(
spec,
docs,
attached_by=runtime.attached_selectors,
workers=workers,
resolved_name=resolved_name,
vertex_filter=vertex_filter,
params=params,
)
except WorkerCastError:
# on_doc_error="fail": a document failed, not the pool.
raise
except Exception as exc:
logger.warning(
"Process-pool casting unavailable (%s: %s); falling back to "
"in-process casting for resource %r.",
type(exc).__name__,
exc,
runtime.name,
)
mode = "thread"
# Dynamic edge feedback mutates the shared edge_config *during* casting;
# spreading documents over threads would race those registrations, so the
# fallback for that configuration is single-threaded, not "thread".
if mode == "thread" and params.dynamic_edges:
mode = "inline"
raw = await self._cast_documents(
runtime, docs, n_cores=workers if mode == "thread" else 1
)
cast_results, failures = self._collect_cast_results(
raw,
docs,
on_doc_error=params.on_doc_error,
resolved_name=resolved_name,
params=params,
)
done = [r for r in cast_results if isinstance(r, ResourceCastResult)]
graph = self._assemble(
runtime, done, vertex_filter=vertex_filter, params=params
)
return CastBatchResult(graph=graph, failures=failures)
async def cast_batch_by_key(
self,
data: Sequence[Any],
keys: Sequence[Hashable],
resource_name: str | None,
*,
params: IngestionParams,
allowed_vertex_names: set[str] | None = None,
) -> tuple[list[tuple[Hashable, GraphContainer]], list[DocCastFailure]]:
"""Cast *data* into one container per distinct key of its records.
``keys[i]`` is the key of ``data[i]``. Each container holds only what
records of its key cast to; containers come in the order their key is
first seen. Failures carry their record's index in *data*.
"""
positions: dict[Hashable, list[int]] = {}
for index, key in enumerate(keys):
positions.setdefault(key, []).append(index)
if len(positions) <= 1:
result = await self.cast_batch(
data,
resource_name,
params=params,
allowed_vertex_names=allowed_vertex_names,
)
return [(key, result.graph) for key in positions], result.failures
runtime = self.ingestion_model.fetch_resource(resource_name)
docs = [_coerce_doc(doc) for doc in data]
mode = self._resolve_cast_mode(params, n_docs=len(docs))
if mode == "process" and self._cast_spec(runtime, params=params) is None:
mode = "thread"
if mode == "process":
# A worker-cast resource has no order-dependent steps, so each key's
# records cast as a batch of their own.
parts: list[tuple[Hashable, GraphContainer]] = []
failures: list[DocCastFailure] = []
for key, indices in positions.items():
result = await self.cast_batch(
[docs[i] for i in indices],
resource_name,
params=params,
allowed_vertex_names=allowed_vertex_names,
)
parts.append((key, result.graph))
failures.extend(
f.model_copy(update={"doc_index": indices[f.doc_index]})
for f in result.failures
)
return parts, failures
# In process, records cast once in source order (dynamic edges depend
# on it), and the results are split by key afterwards.
if mode == "thread" and params.dynamic_edges:
mode = "inline"
workers = max(1, params.n_cores)
raw = await self._cast_documents(
runtime, docs, n_cores=workers if mode == "thread" else 1
)
cast_results, failures = self._collect_cast_results(
raw,
docs,
on_doc_error=params.on_doc_error,
resolved_name=runtime.name,
params=params,
)
cast_at = dict(
zip(
(
i
for i, item in enumerate(raw)
if not isinstance(item, BaseException)
),
cast_results,
)
)
vertex_filter = resolve_effective_vertex_names(
runtime.collect_vertex_names(),
allowed_vertex_names=allowed_vertex_names,
)
parts = []
for key, indices in positions.items():
done = [
result
for i in indices
if isinstance(result := cast_at.get(i), ResourceCastResult)
]
parts.append(
(
key,
self._assemble(
runtime, done, vertex_filter=vertex_filter, params=params
),
)
)
return parts, failures
@staticmethod
def _assemble(
runtime: ResourceRuntime,
done: list[ResourceCastResult],
*,
vertex_filter: set[str] | None,
params: IngestionParams,
) -> GraphContainer:
"""Fold cast records into one container, filtered as the params ask."""
graph = GraphContainer.from_docs_list(
[r.entities for r in done],
attached=[r.attached for r in done],
attached_by=runtime.attached_selectors,
)
filter_graph_container_by_vertices_inplace(
graph, allowed_vertex_names=vertex_filter
)
if params.drop_empty_identity_docs:
filter_graph_container_drop_empty_identity_inplace(
graph,
vertex_config=runtime.vertex_config,
edge_derivation=runtime.edge_derivation,
)
return graph
def _resolve_cast_mode(
self, params: IngestionParams, *, n_docs: int
) -> Literal["inline", "thread", "process"]:
"""Resolve ``auto`` to a concrete executor for this batch.
``auto`` reaches for processes only when there is real CPU work to win
back: more than one core requested, no dynamic-edge feedback (which
cannot cross a process boundary), and enough documents that each worker
gets at least a full chunk to amortize pickling documents out and the
folded entities back. Workers fold and filter their own chunks, so the
parent-side serial cost that used to cancel out the process-pool gain
(~3x on casting over 8 workers) no longer applies.
"""
mode = params.cast_executor
if mode != "auto":
return mode
if (
params.n_cores > 1
and not params.dynamic_edges
and n_docs >= params.n_cores * _MIN_DOCS_PER_CAST_CHUNK
):
return "process"
return "inline"
def _collect_cast_results(
self,
raw: list[ResourceCastResult | Exception],
docs: list[dict[str, Any]],
*,
on_doc_error: Literal["fail", "skip"],
resolved_name: str,
params: IngestionParams,
) -> tuple[list[ResourceCastResult | BaseException], list[DocCastFailure]]:
if on_doc_error == "fail":
# Preserve "the first failing document aborts the batch", by document
# order rather than by whichever worker happened to fail first.
for item in raw:
if isinstance(item, BaseException):
raise item
cast_results: list[ResourceCastResult | BaseException] = []
failures: list[DocCastFailure] = []
for i, item in enumerate(raw):
doc = docs[i]
if isinstance(item, BaseException):
failures.append(
_doc_failure_from_exception(
resource_name=resolved_name,
doc_index=i,
doc=doc,
exc=item,
doc_keys=params.doc_error_preview_keys,
doc_preview_max_bytes=params.doc_error_preview_max_bytes,
)
)
continue
failures.extend(
_transform_failures_to_doc_cast_failures(
resource_name=resolved_name,
doc_index=i,
doc=doc,
transform_failures=item.transform_failures,
doc_keys=params.doc_error_preview_keys,
doc_preview_max_bytes=params.doc_error_preview_max_bytes,
)
)
cast_results.append(item)
return cast_results, failures
def _cast_spec(
self, runtime: ResourceRuntime, *, params: IngestionParams
) -> CastSpec | None:
"""Serializable description of *runtime*, or ``None`` if it cannot cross.
Dynamic edge feedback registers new edges on the shared ``edge_config``
*during* casting. In a worker process those registrations would be invisible
to the parent's db-aware projection, so that configuration stays in-process.
So does a resource with a step that names its edges per document: each
worker would register them on its own copy of the runtime.
"""
if params.dynamic_edges or runtime.has_dynamic_edge_steps:
return None
with self._cache_lock:
if runtime.name in self._spec_cache:
return self._spec_cache[runtime.name]
spec = self._build_cast_spec(runtime, params=params)
with self._cache_lock:
return self._spec_cache.setdefault(runtime.name, spec)
def _build_cast_spec(
self, runtime: ResourceRuntime, *, params: IngestionParams
) -> CastSpec | None:
try:
return CastSpec(
resource=runtime.config.to_dict(skip_defaults=False),
vertex_config=runtime.vertex_config.to_dict(skip_defaults=False),
edge_config=runtime.edge_config.to_dict(skip_defaults=False),
transforms=[
proto.to_dict(skip_defaults=False)
for proto in self.ingestion_model.transforms
],
strict_references=params.strict_references,
target_db_flavor=_flavor_value(runtime.target_db_flavor),
)
except Exception as exc:
logger.warning(
"Resource %r cannot be described for worker processes (%s: %s); "
"casting it in process instead.",
runtime.name,
type(exc).__name__,
exc,
)
return None
async def _cast_batch_in_processes(
self,
spec: CastSpec,
docs: list[dict[str, Any]],
*,
attached_by: dict[str, str],
workers: int,
resolved_name: str,
vertex_filter: set[str] | None,
params: IngestionParams,
) -> CastBatchResult:
"""Cast *docs* across worker processes, preserving document order.
Workers fold and filter their own contiguous chunks; the parent only
concatenates the partial containers in chunk order, which yields the
identical graph to a serial fold over the whole batch.
"""
n_workers = default_worker_count(workers)
size = -(-len(docs) // n_workers)
chunks = [docs[i : i + size] for i in range(0, len(docs), size)]
offsets = list(range(0, len(docs), size))
filter_names = sorted(vertex_filter) if vertex_filter is not None else None
loop = asyncio.get_running_loop()
pool = self._get_pool(spec, n_workers, resource_name=resolved_name)
# Only documents cross here: the workers were primed with the spec.
parts = await asyncio.gather(
*[
loop.run_in_executor(
pool,
cast_chunk,
chunk,
filter_names,
params.drop_empty_identity_docs,
)
for chunk in chunks
]
)
if params.on_doc_error == "fail":
# Chunks are contiguous and in order, so the first chunk reporting an
# error holds the failing document with the smallest global index.
for part in parts:
errors = part[4]
if errors:
_, error = min(errors, key=lambda e: e[0])
raise WorkerCastError(error)
vertices: dict[str, list] = {}
edges: dict[Any, list] = {}
linear: list[Any] = []
attached: dict[str, list] = {}
failures: list[DocCastFailure] = []
for offset, part in zip(offsets, parts):
(
part_vertices,
part_edges,
part_linear,
transform_failures,
errors,
part_attached,
) = part
for k, v in part_vertices.items():
vertices.setdefault(k, []).extend(v)
for k, v in part_attached.items():
attached.setdefault(k, []).extend(v)
for k, v in part_edges.items():
edges.setdefault(k, []).extend(v)
linear.extend(part_linear)
# A document either failed outright or carries transform failures,
# never both; merging by local index restores document order.
per_doc: list[tuple[int, tuple[str, str, str] | None, list]] = [
(i, error, []) for i, error in errors
] + [(i, None, tfails) for i, tfails in transform_failures]
for local_idx, error, tfails in sorted(per_doc, key=lambda t: t[0]):
doc = docs[offset + local_idx]
if error is not None:
failures.append(
_doc_failure_from_exception(
resource_name=resolved_name,
doc_index=offset + local_idx,
doc=doc,
exc=WorkerCastError(error),
doc_keys=params.doc_error_preview_keys,
doc_preview_max_bytes=params.doc_error_preview_max_bytes,
)
)
continue
failures.extend(
_transform_failures_to_doc_cast_failures(
resource_name=resolved_name,
doc_index=offset + local_idx,
doc=doc,
transform_failures=tfails,
doc_keys=params.doc_error_preview_keys,
doc_preview_max_bytes=params.doc_error_preview_max_bytes,
)
)
graph = GraphContainer(
vertices=vertices,
edges=edges,
linear=linear,
attached=attached,
attached_by=dict(attached_by),
)
return CastBatchResult(graph=graph, failures=failures)
async def _cast_documents(
self,
runtime: ResourceRuntime,
docs: list[dict[str, Any]],
*,
n_cores: int,
) -> list[ResourceCastResult | Exception]:
"""Cast *docs* in document order, spreading contiguous slices over workers.
A slice per worker rather than a task per document: ``cast_document`` is
GIL-bound, so the dispatch was costing more than the work. Anything beyond
a trivially small batch still crosses into a worker thread once, to keep
the event loop free for batch prefetch.
"""
if not docs:
return []
workers = max(1, n_cores)
if workers == 1 and len(docs) <= _MIN_DOCS_PER_CAST_CHUNK:
return _cast_chunk(runtime, docs)
max_useful = -(-len(docs) // _MIN_DOCS_PER_CAST_CHUNK)
n_chunks = max(1, min(workers, max_useful))
size = -(-len(docs) // n_chunks)
chunks = [docs[i : i + size] for i in range(0, len(docs), size)]
parts = await asyncio.gather(
*[asyncio.to_thread(_cast_chunk, runtime, chunk) for chunk in chunks],
return_exceptions=True,
)
out: list[ResourceCastResult | Exception] = []
for part in parts:
if isinstance(part, BaseException):
# A chunk raised outside per-document capture (cancellation, exit):
# abort the batch rather than reporting it against one document.
raise part
out.extend(part)
return out