[SPARK-3347] [yarn] Fix yarn-alpha compilation.

Missing import. Oops.