Preserve type of decorated methods/classes

[nschank]

What does this PR do ?

Use a TypeVar to ensure that decorators don't obfuscate the structure of the thing they're decorating.

Before this change, the type-checker and IDEs would treat any functions/classes decorated by these decorators as completely opaque (no known signature/methods/anything); now, these decorators are transparent and the type-checker knows the underlying class/type is unchanged.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[janEbert]

In general, there are also other wrapper functions/decorators in this file. Is there a reason they weren't similarly modified? :)

[janEbert]

Shouldn't this still return a Callable, or more specifically, a Callable[[_Wrapped], _Wrapped]?

[nschank]

Yes, but I was leaning on type inference - for decorators you often need to, unfortunately. I updated these two as requested, but experimental_fn is an example where actually marking the return type would obfuscate the optional additional parameter in the returned function - in this case we could technically still get around it with a Protocol (although very excessive lol), but for complex decorators you can end up not being able to specify the type hint correctly at all!

[janEbert]

Oh curious, okay. So in the case of experimental_fn, it's better to have Any/Unknown, as the type checker figures things out better by itself? I don't really understand why it works for other places but not there.

Do you have an easy test setup that you're using for this? You made me also want to play around to get a better grasp. :)

[nschank]

I use pyright locally! Specifically I have pip install pyright==1.1.408 in my devcontainer setup, then the Pylance extension should start giving you the red squigglies. For this particular PR, I mostly ended up comparing what VSCode autocomplete/mouse-over showed me for decorated functions - before, it would have no idea what the arguments were, and now it does.

In the case of experimental_fn, we actually want to not specify a return type at all. The type checker will infer the return type from the contents of the function, and the "name" of the type is too hard/annoying to spell out.

Here's an example of a function which cannot have its return type decorated, because python's type system won't let you "write" the type correctly (without https://peps.python.org/pep-0827/):

def without_parameters[T, R](decorator: Callable[[T], R], /):  # noqa: ANN201
  """Converts a method which directly accepts an object into a decorator.

  The provided decorator should accept a decorated object as its first argument.
  The returned decorator can then be used either directly on an object (no
  parentheses) or with parentheses.

  Example:
    ```
    @decorator_builder.without_parameters
    def my_decorator(func: Callable[[int], int], /) -> Callable[[int], int]:
       ...

    @my_decorator
    def func1(x: int) -> int:
      return x + 1

    @my_decorator()
    def func2(x: int) -> int:
      return x + 2
    ```
  """

  @overload
  def wrapper(decorated: None = None, /) -> Callable[[T], R]: ...

  @overload
  def wrapper(decorated: T, /) -> R: ...

  @functools.wraps(decorator)
  def wrapper(decorated: T | None = None, /) -> R | Callable[[T], R]:
    if decorated is None:
      return decorator
    else:
      return decorator(decorated)

  return wrapper

If you just leave the annotation unset, the type-checker infers the correct type for you. Any would be bad because it would just tell the type checker it shouldn't attempt to infer the type!

[janEbert]

Thanks a lot! That gives me a much better sense of the issues, although I still don't understand the difference between the functions, where it works for some, but not others. I'll play around myself a bit, thanks again for the detailed summary!

[nschank]

Haha mostly because this file is 2500 lines so I didn't really want to go yak shaving too hard, but fair. I updated a few other ones where it was immediately obvious how to, but there's some decorators which would require more work so I skipped them.

[janEbert]

Thank you!