Not sure what your "this" refers to.  There shouldn't be be stages left in the pendingTasks, waiting, running or failed queues at this point, so if there are, then I should at least log the unexpected occurrence.  If the unexpected occurs, do you want to fail, adopt this "clean up for a mistake someplace else" approach, or try to do something like TimeStampedHashMap to periodically clean up the leaks?

Within jobIdToStageIdsRemove, I'm also now looking at cleaning up stageIdToJobIds, stageToInfos and shuffleToMapStage so that they don't continue to grow.  The first two are easy, but I've got some work to do on the third.