Track when a particular function (like _run_ddp_forward) is in the stack

[anijain2305]

In some cases, it is useful to know if a particular frame is running under a specific context. For example, we would like to know _run_ddp_forward is ancestor of the current frame we are analyzing. This enables us to skip/specialize on this additional information during Dynamo bytecode processing.

Conceptually, it could look something like this

Before Dynamo

def fn():
    .....
    .....

After Dynamo

def wrapper(fn):
    def inner(*args, **kwargs):
        torchdyanmo.global_ctx.add("ddp_on")
        out = torchdynamo.optimize(my_compiler)(fn)(*args, **kwargs)
        torchdyanmo.global_ctx.remove("ddp_on")
        return out
    return inner

cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh

[wconstab]

I tried building a ctx mgr inside dynamo that i can call from ddp but that leaves me with circular imports...

i guess i could make a global ddp context in ddp land, and just query that from the dynamo side..

but if there is a way to do this infrastructurally within dynamo that'd be awesome

I was able to do something equivalent but slow:

        def _inside_ddp_callstack(frame):
            up = frame.f_back
            while up:
                if up.f_code.co_name == "_run_ddp_forward":
                    # print("Found DDP in parent trace")
                    return True
                up = up.f_back
            return False

I can call this inside 'catch_errors' and potentially hijack the compiler there, but I haven't figured it out yet.

[wconstab]

for context, this is the DDP thing I am working on that surfaced these questions:
pytorch/torchdynamo#628