Big debugging update

[samlurye]

Summary:
This diff contains several updates to the monarch actor mesh debugging experience.

• Improved debug input parsing, with the new cast command supporting more sophisticated rank selection grammar:
  • cast ranks(3) pdb_command: send pdb_command to rank 3.
  • cast ranks(1,3,5) pdb_command: send pdb_command to ranks 1, 3 and 5.
  • cast ranks(1:10:2) pdb_command: send pdb_command to the ranks in range(start=1, stop=10, step=2).
  • cast ranks(pp=2, dp=(1,3), tp=2:8) pdb_command: send pdb_command to ranks with pp dim 2, dp dim 1 or 3, and tp dim in range(2,8).
• The debug client is now automatically registered with an actor mesh when that actor mesh is spawned. This means calling init_debugging(actor_mesh) is no longer necessary.
• Debugging now works with MAST jobs, by enforcing that breakpoints aren't set in __main__, and that the file containing the breakpoint exists on the remote host.
  • The first requirement is due to how cloudpickle works -- if an actor endpoint is defined inside __main__, cloudpickle will serialize it by value instead of by reference. When the code then runs on the remote host, it thinks the location of the code is the user's local __main__ file, which confuses pdb, because the file doesn't exist at the same path (or may not exist at all) on the remote host.
  • The second requirement is due to important parts of pdb's implementation relying on the ability to search for the file being debugged on the remote host's file system.
• A debugging session for a specific rank is now forcefully exited once the endpoint finishes execution. This contains the debugging experience within user-authored code. It is also necessary for preventing hangs, because if pdb is allowed to continue indefinitely, then control flow will eventually bubble back up to the main asyncio event loop on the worker, at which point everything breaks.
• Hitting a breakpoint now automatically enables post-mortem debugging, so any rank that encounters an exception after hitting a breakpoint will automatically stop at the exception. Attaching the debugger to that rank should then provide an experience like pdb.post_mortem().

Next steps/gaps I'm aware of (reviewers please read):

• Indexing debug sessions by rank isn't sustainable, because two actor meshes may simultaneously hit breakpoints on the same rank and cause a collision inside the debug client.
• Entering the debug client should happen automatically, rather than requiring the user to do await debug_client().enter.call_one().
• Casting pdb commands should ideally leverage MeshTrait rather than reimplementing the selection logic.
• If a mesh was reshaped/renamed so that its dimension names aren't hosts and gpus anymore, the debugger should reflect the new shape/names.
• The user should be able to enable post-mortem debugging without having to hit a separate breakpoint first.

Differential Revision: D77568423

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D77568423

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D77568423

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D77568423