try api call logging with decorator

[pmeier]

We have a lot of these sprinkled in throughout our code base:

vision/torchvision/transforms/functional.py

Lines 419 to 420 in 6fcf0a2

	if not torch.jit.is_scripting() and not torch.jit.is_tracing():
	_log_api_usage_once(resize)

This is almost a text book case for using a decorator. I was always told that decorators are not supported by torch.jit.script and thus we have to go the "long way". But it seems decorators are supported

import functools

import torch


def logger(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        print(f"Logging call to {fn.__name__}")
        return fn(*args, **kwargs)

    return wrapper


@logger
def foo(input: torch.Tensor) -> torch.Tensor:
    print("inside foo()")
    return input.abs()


input = torch.tensor(-1)
print("eager")
print(foo(input))
print("scripted")
print(torch.jit.script(foo)(input))

eager
Logging call to foo
inside foo()
tensor(1)
scripted
inside foo()
tensor(1)

Since we get no print from the logger during scripted mode, maybe it picks up on the functools.wraps decorator and only executes the underlying function.

[facebook-github-bot]

💊 CI failures summary and remediations

As of commit 1addc91 (more details on the Dr. CI page):

---

💚 💚 Looks good so far! There are no failures yet. 💚 💚

---

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

[pmeier]

@kazhang in #5052 you stated that one cannot

• use decorator
  • For example: @log_api_usage in use decorator to log API usage #4976
  • doesn’t work with TorchScript since decorator needs to use kwargs, which is not supported in TorchScript;

to log the API calls. Either there was support added or I'm missing something here, since it seems to work out fine. I've added the _log_api_usage_once call to a decorator and applied it to all our functional transforms. JIT scriptability is tested by our test suite, but there no errors.

[kazhang]

Hi @pmeier , thanks for the PR! looks like decorator (more specifically kwargs) is supported now. Initially I tried it out in #4976 but got this error from JIT. I'm glad that it works now!

[kazhang]

shall we move this decorator to utils? or turn _log_api_usage_once into a decorator if it works for class as well?

[datumbox]

@pmeier I believe adding decorators previously broke JIT-scriptability. At least that's what I recall from Kai's earlier experiments (see #4976 and #5052). But then the CI tests are not currently failing which is very strange.

@kazhang any thoughts on whether the approach followed by Philip actually works OR if we have a gap on our testing strategy?

Edit: I dug a bit to see what was done previously and it seems that the linked PR has modified and doesn't show the correct experiment. I believe this is the commit that was failing JIT. The approach is very similar to what we do here, so I'm not sure why it works now. Specifically the fn(*args, **kwargs) part was one of the key limitations on JIT. Does it means that now JIT supports it? I've reached out to people from JIT to see if they can provide more clarity.

[kazhang]

@datumbox @pmeier
Now I remember why decorator didn't work for me in #4976
See my comment:

Unfortunately, torchscript doesn't like decorator. When using the decorator on boxes functions, JIT looks for helper functions definition(e.g. _upcast) in utils.py

In this PR, we define the decorator in functional.py instead of utils.py, so JIT is still able to find helper functions used by the function to be scripted. So basically we'd need to define this decorator everywhere.

[kazhang]

In #5436 I tried to move the decorator to utils.py, and reproduced the error: https://app.circleci.com/pipelines/github/pytorch/vision/14851/workflows/ad9beb44-ba9e-472e-a71b-8775e909f8b5/jobs/1195207

[pmeier]

Looking through the issue tracker, I found the long standing issue pytorch/pytorch#38964 that should have been fixed in pytorch/pytorch#63248. Maybe @qihqi can have a look if the fix was incomplete or our setup is different?

[datumbox]

I'm going to mark as blocked to avoid any accidental merges because it's possible that this PR breaks JIT-scriptability without the tests catching it. If that's not the case, then that's great and we can unblock. We just need to do a bit more due diligence to ensure this is not due to a gap in our testing strategy.

[datumbox]

@pmeier I believe adding decorators previously broke JIT-scriptability. At least that's what I recall from Kai's earlier experiments (see #4976 and #5052). But then the CI tests are not currently failing which is very strange.

@kazhang any thoughts on whether the approach followed by Philip actually works OR if we have a gap on our testing strategy?

Edit: I dug a bit to see what was done previously and it seems that the linked PR has modified and doesn't show the correct experiment. I believe this is the commit that was failing JIT. The approach is very similar to what we do here, so I'm not sure why it works now. Specifically the fn(*args, **kwargs) part was one of the key limitations on JIT. Does it means that now JIT supports it? I've reached out to people from JIT to see if they can provide more clarity.

[pmeier]

This is the official way to declare decorators that preserve the signature. Let's see if the test suite is still happy.

[datumbox]

I spoke with @gmagogsfm and he suspects that TorchScript completely ignored the logger decorator. He said that this is likely an unintended use case and suggested it is not safe to depend on it.

[vadimkantorov]

related: #4957, pytorch/pytorch#70051