It's fine if you want to leave this class but it might be worth taking it out since it could be interpreted as a public API. I think in general we'll want a standalone thing that runs in YARN deployments but this is just a small example. Also we could maybe move this to the examples directory.