NVIDIA / NVIDIA/cudf

[FEA] `EngineContext` base class to pass distributed context into I/O source plguins

Open
#23,042 1 comment 0 reactions 1 assignee Claimed by @madsbk View on GitHub
cudf-polars feature request
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

**Is your feature request related to a problem? Please describe.**

`register_io_source` sources have no way to receive runtime context from the engine. Polars calls the source with a fixed signature `(with_columns, predicate, n_rows, batch_size)` and there is no hook for an engine to inject additional arguments like which rank the source is running on.

cudf-polars works around this today by inspecting `__closure__` of the Polars-generated wrapper function to find a `RankAwareSource` instance and calling it directly with `rank`/`nranks`. This depends on undocumented Polars internals and will break if the wrapper's closure layout changes.

**Describe the solution you'd like**

Polars should define a minimal `EngineContext` base class and pass it into the source call when an engine provides one:

```python
# In Polars
class EngineContext:
pass
```

Engines subclass it with their own fields:

```python
# In cudf-polars
@dataclass
class CudfEngineContext(EngineContext):
rank: int = 0
nranks: int = 1
...
```

Polars passes it as an optional kwarg when an engine provides one:

```python
chunks = source(with_columns, predicate, n_rows, batch_size, engine_context=context)
```

Sources opt in by declaring the parameter:

```python
def my_source(with_columns, predicate, n_rows, batch_size,
engine_context: EngineContext = EngineContext()):
rank = engine_context.rank if isinstance(engine_context, CudfEngineContext) else 0
yield data.slice(rank)
```

Sources that do not declare `engine_context` are called without it, so existing sources are unaffected.

**Describe alternatives you've considered**

`RankAwareSource` (current workaround in [cudf#22867](https://github.com/rapidsai/cudf/pull/22867)): requires subclassing, forbids wrapping the instance, and relies on closure inspection.

**Additional context**
N/A

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.