Merge pull request #1129 from alexarchambault/run-spark-hadoop

Add --spark / --standalone-spark / --hadoop to the run command

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,19 @@ 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("Run as a Spark job, using a vanilla Spark distribution downloaded by Scala CLI")
+  @ExtraName("sparkStandalone")
+    standaloneSpark: Option[Boolean] = None,
+  @Group("Run")
+  @HelpMessage("Run as a Hadoop job, using the \"hadoop jar\" command")
+  @ExtraName("hadoop")
+    hadoopJar: Boolean = false,
+  @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
@@ -760,7 +760,7 @@ object Package extends ScalaCommand[PackageOptions] {
     providedFiles
   }
 
-  private def assembly(
+  def assembly(
     build: Build.Successful,
     destPath: os.Path,
     mainClassOpt: Option[String],

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

@@ -10,18 +10,30 @@ 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.{RunHadoop, 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.standaloneSpark.getOrElse(false) && !options.sparkSubmit.contains(false))
+      RunMode.StandaloneSparkSubmit
+    else if (options.sparkSubmit.getOrElse(false))
+      RunMode.SparkSubmit
+    else if (options.hadoopJar)
+      RunMode.HadoopJar
+    else
+      RunMode.Default
+
   private def scratchDirOpt(options: RunOptions): Option[os.Path] =
     options.scratchDir
       .filter(_.trim.nonEmpty)
@@ -48,7 +60,31 @@ 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.StandaloneSparkSubmit | RunMode.SparkSubmit | RunMode.HadoopJar =>
+              Some("8")
+            case RunMode.Default => None
+          }
+        }
+      ),
+      internalDependencies = baseOptions.internalDependencies.copy(
+        addRunnerDependencyOpt = baseOptions.internalDependencies.addRunnerDependencyOpt.orElse {
+          runMode(options) match {
+            case RunMode.StandaloneSparkSubmit | RunMode.SparkSubmit | RunMode.HadoopJar =>
+              Some(false)
+            case RunMode.Default => None
+          }
+        }
+      ),
+      internal = baseOptions.internal.copy(
+        keepResolution = baseOptions.internal.keepResolution || {
+          runMode(options) match {
+            case RunMode.StandaloneSparkSubmit | RunMode.SparkSubmit | RunMode.HadoopJar => true
+            case RunMode.Default                                                         => false
+          }
+        }
       ),
       notForBloopOptions = baseOptions.notForBloopOptions.copy(
         runWithManifest = options.useManifest
@@ -76,6 +112,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 +132,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 +209,7 @@ object Run extends ScalaCommand[RunOptions] {
             val maybeProcess = maybeRun(
               s,
               allowTerminate = false,
+              runMode = runMode(options),
               showCommand = options.command,
               scratchDirOpt = scratchDirOpt(options)
             )
@@ -206,6 +246,7 @@ object Run extends ScalaCommand[RunOptions] {
           val res = maybeRun(
             s,
             allowTerminate = true,
+            runMode = runMode(options),
             showCommand = options.command,
             scratchDirOpt = scratchDirOpt(options)
           )
@@ -226,9 +267,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 +292,7 @@ object Run extends ScalaCommand[RunOptions] {
       finalArgs,
       logger,
       allowExecve,
+      runMode,
       showCommand,
       scratchDirOpt
     )
@@ -262,9 +305,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 +352,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 +371,74 @@ 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
+              )
+            }
+          case RunMode.StandaloneSparkSubmit =>
+            value {
+              RunSpark.runStandalone(
+                build,
+                mainClass,
+                args,
+                logger,
+                allowExecve,
+                showCommand,
+                scratchDirOpt
+              )
+            }
+          case RunMode.HadoopJar =>
+            value {
+              RunHadoop.run(
+                build,
+                mainClass,
+                args,
+                logger,
+                allowExecve,
+                showCommand,
+                scratchDirOpt
+              )
+            }
         }
     }
   }

modules/cli/src/main/scala/scala/cli/commands/packaging/Spark.scala

@@ -23,4 +23,10 @@ object Spark {
 
   def sparkModules: Seq[AnyModule] =
     names.map(name => mod"org.apache.spark::spark-$name")
+
+  def hadoopModules: Seq[AnyModule] =
+    Seq(
+      // TODO Add more for Hadoop 2, maybe for 3 too
+      mod"org.apache.hadoop:hadoop-client-api"
+    )
 }

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

@@ -0,0 +1,10 @@
+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
+  case object StandaloneSparkSubmit extends RunMode
+  case object HadoopJar             extends RunMode
+}

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

@@ -0,0 +1,71 @@
+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
+
+object RunHadoop {
+
+  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.hadoopModules via provided settings?
+    val providedModules = Spark.hadoopModules
+    scratchDirOpt.foreach(os.makeDir.all(_))
+    val assembly = os.temp(
+      dir = scratchDirOpt.orNull,
+      prefix = "hadoop-job",
+      suffix = ".jar",
+      deleteOnExit = scratchDirOpt.isEmpty
+    )
+    value {
+      PackageCmd.assembly(
+        build,
+        assembly,
+        // "hadoop jar" doesn't accept a main class as second argument if the jar as first argument has a main class in its manifest…
+        None,
+        providedModules,
+        withPreamble = false,
+        () => (),
+        logger
+      )
+    }
+
+    val javaOpts = build.options.javaOptions.javaOpts.toSeq.map(_.value.value)
+    val extraEnv =
+      if (javaOpts.isEmpty) Map[String, String]()
+      else
+        Map(
+          "HADOOP_CLIENT_OPTS" -> javaOpts.mkString(" ") // no escaping…
+        )
+    val hadoopJarCommand = Seq("hadoop", "jar")
+    val finalCommand =
+      hadoopJarCommand ++ Seq(assembly.toString, mainClass) ++ args
+    if (showCommand)
+      Left(Runner.envCommand(extraEnv) ++ finalCommand)
+    else {
+      val proc =
+        if (allowExecve)
+          Runner.maybeExec("hadoop", finalCommand, logger, extraEnv = extraEnv)
+        else
+          Runner.run(finalCommand, logger, extraEnv = extraEnv)
+      Right((
+        proc,
+        if (scratchDirOpt.isEmpty) Some(() => os.remove(assembly, checkExists = true))
+        else None
+      ))
+    }
+
+  }
+
+}