Skip to content

[FEA]: Precise types for Graph.__getitem__ executable node views #2482

Description

@Andy-Jost

Is this a duplicate?

Not a duplicate. This is a specific sub-item of the typing EPIC #469.

Area

cuda.core

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

Graph.__getitem__ is typed as returning the base class ExecutableGraphNode, but at run time it returns a specific view: ExecutableKernelNode, ExecutableMemsetNode, ExecutableMemcpyNode, ExecutableChildGraphNode, ExecutableEventRecordNode, ExecutableEventWaitNode, or ExecutableHostCallbackNode.

The base class declares no update, and is_enabled, enable, and disable exist only on the kernel, memset, and memcpy views, because CUDA restricts cuGraphNodeSetEnabled to those node types. Since cuda.core ships py.typed alongside its generated stubs, type checkers reject correct code:

error: "ExecutableGraphNode" has no attribute "update"  [attr-defined]
error: "ExecutableGraphNode" has no attribute "disable"  [attr-defined]

Editors cannot complete update, enable, or disable after graph[node]. either, so the feature is hard to discover. Affected users must cast:

view = cast(ExecutableKernelNode, graph[node])
view.update(config=cfg, kernel=k, args=(a, b))

Nothing is wrong at run time. Only the advertised type is vaguer than the real one.

Describe the solution you'd like

Declare Graph.__getitem__ with @overload, keyed on the source node type, one for each of the seven supported mappings, KernelNode to ExecutableKernelNode and so on. Indexing with a node type that has no executable view, which raises TypeError today, would then also be reported statically.

Describe alternatives you've considered

  • A permissive update(*args, **kwargs) on the base class. This silences the checker but discards the precise per-node signatures, which are the valuable part of the design.
  • Hand-editing the generated .pyi. Not viable: the stub generator rewrites whole files and offers no preserve or skip option.
  • Moving the implementation into a base cdef class so that only the overloads remain in the derived class. This breaks at run time, because the last @overload placeholder becomes the real slot.

Additional context

The work is blocked upstream, not by design. Writing the overloads in the .pyx does work: Cython compiles the pattern, the final implementation wins at run time, and the generator copies the decorators into the stub. The generator also copies the implementation, which mypy rejects with An implementation for an overloaded function is not allowed in a stub file. Reported as jon-edward/stubgen-pyx#48, with a fix proposed in jon-edward/stubgen-pyx#49.

Adoption after that lands takes two steps, best kept separate. First, raise the stubgen-pyx pin in .pre-commit-config.yaml, currently 0.2.6 against 0.2.17 upstream; that jump will churn stubs from unrelated generator improvements and deserves its own PR. Second, add the overloads. No API change or deprecation is involved, since overloads only narrow the return type.

Metadata

Metadata

Assignees

Labels

P2Low priority - Nice to havecuda.coreEverything related to the cuda.core moduleenhancementAny code-related improvements

Type

No type

Projects

No projects

Relationships

None yet

Development

No branches or pull requests

Issue actions