[FEA] `EngineContext` base class to pass distributed context into I/O source plguins
- 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
Assessment
This issue has not been assessed yet.