Your Pr desc mentions less clear error messages in dynamo, can you explain how that relates to the changes? It seems like this change should just cause more cases to miss cache and there would be no associated error.

Relatedly, if we silently miss cache for unhashable type errors, we'll not get the signal we need to fix those cases. Do you have a plan for that?

Originally we had landed #156868 to purposely skip the cache during compile, since there was the possibility of symints in the op schema, and existing has_symints would miss them if the schema was nested (e.g. an argument is a list of tensors). The goal was to avoid the additional overhead of doing a nested check and instead just disabling on compile.

However, that led to #159590, where dynamo is disabled for autograd with dynamic shapes so _are_we_tracing was not sufficient, and #159601 where the statefulness of _MaskPartial placement led to and issue with the propagate_op_sharding cache and the _gen_transform_infos not being aligned, so the mask materialization would not happen in the correct object.

So the additional cache misses are intentional. If there are errors from the caches, we should also expect to see them in the non-cached version. Else there is an error in functools, or the function we are using cacheing for mutates arguments, in which case we shouldn't be using cacheing for it. I do still feel a bit wary about try/except clauses in general, so we could additionaly filter the except on the exception message to only include "unshable type" in the error message, but that assumes the messsage will never change in functools.

Discussed further with @bdhirsh and instead of the try/except, I changed it to go back to the compiled autograd flag, with an addition of the the PythonDispatcher being invoked in compile_inner, and also checked in _are_we_tracing.

yep i would start out making this a bit lower blast radius - probably by moving it just around the code in dynamo that performs FakeTensor propagation

Moved it to https://github.com/pytorch/pytorch/pull/160798/files#diff-9f0663783bcd93e948e0491ef61b48123bdc9977bcc632fd707da578df13bfa1R2147 and looks like it's resolved the issues. Does this look appropriate?

So we'd need to cover every fake tensor tracing callsite with python dispatcher? If so, those are spread across a bit. You probably want to move the enable inside of wrap_to_fake_tensor_and_record as it is used a lot, and another one is wrap_fake_exception.

Modified it to wrap the entire of wrap_to_fake_tensor_and_record and for the fn call in wrap_fake_exception.

nit: we could probably do even better, but right now most of our compile + dtensor related tests live in test_dtensor_compile.py. I would probably vote to put the test there instead of in this file, so we can keep most of the eager-only DTensor test files as being actually eager-only.

also, just confirming - does this test fail if you remove the enable_python_dispatcher bits?

Good point - I initially had correctness issues with the fake process group, but saw that TestDTensorCompileE2E had comms setup for the tests, so I moved it in there now.

This test is actually for the issue brought up in #159601, so not impacted by the enable_python_dispatcher changes. Add the _are_we_tracing() check to redistribute.py is what makes this pass.

i think it was pretty non-obvious that we needed this for the dynamo case that you found. do you mind adding some comments so someone reading this code later can understand why we need it?

For this kind of situation I like add two comments that reference one another in the two different regions of code, so someone reading the code has some breadcrumbs:

e.g.

....in the region of dynamo where you enable_python_dispatcher
# Note [enable_python_dispatcher in dynamo]
# Dynamo disables itself when it runs fake tensor prop, which means that tensor subclasses
# have no way to know (purely based off of global state) if they are currently being run under compile or not.
# we use enable_python_dispatcher mainly to tweak the DispatchKeyState so that subclass authors
# can check it to know if they are running in an eager context or not

... in the _are_we_tracing function
# See Note [enable_python_dispatcher in dynamo]

Added the comments, appreciate it!

are there any functional changes to this file other than shifting everything here under the enable_python_dispatcher block? (if so, it would be easier to review if you can keep the rest of the code unchanged so the PR diff only shows up as an indent)

Yeah, it looks like the linter adds some of these changes. I'll check if there's smaller sections to enable the dispatcher as well as clean this up.

I figured it'd actually much cleaner to use a decorator, so I added wrap_python_dispatch. Not sure if there's an existing naming precedent for such a decorator though, so let me know if it needs to be changed up anyway.

One comment I have is that:

(1) this enables enable_python_dispatcher for FakeTensorProp in dynamo

(2) technically, the "bad" case we saw only happened in the dynamo code that converted a graph input tensor into a real tensor. (the call to fake_mode.from_tensor(...))

Converting (real tensor) graph inputs into FakeTensors happens once for every input, while FakeTensorProp happens for every single tensor operation in the graph, so (1) is much more common than (2). We know that historically, FakeTensorProp takes a meaningful amount of compile time (we introduced caching inside of FakeTensorProp a while ago to help). So there is some risk of enable_python_dispatcher making FakeTensorProp in dynamo (and thus compile times) a bit slower.

@bdhirsh I see, so wrap_fake_exception is used in FakeTensorProp? In that case is wrap_to_fake_tensor_and_record used for FakeTensorProp? I think I'm lacking some context to figure out where the right place for the dispatcher is.

discussed offline - we found a way to only turn on the python dispatcher for calls to fake_mode.from_tensor, and not during all of fake tensor prop