Adding retry to future sequence for running Databricks notebooks in parallel in Scala

Here is one attempt. Since your code does not compile, I inserted a few dummy classes.

Also, you did not fully specify the desired behavior, so I made some assumptions. Only five retries will be made for each connection. If any of the Futures are still failing after five retries, then the entire Future is failed. Both of these behaviors can be changed, but since you did not specify, I am not sure what it is you want.

If you have questions or would like me to make an alteration to the program, let me know in the comments section.

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
}

I found this to work:

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}"
      }
    }
  )
}