TestAutograd.test_thread_shutdown may fail sometimes

[Flamefire]

🐛 Describe the bug

When running this test e.g. via python test_autograd.py TestAutograd.test_thread_shutdown it may fail with the following sequence of events (i.e. this is a race condition):

• The test expects "PYTORCH_API_USAGE torch.autograd.thread_shutdown" at

  pytorch/test/test_autograd.py

  Line 4556 in d561aa9

  self.assertRegex(s, "PYTORCH_API_USAGE torch.autograd.thread_shutdown")

• This should be done by:

  pytorch/torch/csrc/autograd/engine.cpp

  Line 445 in d561aa9

  C10_LOG_API_USAGE_ONCE("torch.autograd.thread_shutdown");

• That in turn is only done via a "ShutdownTask" which (on a successful test) will be added via Engine::~Engine()

  pytorch/torch/csrc/autograd/engine.cpp

  Line 259 in d561aa9

  stop();

  calling stop()

  pytorch/torch/csrc/autograd/engine.cpp

  Line 283 in d561aa9

  queue->pushShutdownTask();

However that ShutdownTask is only pushed when all queues are empty, see

pytorch/torch/csrc/autograd/engine.cpp

Line 279 in d561aa9

noBackward = noBackward && queue->empty();

As the test code is very simple it is possible that the the "run_backward" task pushed at the end was not yet picked up by the/a thread and hence is still in the/a queue when the Engine destructor is run during program shutdown. In that case the ShutdownTask will never be pushed and the test fails with

FAIL: test_thread_shutdown (__main__.TestAutograd)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "test_autograd.py", line 4360, in test_thread_shutdown
    self.assertRegex(s, "PYTORCH_API_USAGE torch.autograd.thread_shutdown")
AssertionError: Regex didn't match: 'PYTORCH_API_USAGE torch.autograd.thread_shutdown' not found in 'PYTORCH_API_USAGE torch.python.import\nPYTORCH_API_USAGE c10d.python.import\nPYTORCH_API_USAGE tensor.create\n'

Versions

Noticed with PyTorch 1.9.0 on an Intel Cascade Lake system with V100 GPUs: easybuilders/easybuild-easyconfigs#16233

But may still apply as the relevant code is unchanged. See the referenced parts from current master above.

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @lezcano @Varal7

[ezyang]

This is a nice diagnosis. Do you have any sense for how you would fix it?

[albanD]

I am still a bit confused.
Since the backward() call finishes successfully before we destroy the process. The queue should have been emptied already.
Or does this happen when there is an error during the backward in this test and you see this regex mismatch instead of the actual error that happened in the worker process?

[Flamefire]

Since the backward() call finishes successfully before we destroy the process. The queue should have been emptied already.

I'm not 100% sure what MyFunction.apply(v).backward() is actually doing. A bit of digging and printing of stacktraces makes me think it basically only pushes a "run_backward" task onto the queue and returns. I.e. it is an asynchronous call.
Is that correct? Who may be able to confirm this?

@ezyang So if my reasoning above is correct than the answer to how to fix it would be to make sure the task is actually finished, i.e. add some sort of barrier.

But now after digging deeper Engine::execute already does wait for the result so I'm not sure why it still does fail. I can't currently debug this further because on the machines I have access to it doesn't fail.

Or does this happen when there is an error during the backward in this test and you see this regex mismatch instead of the actual error that happened in the worker process?

I doubt so. The stderr as reported by the test is PYTORCH_API_USAGE torch.python.import\nPYTORCH_API_USAGE c10d.python.import\nPYTORCH_API_USAGE tensor.create\n so it simply is lacking the shutdown log.

[albanD]

I.e. it is an asynchronous call. Is that correct?

No it is not. This is a blocking call that will wait for all gradients to be computed and all .grad field to be populated before returning.

[Flamefire]

I further checked that. I can verify that the shutdown task is pushed during the destructor call of PythonEngine. It then goes on into the destructor of Engine and terminates. So not sure why it fails

[Flamefire]

@albanD I used another approach, due to not having direct access to the machine where this fails I let them add a patch with some prints: debug_pt19.patch.txt

test_thread_shutdown (__main__.TestAutograd) ... Output: PYTORCH_API_USAGE torch.python.import
PYTORCH_API_USAGE c10d.python.import
PYTORCH_API_USAGE tensor.create
Impl cpu reports 1 devices
Impl cuda reports 1 devices
Impl meta reports 1 devices
Started with device -1
Got other task-1
Got other task-1
Ended with device -1
Starting wait in execute
End wait in execute
Started with device -1
Got other task-1
Got other task-1
Ended with device -1
Starting wait in execute
End wait in execute
stop from frame #0: torch::autograd::Engine::stop() + 0x45 (0x154700a67125 in /tmp/eb-cirgbre5/tmpvgi_ms38/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so)
frame #1: torch::autograd::python::PythonEngine::~PythonEngine() + 0x17 (0x154706ff3197 in /tmp/eb-cirgbre5/tmpvgi_ms38/lib/python3.8/site-packages/torch/lib/libtorch_python.so)
frame #2: <unknown function> + 0x39b0c (0x15473252bb0c in /lib64/libc.so.6)
frame #3: on_exit + 0 (0x15473252bc40 in /lib64/libc.so.6)
frame #4: __libc_start_main + 0xfa (0x15473251549a in /lib64/libc.so.6)
frame #5: _start + 0x2e (0x4006de in /user/gent/400/vsc40023/eb_scratch/RHEL8/cascadelake-volta-ib/software/Python/3.8.6-GCCcore-10.2.0/bin/python)

noBackward=1 wait_duration=10
Pushing shutdown
Waiting for 1 threads
Stopping with timeout
End waiting for all threads: 1

So from what I can tell with this what happens is:

• Initial state: No threads running
• Engine::execute gets called through some layers from Python
• Ultimately this calls execute_with_graph_task which at the start calls initialize_device_threads_pool starting 1 device thread
• as the current worker_device == NO_DEVICE execute_with_graph_task calls into thread_main with the current device set to CPU_DEVICE/-1
• Seemingly the graph task is scheduled to the CPU queue so the thread_main finds 2 tasks, executes them and returns
• This happens twice (I assume for forward and backward pass of the test)
• The program is now finished with no work pushed to the device/CUDA queue

So execution is finished while the thread started at step 3 is likely somewhere on this lines, maybe at::init_num_threads() takes a bit longer due to some locking.

But anyway as the thread is detached it might be anywhere and the timeout of 10s seems to be to short for the thread to fully initialize, call thread_main, get the shutdown task and terminate. A classic race condition as presumed earlier. So the thread is indeed "leaked".

So only "fix" I can imagine is to make sure execution happens NOT on the CPU, so the apply(v).backward() force the thread to init and handle the task hoping it gets to handle the shutdown in the 10s timeout.

[albanD]

10s to initialize the thread sounds quite a lot...

So only "fix" I can imagine is to make sure execution happens NOT on the CPU, so the apply(v).backward() force the thread to init and handle the task hoping it gets to handle the shutdown in the 10s timeout.

I agree with that. Could you send a patch doing that?
Leaving a small comment there that this has to run on the GPU (if there is one) to ensure the worker thread is initialized (if it exists).

[Flamefire]

I agree that 10s is a lot, so not sure where the thread is hanging, maybe it doesn't get to set the OMP threads while the other (CPU) thread is doing work and is suspended and the OS doesn't let it wake up earlier... So far I've only seen it on 1 system (to which I don't have direct access to making debugging/fixing this hard)

I'm afraid I can't send a patch as I'm not directly sure how to do that.

I hope my analysis helps to get this issue fixed though.