diff --git a/yarn/common/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala b/yarn/common/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
index a3c43b43848d2bcf9b9997449be977265d94e0d5..e6fe0265d8811bdd5d2d5d0d410525b4cef3b02b 100644
--- a/yarn/common/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
+++ b/yarn/common/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
@@ -92,6 +92,7 @@ private[spark] class ApplicationMaster(args: ApplicationMasterArguments,
 
       logInfo("ApplicationAttemptId: " + appAttemptId)
 
+      val fs = FileSystem.get(yarnConf)
       val cleanupHook = new Runnable {
         override def run() {
           // If the SparkContext is still registered, shut it down as a best case effort in case
@@ -115,7 +116,7 @@ private[spark] class ApplicationMaster(args: ApplicationMasterArguments,
             // we only want to unregister if we don't want the RM to retry
             if (finalStatus == FinalApplicationStatus.SUCCEEDED || isLastAttempt) {
               unregister(finalStatus, finalMsg)
-              cleanupStagingDir()
+              cleanupStagingDir(fs)
             }
           }
         }
@@ -303,8 +304,7 @@ private[spark] class ApplicationMaster(args: ApplicationMasterArguments,
   /**
    * Clean up the staging directory.
    */
-  private def cleanupStagingDir() {
-    val fs = FileSystem.get(yarnConf)
+  private def cleanupStagingDir(fs: FileSystem) {
     var stagingDirPath: Path = null
     try {
       val preserveFiles = sparkConf.get("spark.yarn.preserve.staging.files", "false").toBoolean