How to stop a notebook streaming job gracefully?

问题

I have a streaming application which is running into a Databricks notebook job (https://docs.databricks.com/jobs.html). I would like to be able to stop the streaming job gracefully using the stop() method of the StreamingQuery class which is returned by the stream.start() method. That of course requires to either have access to the mentioned streaming instance or to access the context of the running job itself. In this second case the code could look as next:

spark.sqlContext.streams.get("some_streaming_uuid").stop()

The above code should be executed from a different notebook job lets call it stop_streaming_job although I wasn't able to find a way to access the job context and to execute the above scala code. Is there any way to achieve that with the databricks notebooks?

回答1:

One way to solve this is using the databricks file system (dbfs) or your local file system. The idea is to extend the functionality of the Spark StreamingQuery class by implementing a new function called awaitExternalTermination. The solution creates a new file in the given DBFS directory which acts as a flag responsible for the lifetime of the streaming job. The job will continue running as long as the file exists in the given directory. Next follows the implementation of the file watcher which is an extension method of the StreamingQuery class and uses Scala futures:

object extensions {
  import fs._
  object FileSystemType extends Enumeration {
    val DBFS, LocalFileSystem = Value
  }

  implicit class FileSystemStopStreamingQuery(val self :StreamingQuery) extends AnyVal {
    /**
     * Extension method for StreamingQuery, it waits for an external call to delete the streaming file. When that happens it will call the stop method
     * of the current StreamingQuery instance.
     *
     * @param streamStopDir dir to be watched
     * @param jobName the job unique identifier/the file name
     * @param fsType DFFS or LocalFileSystem
     */
    def awaitExternalTermination(streamStopDir :String, jobName :String, fsType : FileSystemType.Value): Unit ={

      if(streamStopDir == null || streamStopDir.isEmpty)
        throw new IllegalArgumentException("streamStopDir can't be null or empty.")

      if(jobName == null || jobName.isEmpty)
        throw new IllegalArgumentException("jobName can't be null or empty.")

      val fsWrapper :FileSystemWrapper = fsType match {
        case FileSystemType.DBFS => new DbfsWrapper(streamStopDir, jobName)
        case FileSystemType.LocalFileSystem => new LocalFileSystemWrapper(streamStopDir, jobName)
        case _ => throw new IllegalArgumentException("Invalid file system provided.")
      }

      val stopWatchFuture: Future[Boolean] = Future {

        if(!fsWrapper.targetFileExists)
            fsWrapper.createTargetFile(self.id.toString)

        while (self.isActive && fsWrapper.targetFileExists){
          val random: ThreadLocalRandom = ThreadLocalRandom.current()
          val r = random.nextLong(10, 100 + 1) // returns value between 10 and 100
          Thread.sleep(r)
        }

        if(!fsWrapper.targetFileExists){
          self.stop()
          true
        }
        else
          false
      }

      var output = "success"
      stopWatchFuture onComplete {
        case Success(result : Boolean) => if (!result) {
          output = s"failure: file not found."
        }
        case Failure(t) => output = s"failure: ${t.getMessage}."
      }

      self.awaitTermination()
    }
  }
}

And the implementation of the DBFS wrapper class:

import com.databricks.dbutils_v1.DBUtilsHolder.dbutils

class DbfsWrapper(val stopDir: String, val targetFile: String) extends FileSystemWrapper {
  override def targetFileExists(): Boolean = {
    try {
      dbutils.fs.ls(targetPath).size > 0
    }
    catch {
      case _: java.io.FileNotFoundException => false
    }
  }

  override def createTargetFile(content: String): Unit = {
    dbutils.fs.put(targetPath, content)
  }
}

To stop the streaming job just remove the mentioned file with %fs rm -r your_path when using DBFS or just rm -r your_path for the local FS.

The complete code can be found here.

来源：https://stackoverflow.com/questions/59310665/how-to-stop-a-notebook-streaming-job-gracefully