Fix ty suppressions from #405 with real type fixes (#408)

[lucianlavric]

Summary

#405 introduced ty: ignore suppressions (tracked in #408) where the type fixes were not trivial. This removes every #408-linked suppression by fixing the underlying types instead: declared buffer types instead of __getattr__ unions, fail-fast narrowing instead of ignores, and one real mismatch — infer_task_inputs declared decoder: LinkPredictionDecoder but assigns the bound decode method; it is now typed as the callable it actually is.

Approach

• Declare attribute types for registered buffers (LightGCN._layer_weights, the early-stop test's _DummyModel.foo) so ty resolves them as Tensor rather than the Tensor | Module union from nn.Module.__getattr__.
• Make the GnnModel protocol runtime_checkable and narrow through it in InferencerV1; keep the dynamic graph_backend attach in the GraphSAGE spec explicit via setattr.
• Narrow with fail-fast asserts where the runtime type is guaranteed by construction: the model in infer_task_inputs (LinkPredictionGNN), and ModuleDict values in NodeAnchorBasedLinkPredictionTasks._get_all_tasks.
• FeatureInteraction.reset_parameters: bind via getattr so callable() narrows the optional method.
• KDD tutorial: from torch.multiprocessing.spawn import spawn — the dotted call resolves to the same-named submodule, which type checkers reject as non-callable.

User-facing changes

None intended. Invalid model types now fail fast with descriptive assert messages instead of failing deeper in the call.

Verification

• make type_check: no new diagnostics versus main (the same 48 pre-existing macOS-environment diagnostics appear on both — uninstallable-on-macOS imports and two Linux-only os/psutil APIs).
• make check_format_py: clean.
• Full make unit_test_py is not runnable on macOS (graphlearn_torch ships in the Docker/dev images only; the darwin env lacks tensorflow_metadata/torchrec), unchanged from main — relying on PR CI for the test gate.
• Runtime patterns validated standalone with pure torch: buffer registration and += behavior unchanged with the declared annotation; runtime_checkable GnnModel isinstance correctly passes after setattr and fails on modules without the attribute; the explicitly imported spawn is callable.

[kmontemayor2-sc]

Thanks for removing the ty suppressions, and sorry for the late review 😅 . I found one blocking runtime regression in the current implementation: the new LinkPredictionGNN import creates a circular dependency that prevents the node-anchor link-prediction modeling spec from importing.

[kmontemayor2-sc]

This import introduces a circular dependency:

utils/infer.py
  → models/pyg/link_prediction.py
  → models/layers/task.py
  → utils/infer.py

The same import succeeds at the PR’s base commit. This blocks NodeAnchorBasedLinkPredictionModelingTaskSpec from loading.
Please avoid importing the concrete model here at runtime. A small structural Protocol describing the required decode and tasks.result_types members would both avoid the cycle and preserve the existing model contract. A TYPE_CHECKING-only import plus a cast would also avoid the cycle, although it would not provide runtime validation.

[lucianlavric]

Confirmed the cycle (infer → link_prediction → task → infer) and removed the import. infer_task_inputs now narrows against a module-private runtime_checkable Protocol declaring decode and tasks.result_types, defined in infer.py, so no concrete model is imported. Verified NodeAnchorBasedLinkPredictionModelingTaskSpec imports again.

[kmontemayor2-sc]

This changes the function from accepting any compatible nn.Module to requiring one specific legacy LinkPredictionGNN class.
Previously, a custom model or wrapper worked as long as it exposed decode and tasks.result_types, which matches the function’s generic nn.Module | DDP signature. The new check rejects compatible custom implementations, and the required concrete class is itself marked deprecated.
Could we type this against a structural protocol instead? If runtime validation is required, please raise an explicit TypeError rather than using assert, since the assertion disappears under python -O.

[lucianlavric]

Switched to the structural Protocol, so any module exposing decode and tasks.result_types is accepted again, and replaced the assert with raise TypeError. Made the same assert → TypeError change to the two other guards this PR introduced (InferencerV1 and _get_all_tasks) for consistency.

[kmontemayor2-sc]

This declaration is safe because decoder is assigned unconditionally below, but I would initialize it where it is declared rather than leave a declaration-only local:

decoder: Callable[
    [torch.Tensor, torch.Tensor], torch.Tensor
] = base_model.decode

Alternatively, if the narrowed model type exposes the correct signature, decoder = base_model.decode should be sufficient.
I would not add a None or no-op default here: the decoder is required and model-specific, so such a default would weaken the invariant and could hide an invalid model.

[lucianlavric]

Dropped the declaration; decoder and batch_result_types are inferred from the Protocol now. One wrinkle: after the isinstance narrowing ty types the value as Module & _LinkPredictionModel and resolves attributes through Module.getattr as well, giving set | Tensor | Module, so a cast to the Protocol after the runtime check is still needed. Same reason InferencerV1 casts to GnnModel. Comment in the code explains it.

[kmontemayor2-sc]

I have a preference against declaration only statements like this, can we avoid these?

[lucianlavric]

Replaced the annotation + register_buffer pair with self._layer_weights = nn.Buffer(...) (torch ≥ 2.4, repo pins 2.8). Assignment registers the buffer and ty infers Tensor from the value. Same change to _DummyModel.foo in early_stop_test.py. Checked state_dict keys, persistence, requires_grad, and in-place += are unchanged; early-stop tests pass.

[lucianlavric]

Addressed all four. The blocking cycle is fixed and verified by importing the node-anchor spec. make type_check matches main's baseline and make check_format_py is clean. Ready for another look.