Add run --spark

Author: alexarchambault

modules/build/src/main/scala/scala/build/internal/Runner.scala

@@ -99,7 +99,7 @@ object Runner {
     }
   }
 
-  private def envCommand(env: Map[String, String]): Seq[String] =
+  def envCommand(env: Map[String, String]): Seq[String] =
     env.toVector.sortBy(_._1).map {
       case (k, v) =>
         s"$k=$v"

modules/cli-options/src/main/scala/scala/cli/commands/RunOptions.scala

@@ -25,6 +25,11 @@ final case class RunOptions(
   @Recurse
     mainClass: MainClassOptions = MainClassOptions(),
   @Group("Run")
+  @Hidden
+  @HelpMessage("Run as a Spark job, using the spark-submit command")
+  @ExtraName("spark")
+    sparkSubmit: Option[Boolean] = None,
+  @Group("Run")
   @HelpMessage("Print the command that would have been run (one argument per line), rather than running it")
     command: Boolean = false,
   @Group("Run")

modules/cli/src/main/scala/scala/cli/commands/Package.scala

@@ -719,7 +719,7 @@ object Package extends ScalaCommand[PackageOptions] {
     * provided modules, and might have been bumped by other modules. This is strictly a subset of
     * the whole dependency graph.
     */
-  private def providedFiles(
+  def providedFiles(
     build: Build.Successful,
     provided: Seq[dependency.AnyModule],
     logger: Logger

modules/cli/src/main/scala/scala/cli/commands/Run.scala

@@ -10,18 +10,26 @@ import scala.build.internal.{Constants, Runner, ScalaJsLinkerConfig}
 import scala.build.options.{BuildOptions, JavaOpt, Platform}
 import scala.build.{Build, BuildThreads, Inputs, Logger, Positioned}
 import scala.cli.CurrentParams
+import scala.cli.commands.run.RunMode
 import scala.cli.commands.util.MainClassOptionsUtil._
 import scala.cli.commands.util.SharedOptionsUtil._
 import scala.cli.internal.ProcUtil
 import scala.util.Properties
 import scala.cli.config.{ConfigDb, Keys}
 import scala.cli.commands.util.CommonOps.SharedDirectoriesOptionsOps
+import scala.cli.commands.util.RunSpark
 
 object Run extends ScalaCommand[RunOptions] {
   override def group = "Main"
 
   override def sharedOptions(options: RunOptions): Option[SharedOptions] = Some(options.shared)
 
+  private def runMode(options: RunOptions): RunMode =
+    if (options.sparkSubmit.getOrElse(false))
+      RunMode.SparkSubmit
+    else
+      RunMode.Default
+
   private def scratchDirOpt(options: RunOptions): Option[os.Path] =
     options.scratchDir
       .filter(_.trim.nonEmpty)
@@ -48,7 +56,29 @@ object Run extends ScalaCommand[RunOptions] {
       javaOptions = baseOptions.javaOptions.copy(
         javaOpts =
           baseOptions.javaOptions.javaOpts ++
-            sharedJava.allJavaOpts.map(JavaOpt(_)).map(Positioned.commandLine)
+            sharedJava.allJavaOpts.map(JavaOpt(_)).map(Positioned.commandLine),
+        jvmIdOpt = baseOptions.javaOptions.jvmIdOpt.orElse {
+          runMode(options) match {
+            case RunMode.SparkSubmit => Some("8")
+            case RunMode.Default     => None
+          }
+        }
+      ),
+      internalDependencies = baseOptions.internalDependencies.copy(
+        addRunnerDependencyOpt = baseOptions.internalDependencies.addRunnerDependencyOpt.orElse {
+          runMode(options) match {
+            case RunMode.SparkSubmit => Some(false)
+            case RunMode.Default     => None
+          }
+        }
+      ),
+      internal = baseOptions.internal.copy(
+        keepResolution = baseOptions.internal.keepResolution || {
+          runMode(options) match {
+            case RunMode.SparkSubmit => true
+            case RunMode.Default     => false
+          }
+        }
       ),
       notForBloopOptions = baseOptions.notForBloopOptions.copy(
         runWithManifest = options.useManifest
@@ -76,6 +106,7 @@ object Run extends ScalaCommand[RunOptions] {
     def maybeRun(
       build: Build.Successful,
       allowTerminate: Boolean,
+      runMode: RunMode,
       showCommand: Boolean,
       scratchDirOpt: Option[os.Path]
     ): Either[BuildException, Option[(Process, CompletableFuture[_])]] = either {
@@ -95,15 +126,17 @@ object Run extends ScalaCommand[RunOptions] {
             allowExecve = allowTerminate,
             jvmRunner = build.artifacts.hasJvmRunner,
             potentialMainClasses,
+            runMode,
             showCommand,
             scratchDirOpt
           )
         }
 
         processOrCommand match {
-          case Right(process) =>
+          case Right((process, onExitOpt)) =>
             val onExitProcess = process.onExit().thenApply { p1 =>
               val retCode = p1.exitValue()
+              onExitOpt.foreach(_())
               if (retCode != 0)
                 if (allowTerminate)
                   sys.exit(retCode)
@@ -170,6 +203,7 @@ object Run extends ScalaCommand[RunOptions] {
             val maybeProcess = maybeRun(
               s,
               allowTerminate = false,
+              runMode = runMode(options),
               showCommand = options.command,
               scratchDirOpt = scratchDirOpt(options)
             )
@@ -206,6 +240,7 @@ object Run extends ScalaCommand[RunOptions] {
           val res = maybeRun(
             s,
             allowTerminate = true,
+            runMode = runMode(options),
             showCommand = options.command,
             scratchDirOpt = scratchDirOpt(options)
           )
@@ -226,9 +261,10 @@ object Run extends ScalaCommand[RunOptions] {
     allowExecve: Boolean,
     jvmRunner: Boolean,
     potentialMainClasses: Seq[String],
+    runMode: RunMode,
     showCommand: Boolean,
     scratchDirOpt: Option[os.Path]
-  ): Either[BuildException, Either[Seq[String], Process]] = either {
+  ): Either[BuildException, Either[Seq[String], (Process, Option[() => Unit])]] = either {
 
     val mainClassOpt = build.options.mainClass.filter(_.nonEmpty) // trim it too?
       .orElse {
@@ -250,6 +286,7 @@ object Run extends ScalaCommand[RunOptions] {
       finalArgs,
       logger,
       allowExecve,
+      runMode,
       showCommand,
       scratchDirOpt
     )
@@ -262,9 +299,10 @@ object Run extends ScalaCommand[RunOptions] {
     args: Seq[String],
     logger: Logger,
     allowExecve: Boolean,
+    runMode: RunMode,
     showCommand: Boolean,
     scratchDirOpt: Option[os.Path]
-  ): Either[BuildException, Either[Seq[String], Process]] = either {
+  ): Either[BuildException, Either[Seq[String], (Process, Option[() => Unit])]] = either {
 
     build.options.platform.value match {
       case Platform.JS =>
@@ -308,7 +346,7 @@ object Run extends ScalaCommand[RunOptions] {
                 esModule = esModule
               )
               process.onExit().thenApply(_ => if (os.exists(jsDest)) os.remove(jsDest))
-              Right(process)
+              Right((process, None))
             }
           }
         value(res)
@@ -327,33 +365,50 @@ object Run extends ScalaCommand[RunOptions] {
               logger,
               allowExecve = allowExecve
             )
-            Right(proc)
+            Right((proc, None))
           }
         }
       case Platform.JVM =>
-        if (showCommand) {
-          val command = Runner.jvmCommand(
-            build.options.javaHome().value.javaCommand,
-            build.options.javaOptions.javaOpts.toSeq.map(_.value.value),
-            build.fullClassPath,
-            mainClass,
-            args,
-            useManifest = build.options.notForBloopOptions.runWithManifest
-          )
-          Left(command)
-        }
-        else {
-          val proc = Runner.runJvm(
-            build.options.javaHome().value.javaCommand,
-            build.options.javaOptions.javaOpts.toSeq.map(_.value.value),
-            build.fullClassPath,
-            mainClass,
-            args,
-            logger,
-            allowExecve = allowExecve,
-            useManifest = build.options.notForBloopOptions.runWithManifest
-          )
-          Right(proc)
+        runMode match {
+          case RunMode.Default =>
+            if (showCommand) {
+              val command = Runner.jvmCommand(
+                build.options.javaHome().value.javaCommand,
+                build.options.javaOptions.javaOpts.toSeq.map(_.value.value),
+                build.fullClassPath,
+                mainClass,
+                args,
+                useManifest = build.options.notForBloopOptions.runWithManifest,
+                scratchDirOpt = scratchDirOpt
+              )
+              Left(command)
+            }
+            else {
+              val proc = Runner.runJvm(
+                build.options.javaHome().value.javaCommand,
+                build.options.javaOptions.javaOpts.toSeq.map(_.value.value),
+                build.fullClassPath,
+                mainClass,
+                args,
+                logger,
+                allowExecve = allowExecve,
+                useManifest = build.options.notForBloopOptions.runWithManifest,
+                scratchDirOpt = scratchDirOpt
+              )
+              Right((proc, None))
+            }
+          case RunMode.SparkSubmit =>
+            value {
+              RunSpark.run(
+                build,
+                mainClass,
+                args,
+                logger,
+                allowExecve,
+                showCommand,
+                scratchDirOpt
+              )
+            }
         }
     }
   }

modules/cli/src/main/scala/scala/cli/commands/run/RunMode.scala

@@ -0,0 +1,8 @@
+package scala.cli.commands.run
+
+sealed abstract class RunMode extends Product with Serializable
+
+object RunMode {
+  case object Default     extends RunMode
+  case object SparkSubmit extends RunMode
+}

modules/cli/src/main/scala/scala/cli/commands/util/RunSpark.scala

@@ -0,0 +1,76 @@
+package scala.cli.commands.util
+
+import scala.build.EitherCps.{either, value}
+import scala.build.{Build, Logger}
+import scala.build.errors.BuildException
+import scala.build.internal.Runner
+import scala.cli.commands.{Package => PackageCmd}
+import scala.cli.commands.packaging.Spark
+import scala.cli.commands.run.RunMode
+import scala.cli.packaging.Library
+import scala.util.Properties
+
+object RunSpark {
+
+  def run(
+    build: Build.Successful,
+    mainClass: String,
+    args: Seq[String],
+    logger: Logger,
+    allowExecve: Boolean,
+    showCommand: Boolean,
+    scratchDirOpt: Option[os.Path]
+  ): Either[BuildException, Either[Seq[String], (Process, Option[() => Unit])]] = either {
+
+    // FIXME Get Spark.sparkModules via provided settings?
+    val providedModules = Spark.sparkModules
+    val providedFiles =
+      value(PackageCmd.providedFiles(build, providedModules, logger)).toSet
+    val depCp        = build.dependencyClassPath.filterNot(providedFiles)
+    val javaHomeInfo = build.options.javaHome().value
+    val javaOpts     = build.options.javaOptions.javaOpts.toSeq.map(_.value.value)
+    val ext          = if (Properties.isWin) ".cmd" else ""
+    val submitCommand: String =
+      Option(System.getenv("SPARK_HOME"))
+        .map(os.Path(_, os.pwd))
+        .map(_ / "bin" / s"spark-submit$ext")
+        .filter(os.exists(_))
+        .map(_.toString)
+        .getOrElse(s"spark-submit$ext")
+    val jarsArgs =
+      if (depCp.isEmpty) Nil
+      else Seq("--jars", depCp.mkString(","))
+
+    scratchDirOpt.foreach(os.makeDir.all(_))
+    val library = os.temp(
+      Library.libraryJar(build),
+      dir = scratchDirOpt.orNull,
+      deleteOnExit = scratchDirOpt.isEmpty,
+      prefix = "spark-job",
+      suffix = ".jar"
+    )
+
+    val finalCommand =
+      Seq(submitCommand, "--class", mainClass) ++
+        jarsArgs ++
+        javaOpts.flatMap(opt => Seq("--driver-java-options", opt)) ++
+        Seq(library.toString) ++
+        args
+    val envUpdates = javaHomeInfo.envUpdates(sys.env)
+    if (showCommand)
+      Left(Runner.envCommand(envUpdates) ++ finalCommand)
+    else {
+      val proc =
+        if (allowExecve)
+          Runner.maybeExec("spark-submit", finalCommand, logger, extraEnv = envUpdates)
+        else
+          Runner.run(finalCommand, logger, extraEnv = envUpdates)
+      Right((
+        proc,
+        if (scratchDirOpt.isEmpty) Some(() => os.remove(library, checkExists = true))
+        else None
+      ))
+    }
+  }
+
+}