Добавление повтора в будущую последовательность для параллельного запуска записных книжек Databricks в Scala

Question

Я использую приведенный ниже код из самого Databricks, чтобы узнать, как запустить его ноутбук параллельно в Scala, https://docs.databricks.com/notebooks/notebook-workflows.html#run -multiple-notebooks-concurrently . Я пытаюсь добавить функцию повтора, при которой, если одна из записных книжек в последовательности выйдет из строя, она повторит попытку этой записной книжки на основе значения повтора, которое я ему передал.

Вот код параллельной записной книжки из Databricks:

//parallel notebook code

import scala.concurrent.{Future, Await}
import scala.concurrent.duration._
import scala.util.control.NonFatal

case class NotebookData(path: String, timeout: Int, parameters: Map[String, String] = Map.empty[String, String])

def parallelNotebooks(notebooks: Seq[NotebookData]): Future[Seq[String]] = {
  import scala.concurrent.{Future, blocking, Await}
  import java.util.concurrent.Executors
  import scala.concurrent.ExecutionContext
  import com.databricks.WorkflowException

  val numNotebooksInParallel = 5
  // If you create too many notebooks in parallel the driver may crash when you submit all of the jobs at once. 
  // This code limits the number of parallel notebooks.
  implicit val ec = ExecutionContext.fromExecutor(Executors.newFixedThreadPool(numNotebooksInParallel))
  val ctx = dbutils.notebook.getContext()
  
  Future.sequence(
    notebooks.map { notebook => 
      Future {
        dbutils.notebook.setContext(ctx)
        if (notebook.parameters.nonEmpty)
          dbutils.notebook.run(notebook.path, notebook.timeout, notebook.parameters)
        else
          dbutils.notebook.run(notebook.path, notebook.timeout)
      }
      .recover {
        case NonFatal(e) => s"ERROR: ${e.getMessage}"
      }
    }
  )
}

Это пример того, как я вызываю приведенный выше код для запуска нескольких примеров записных книжек:

import scala.concurrent.Await
import scala.concurrent.duration._
import scala.language.postfixOps
val notebooks = Seq(
  NotebookData("Notebook1", 0, Map("client"->client)),
  NotebookData("Notebook2", 0, Map("client"->client))
)
val res = parallelNotebooks(notebooks)
Await.result(res, 3000000 seconds) // this is a blocking call.
res.value

Answer 1

Вот одна попытка. Поскольку ваш код не компилируется, я вставил несколько фиктивных классов.

Кроме того, вы не полностью указали желаемое поведение, поэтому я сделал некоторые предположения. Для каждого подключения будет выполнено всего пять попыток. Если какой-либо из Futures по-прежнему не работает после пяти попыток, то все Future терпит неудачу. Оба эти поведения можно изменить, но, поскольку вы не указали, я не уверен, что именно вы хотите.

Если у вас есть вопросы или вы хотите, чтобы я внес изменения в программу, дайте мне знать в разделе комментариев.

object TestNotebookData extends App{
  //parallel notebook code

  import scala.concurrent.{Future, Await}
  import scala.concurrent.duration._
  import scala.util.control.NonFatal

  case class NotebookData(path: String, timeout: Int, parameters: Map[String, String] = Map.empty[String, String])

  case class Context()

  case class Notebook(){
    def getContext(): Context = Context()
    def setContext(ctx: Context): Unit = ()
    def run(path: String, timeout: Int, paramters: Map[String, String] = Map()): Seq[String] = Seq()
  }
  case class Dbutils(notebook: Notebook)

  val dbutils = Dbutils(Notebook())


  def parallelNotebooks(notebooks: Seq[NotebookData]): Future[Seq[Seq[String]]] = {
    import scala.concurrent.{Future, blocking, Await}
    import java.util.concurrent.Executors
    import scala.concurrent.ExecutionContext

    // This code limits the number of parallel notebooks.
    implicit val ec = ExecutionContext.fromExecutor(Executors.newFixedThreadPool(numNotebooksInParallel))
    val ctx = dbutils.notebook.getContext()

    val isRetryable = true
    val retries = 5

    def runNotebook(notebook: NotebookData): Future[Seq[String]] = {
      def retryWrapper(retry: Boolean, current: Int, max: Int): Future[Seq[String]] = {
        val fut = Future {runNotebookInner}
        if (retry && current < max) fut.recoverWith{ _ => retryWrapper(retry, current + 1, max)}
        else fut
      }

      def runNotebookInner() = {
        dbutils.notebook.setContext(ctx)
        if (notebook.parameters.nonEmpty)
          dbutils.notebook.run(notebook.path, notebook.timeout, notebook.parameters)
        else
          dbutils.notebook.run(notebook.path, notebook.timeout)
      }

      retryWrapper(isRetryable, 0, retries)
    }


    Future.sequence(
      notebooks.map { notebook =>
        runNotebook(notebook)
      }
    )
  }

  val notebooks = Seq(
    NotebookData("Notebook1", 0, Map("client"->"client")),
    NotebookData("Notebook2", 0, Map("client"->"client"))
  )
  val res = parallelNotebooks(notebooks)
  Await.result(res, 3000000 seconds) // this is a blocking call.
  res.value
}

Answer 2

Я обнаружил, что это работает:

import scala.util.{Try, Success, Failure}

def tryNotebookRun (path: String, timeout: Int, parameters: Map[String, String] = Map.empty[String, String]): Try[Any] = {
  Try(
    if (parameters.nonEmpty){
      dbutils.notebook.run(path, timeout, parameters)
    }
    else{
      dbutils.notebook.run(path, timeout)
    }
  )
}

//parallel notebook code

import scala.concurrent.{Future, Await}
import scala.concurrent.duration._
import scala.util.control.NonFatal


def runWithRetry(path: String, timeout: Int, parameters: Map[String, String] = Map.empty[String, String], maxRetries: Int = 2) = {
  var numRetries = 0
  while (numRetries < maxRetries){
    
    tryNotebookRun(path, timeout, parameters) match {
      case Success(_) => numRetries = maxRetries
      case Failure(_) => numRetries = numRetries + 1      
    }    
  }
}

case class NotebookData(path: String, timeout: Int, parameters: Map[String, String] = Map.empty[String, String])

def parallelNotebooks(notebooks: Seq[NotebookData]): Future[Seq[Any]] = {
  import scala.concurrent.{Future, blocking, Await}
  import java.util.concurrent.Executors
  import scala.concurrent.ExecutionContext
  import com.databricks.WorkflowException

  val numNotebooksInParallel = 5
  // If you create too many notebooks in parallel the driver may crash when you submit all of the jobs at once. 
  // This code limits the number of parallel notebooks.
  implicit val ec = ExecutionContext.fromExecutor(Executors.newFixedThreadPool(numNotebooksInParallel))
  val ctx = dbutils.notebook.getContext()
  
  Future.sequence(
    notebooks.map { notebook => 
      Future {
        dbutils.notebook.setContext(ctx)
        runWithRetry(notebook.path, notebook.timeout, notebook.parameters)
      }
      .recover {
        case NonFatal(e) => s"ERROR: ${e.getMessage}"
      }
    }
  )
}