When we chatted about this I remember you saying that this code is to handle the case where the stage runs locally at the driver...but from glancing at the DAGScheduler code, it looks like onStageSubmitted() never gets called for the locally-run tasks, and they never show up at the UI.

When do you need this code / when will the jobIdToStageIds mapping not already be set up correctly by OnJobStart?