tgravescs commented on a change in pull request #26170: [SPARK-29397][core] Extend plugin interface to include the driver. URL: https://github.com/apache/spark/pull/26170#discussion_r338577951
########## File path: core/src/main/scala/org/apache/spark/internal/plugin/PluginContainer.scala ########## @@ -0,0 +1,140 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.internal.plugin + +import scala.collection.JavaConverters._ +import scala.util.{Either, Left, Right} + +import org.apache.spark.{SparkContext, SparkEnv} +import org.apache.spark.api.plugin._ +import org.apache.spark.internal.Logging +import org.apache.spark.internal.config._ +import org.apache.spark.util.Utils + +sealed abstract class PluginContainer { + + def shutdown(): Unit + +} + +private class DriverPluginContainer(sc: SparkContext, plugins: Seq[SparkPlugin]) + extends PluginContainer with Logging { + + private val driverPlugins: Seq[(String, DriverPlugin)] = plugins.flatMap { p => + val driverPlugin = p.driverPlugin() + if (driverPlugin != null) { + val name = p.getClass().getName() + val ctx = new PluginContextImpl(name, sc.env.rpcEnv, sc.env.metricsSystem, sc.conf, + sc.env.executorId) + + val extraConf = driverPlugin.init(sc, ctx) + if (extraConf != null) { + extraConf.asScala.foreach { case (k, v) => + sc.conf.set(s"${PluginContainer.EXTRA_CONF_PREFIX}$name.$k", v) + } + } + ctx.registerMetrics() + logInfo(s"Initialized driver component for plugin $name.") + Some(p.getClass().getName() -> driverPlugin) + } else { + None + } + } + + if (driverPlugins.nonEmpty) { + sc.env.rpcEnv.setupEndpoint(classOf[PluginEndpoint].getName(), + new PluginEndpoint(driverPlugins.toMap, sc.env.rpcEnv)) + } + + override def shutdown(): Unit = { + driverPlugins.foreach { case (name, plugin) => + try { + logDebug(s"Stopping plugin $name.") + plugin.shutdown() + } catch { + case t: Throwable => + logInfo(s"Exception while shutting down plugin $name.", t) + } + } + } + +} + +private class ExecutorPluginContainer(env: SparkEnv, plugins: Seq[SparkPlugin]) + extends PluginContainer with Logging { + + private val executorPlugins: Seq[(String, ExecutorPlugin)] = { + val allExtraConf = env.conf.getAllWithPrefix(PluginContainer.EXTRA_CONF_PREFIX) + + plugins.flatMap { p => + val executorPlugin = p.executorPlugin() + if (executorPlugin != null) { + val name = p.getClass().getName() + val prefix = name + "." + val extraConf = allExtraConf + .filter { case (k, v) => k.startsWith(prefix) } + .map { case (k, v) => k.substring(prefix.length()) -> v } + .toMap + .asJava + val ctx = new PluginContextImpl(name, env.rpcEnv, env.metricsSystem, env.conf, + env.executorId) + executorPlugin.init(ctx, extraConf) + ctx.registerMetrics() + + logInfo(s"Initialized executor component for plugin $name.") + Some(p.getClass().getName() -> executorPlugin) + } else { + None + } + } + } + + override def shutdown(): Unit = { + executorPlugins.foreach { case (name, plugin) => + try { + logDebug(s"Stopping plugin $name.") + plugin.shutdown() + } catch { + case t: Throwable => + logInfo(s"Exception while shutting down plugin $name.", t) + } + } + } +} + +object PluginContainer { + + val EXTRA_CONF_PREFIX = "spark.plugins.__internal_conf__." Review comment: I get that that internal() doesn't mean anything to the config name, maybe a bad comparison, its more of keeping naming consistent. To me its a lot more obvious if a config has .internal. in its name that its internal to spark. Users should ignore those. that is why I suggested it. If its .internal. I could also programmatically "grep" for all internal configs fairly easily. Not sure why I would want to do this, other than maybe hide them from user. All the other spark configs either follow format x.y.z with the last one optionally camel case, so why not keep that consistent instead of breaking that convention with the __something__ format. I know our internal configs now have no special name on them, which personally I don't like either as its not obvious its meant to be internal. The only benefit to that is you can easily change to not be internal if you want without changing the name. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected] With regards, Apache Git Services --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
