Properly handle decorated functions

[dolfinus]

Motivation

I've implemented some decorator which wraps the main function, accepts argument of type DictConfig to have access to job config, and performs some operations:

from functools import wraps
from omegaconf import DictConfig

from typing import Callable

def enter_context(key: str) -> Callable:
  def wrapper(func: Callable) -> Callable:
    @wraps(func)
    def inner_wrapper(conf: DictConfig):
        with MyContext(**conf[key]):  # create context using config values, and enter it
            return func(conf) # call original function

    return inner_wrapper

  return wrapper

But when I try to use it like this:

import hydra
from omegaconf import DictConfig


@hydra.main(config_path="../../conf", config_name="config")
@enter_context("env.hwm_store")
def main(conf: DictConfig) -> None:
    ...

I get an error:

[2022-07-19 09:41:30,670] {ssh_hook.py:476} WARNING - Traceback (most recent call last):
  File "myapp.py", line 91, in <module>
    main()
  File "/usr/local/lib/python3.7/site-packages/hydra/main.py", line 53, in decorated_main
    config_name=config_name,
  File "/usr/local/lib/python3.7/site-packages/hydra/_internal/utils.py", line 368, in _run_hydra
    lambda: hydra.run(
  File "/usr/local/lib/python3.7/site-packages/hydra/_internal/utils.py", line 214, in run_and_report
    raise ex
  File "/usr/local/lib/python3.7/site-packages/hydra/_internal/utils.py", line 211, in run_and_report
    return func()
  File "/usr/local/lib/python3.7/site-packages/hydra/_internal/utils.py", line 371, in <lambda>
    overrides=args.overrides,
  File "/usr/local/lib/python3.7/site-packages/hydra/_internal/hydra.py", line 91, in run
    run_mode=RunMode.RUN,
  File "/usr/local/lib/python3.7/site-packages/hydra/_internal/hydra.py", line 568, in compose_config
    from_shell=from_shell,
  File "/usr/local/lib/python3.7/site-packages/hydra/_internal/config_loader_impl.py", line 150, in load_configuration
    from_shell=from_shell,
  File "/usr/local/lib/python3.7/site-packages/hydra/_internal/config_loader_impl.py", line 227, in _load_configuration_impl
    self.ensure_main_config_source_available()
  File "/usr/local/lib/python3.7/site-packages/hydra/_internal/config_loader_impl.py", line 135, in ensure_main_config_source_available
    config_name=None, msg=msg, with_search_path=False
  File "/usr/local/lib/python3.7/site-packages/hydra/_internal/config_loader_impl.py", line 109, in _missing_config_error
    missing_cfg_file=config_name, message=add_search_path()
hydra.errors.MissingConfigException: Primary config directory not found.
Check that the config directory '/usr/local/lib/python3.7/conf' exists and readable

This is because hydra tries to detect config path based on file location there function is stored, and instead of main it is searching for enter_context function, which is stored in a file in '/usr/local/lib/python3.7/conf' directory:

hydra/hydra/_internal/utils.py

Line 53 in dd9e83a

calling_file = inspect.getfile(task_function)

To avoid this, I've added a block with upwrapping all the decorators up to the core function, using __wrapped__ attribute added by @functools.wraps:
https://github.com/python/cpython/blob/3f2dd0a7c0b1a5112f2164dce78fcfaa0c4b39c7/Lib/functools.py#L61

Have you read the Contributing Guidelines on pull requests?

Yes

Test Plan

Added app_with_config_decorated and related test

Related Issues and PRs

None

[facebook-github-bot]

Hi @dolfinus!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@fb.com. Thanks!

[facebook-github-bot]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[Jasha10]

Thanks @dolfinus, this PR looks great!
@pixelb, do you think there's reason to fence this off behind version_base for the sake of backward compatibility?

[pixelb]

I think this should be fine unguarded

[Jasha10]

Hmm, I may have been overly eager in giving approval above -- ideally we should add some documentation before merging this.

[dolfinus]

Do you mean adding example of using @hydra.main over a decorated function?

[Jasha10]

Do you mean adding example of using @hydra.main over a decorated function?

I guess so. Or maybe just a sentence (e.g. in the docstring of the hydra.main function) mentioning special treatment of the __wrapped__ attribute that results from using functools.wraps?

I'm worried that if we don't document this somewhere, future users will be ignorant of this special treatment of __wrapped__. Some users may even implement other workarounds/PRs when they could have just used functools.wraps.

[dolfinus]

@Jasha10 done

[Jasha10]

Thanks so much @dolfinus. This looks great.
I've closed the PR accidentally while trying to add a news fragment file (i.e. a file in the news/ folder), and now I'm having trouble reopening the PR. I'll open a new pull request to land this diff.
Edit: I've opened PR #2308.

[dolfinus]

Any chance this change will be released as 1.2.1?

[Jasha10]

@dolfinus I've just uploaded Hydra v1.3.0.dev0 to pypi, which you can install via pip install hydra-core==1.3.0.dev0.