[SPARK-25095][PySpark] Python support for BarrierTaskContext #22085

jiangxb1987 · 2018-08-12T16:07:16Z

What changes were proposed in this pull request?

Add method barrier() and getTaskInfos() in python TaskContext, these two methods are only allowed for barrier tasks.

How was this patch tested?

Add new tests in tests.py

viirya · 2018-08-12T18:52:25Z

core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala

+          logError(message)
+          throw new SparkException(message)
+        } else {
+          logDebug(s"Started GatewayServer to port BarrierTaskContext on port $boundPort.")


When isBarrier is false, I think we don't need show this?

viirya · 2018-08-12T18:54:18Z

python/pyspark/worker.py


        # initialize global state
        taskContext = TaskContext._getOrCreate()
+        isBarrier = read_bool(infile)


Add a comment indicating the following 3 inputs are only for barrier task?

viirya · 2018-08-12T18:55:04Z

python/pyspark/worker.py

        _accumulatorRegistry.clear()

+        if isBarrier:
+            paras = GatewayParameters(port=boundPort, auth_token=secret, auto_convert=True)


Maybe params instead of paras?

viirya · 2018-08-12T18:56:35Z

python/pyspark/taskcontext.py

    """

    _taskContext = None
+    _javaContext = None


_barrierContext?

viirya · 2018-08-12T18:59:29Z

core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala

+        } else {
+          logDebug(s"Started GatewayServer to port BarrierTaskContext on port $boundPort.")
+        }
        // Write out the TaskContextInfo


This comment should be moved too.

SparkQA · 2018-08-12T20:48:09Z

Test build #94650 has finished for PR 22085 at commit 7b48829.

This patch passes all tests.
This patch merges cleanly.
This patch adds no public classes.

mengxr

Should also add tests.

mengxr · 2018-08-13T00:33:57Z

python/pyspark/taskcontext.py

        """
        return self._localProperties.get(key, None)
+
+    def barrier(self):


Create BarrierTaskContext that extends TaskContext and then move those two methods there.

mengxr · 2018-08-13T00:34:38Z

python/pyspark/taskcontext.py

+            raise Exception("Not supported to call getTaskInfos() inside a non-barrier task.")
+        else:
+            java_list = self._javaContext.getTaskInfos()
+            return [h for h in java_list]


Create BarrierTaskInfo class and wrap it over Java object.

mengxr · 2018-08-13T00:35:19Z

core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala

+            .entryPoint(context.asInstanceOf[BarrierTaskContext])
+            .authToken(secret)
+            .javaPort(0)
+            .callbackClient(GatewayServer.DEFAULT_PYTHON_PORT, GatewayServer.defaultAddress(),


Leave a TODO here. We do not have requests from Java to Python.

HyukjinKwon · 2018-08-13T01:21:43Z

core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala

+            .javaPort(0)
+            .callbackClient(GatewayServer.DEFAULT_PYTHON_PORT, GatewayServer.defaultAddress(),
+              secret)
+            .build())


Wait wait.. you guys sure have another Java gateway for each worker? (or did I rush to read this code?) Can you elaborate why this is needed? We should avoid this unless it's super required or necessary.

We have to port BarrierTaskContext from java to python side, otherwise there is no way to call BarrierTaskContext.barrier() from python side. Thus, of course, the JavaGateway is only initiated when the context is a BarrierTaskContext.

Yea, I read and understood if this is only initialised when the context is a BarrierTaskContext but this is super weird we start another Java gateway here. If it's a hard requirement, then I suspect the design issue. Should this be targeted to 2.4.0, @mengxr?

If this should necessarily target 2.4.0, don't block by me since it's a new feature and probably we could consider another approach later but if we can avoid, I would suggest to avoid for now.

Let me try to track the design doc and changes about this. I think I need more time to check why it happened like this and if there's another way.

@HyukjinKwon Could you elaborate your concerns? Is it because resource usage or security?

Mainly the reason is about resource usage, unusual access pattern via Py4J at worker side, and the possibility of allowing JVM access within Python worker.

It pretty much looks an overkill to launch a Java gateway to allow access to call a function assuming from #22085 (comment). This pattern sounds pretty unusual - such cases, we usually send the data manually and read it in Python side, for instance TaskContext. Now, it opens a gateway for each worker if I am not mistaken.

I was thinking if we can avoid this. Can you elaborate why this is required and necessary? I haven't got enough time to look into this so was thinking about taking a look on this weekends.

This also now opens a possibility for an JVM access from worker side via BarrierTaskContext. For instance, I believe we can hack and access to JVM inside of UDFs.

BTW, this is a design change https://cwiki.apache.org/confluence/display/SPARK/PySpark+Internals .. is this really required to open a gate there?

The major issue here is that we want to make the barrier() call blocking, the task shall wait until timeout or succeeded, do we have other ways to achieve this goal other than current approach here?

Not yet. So I asked to hold it on for now since another gateway here looks the last choice, and was wondering if we can avoid to target 2.4.0. If this blocks, please go ahead. Will check it later on this weekends.

HyukjinKwon · 2018-08-13T08:36:32Z

@mengxr, let me leave this link here again - #22085 (comment) in case this is missed by being folded.

SparkQA · 2018-08-13T12:45:03Z

Test build #94681 has finished for PR 22085 at commit 289146d.

This patch passes all tests.
This patch merges cleanly.
This patch adds the following public classes (experimental):
class BarrierTaskContext(TaskContext):
class BarrierTaskInfo(object):

HyukjinKwon · 2018-08-14T05:53:53Z

python/pyspark/taskcontext.py

+
+    def __init__(self):
+        """Construct a BarrierTaskContext, use get instead"""
+        pass


I would throw an exception BTW if this method should rather be banned

This just follows TaskContext.__init__(), shall we update both?

Ah, this is called in _getOrCreate. Sorry, I rushed to read. In this case, frankly I think we can remove this since that's the default constructor injected by Python or monkey patch to disallow the initialization (like we did for ImageSchema) but I guess we don't necessarily be super clever on this.

I'm okay as is.

mengxr · 2018-08-14T06:59:22Z

python/pyspark/taskcontext.py

+    """
+
+    def __init__(self, info):
+        self.address = info.address


should be info.address()

better to rename info to jobj to make it clear this is from Java

mengxr · 2018-08-14T06:59:32Z

python/pyspark/taskcontext.py

+    """
+
+    def __init__(self, info):
+        self.address = info.address


should be info.address()

better to rename info to jobj to make it clear this is from Java

SparkQA · 2018-08-14T13:04:10Z

Test build #94738 has finished for PR 22085 at commit 05c9609.

This patch passes all tests.
This patch merges cleanly.
This patch adds no public classes.

mengxr · 2018-08-15T20:34:43Z

@HyukjinKwon Thanks for the feedback! We will replace the py4j route by a special implementation that can only trigger "context.barrier()" in JVM.

HyukjinKwon · 2018-08-16T07:40:11Z

Thank you so much @mengxr and @jiangxb1987.

jiangxb1987 · 2018-08-17T16:23:51Z

core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala

+  /**
+   * Gateway to call BarrierTaskContext.barrier().
+   */
+  def barrierAndServe(): Unit = {


It's not clear yet how to trigger this.

jiangxb1987 · 2018-08-17T16:25:08Z

python/pyspark/taskcontext.py

+        else:
+            _load_from_socket(self._port, self._secret)
+
+    def getTaskInfos(self):


This is not available temporarily.

SparkQA · 2018-08-17T22:08:14Z

Test build #94901 has finished for PR 22085 at commit e234a0a.

This patch fails from timeout after a configured wait of `340m`.
This patch merges cleanly.
This patch adds no public classes.

SparkQA · 2018-08-20T07:05:01Z

Test build #94942 has finished for PR 22085 at commit 243a5a3.

This patch fails due to an unknown error code, -9.
This patch merges cleanly.
This patch adds no public classes.

SparkQA · 2018-08-20T12:58:19Z

Test build #94951 has finished for PR 22085 at commit ba0ccad.

This patch passes all tests.
This patch merges cleanly.
This patch adds no public classes.

mengxr · 2018-08-20T18:02:06Z

core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala

 import scala.collection.JavaConverters._

-import org.apache.spark._
+import org.apache.spark.{SparkException, _}


_ should include SparkException already

mengxr · 2018-08-20T18:06:28Z

python/pyspark/taskcontext.py

        return self._localProperties.get(key, None)
+
+
+def _load_from_socket(port, auth_secret):


Should document how this is different from the one in context.py.

mengxr · 2018-08-20T18:08:37Z

core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala

+
+  // Authentication helper used when serving method calls via socket from Python side.
+  private lazy val authHelper = {
+    val conf = Option(SparkEnv.get).map(_.conf).getOrElse(new SparkConf())


When SparkEnv.get returns null?

mengxr · 2018-08-20T18:12:13Z

core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala

+        // Init a ServerSocket to accept method calls from Python side.
+        val isBarrier = context.isInstanceOf[BarrierTaskContext]
+        if (isBarrier) {
+          serverSocket = Some(new ServerSocket(0, 1, InetAddress.getByName("localhost")))


minor: useful to add /* port */ and /* backlog */

mengxr · 2018-08-20T18:14:09Z

core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala

+                  sock = serverSocket.get.accept()
+                  sock.setSoTimeout(10000)
+                  val cmdString = readUtf8(sock)
+                  if (cmdString.equals("run")) {


If we do not expect any other command from the socket, we should throw an exception

mengxr · 2018-08-20T18:14:34Z

core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala

+                var sock: Socket = null
+                try {
+                  sock = serverSocket.get.accept()
+                  sock.setSoTimeout(10000)


Should add a comment about this timeout.

mengxr · 2018-08-20T18:15:30Z

core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala

+                    barrierAndServe(sock)
+                  }
+                } catch {
+                  case _: SocketException =>


Is the the timeout exception? I don't see any exception that we could silently ignore.

mengxr · 2018-08-20T18:18:37Z

core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala

  }
+
+  def writeUTF(str: String, dataOut: DataOutputStream) {
+    val bytes = str.getBytes(StandardCharsets.UTF_8)


nit: UTF_8 or always use StandardCharsets

mengxr · 2018-08-20T18:21:14Z

python/pyspark/taskcontext.py

+    @classmethod
+    def _getOrCreate(cls):
+        """Internal function to get or create global BarrierTaskContext."""
+        if cls._taskContext is None:


Q: Does it handle python worker reuse?

IIUC reuse python worker just means we start a python worker from a daemon thread, it shall not affect the input/output files related to worker.py.

HyukjinKwon · 2018-08-21T02:12:11Z

core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala

+import org.apache.spark.security.SocketAuthHelper
 import org.apache.spark.util._

-


tiny nit: I would remove this one back while addressing other comments.

SparkQA · 2018-08-21T21:41:27Z

Test build #95037 has finished for PR 22085 at commit 2a8f3cb.

This patch passes all tests.
This patch merges cleanly.
This patch adds no public classes.

SparkQA · 2018-08-21T21:52:25Z

Test build #95040 has finished for PR 22085 at commit 1cacd40.

This patch passes all tests.
This patch merges cleanly.
This patch adds no public classes.

mengxr · 2018-08-21T22:55:32Z

LGTM. I'm merging this into master. We might need a minor refactor for readability. But it shouldn't block developers testing this new feature. Thanks!

HyukjinKwon · 2018-08-23T02:48:47Z

python/pyspark/taskcontext.py

+    Load data from a given socket, this is a blocking method thus only return when the socket
+    connection has been closed.
+
+    This is copied from context.py, while modified the message protocol.


It would be nicer if we can deduplciate it later.

HyukjinKwon · 2018-08-23T02:49:58Z

Thanks, @jiangxb1987 and @mengxr again.

dbtsai · 2018-08-24T21:07:14Z

core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala

+        }
+        // Close ServerSocket on task completion.
+        serverSocket.foreach { server =>
+          context.addTaskCompletionListener(_ => server.close())


This is failing the Scala 2.12 build

[error] /Users/d_tsai/dev/apache-spark/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:242: ambiguous reference to overloaded definition, [error] both method addTaskCompletionListener in class TaskContext of type [U](f: org.apache.spark.TaskContext => U)org.apache.spark.TaskContext [error] and method addTaskCompletionListener in class TaskContext of type (listener: org.apache.spark.util.TaskCompletionListener)org.apache.spark.TaskContext [error] match argument types (org.apache.spark.TaskContext => Unit) [error] context.addTaskCompletionListener(_ => server.close()) [error] ^ [error] one error found [error] Compile failed at Aug 24, 2018 1:56:06 PM [31.582s]

Addressed in #22229

squito · 2018-08-27T16:21:30Z

core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala

+                  sock = serverSocket.get.accept()
+                  // Wait for function call from python side.
+                  sock.setSoTimeout(10000)
+                  val input = new DataInputStream(sock.getInputStream())


why isn't authentication the first thing which happens on this connection? I don't think anything bad can happen in this case, but it just makes it more likely we leave a security hole here later on.

(I'd also like to do some refactoring of the socket setup code in python, and that can go further if we do authenticaion first here)

Thanks for catching this, yea I agree it would be better to move the authentication before recognising functions.

ok I'm doing this -- SPARK-25253, will open a pr shortly

implement python barrier taskcontext

7b48829

viirya reviewed Aug 12, 2018

View reviewed changes

mengxr requested changes Aug 13, 2018

View reviewed changes

HyukjinKwon reviewed Aug 13, 2018

View reviewed changes

update

289146d

HyukjinKwon reviewed Aug 14, 2018

View reviewed changes

mengxr reviewed Aug 14, 2018

View reviewed changes

update

05c9609

update;WIP

e234a0a

jiangxb1987 changed the title ~~[SPARK-25095][PySpark] Python support for BarrierTaskContext~~ [WIP][SPARK-25095][PySpark] Python support for BarrierTaskContext Aug 17, 2018

jiangxb1987 commented Aug 17, 2018

View reviewed changes

update

243a5a3

update

ba0ccad

jiangxb1987 changed the title ~~[WIP][SPARK-25095][PySpark] Python support for BarrierTaskContext~~ [SPARK-25095][PySpark] Python support for BarrierTaskContext Aug 20, 2018

mengxr reviewed Aug 20, 2018

View reviewed changes

HyukjinKwon reviewed Aug 21, 2018

View reviewed changes

jiangxb1987 added 2 commits August 22, 2018 00:38

update

2a8f3cb

update

1cacd40

asfgit closed this in ad45299 Aug 21, 2018

HyukjinKwon reviewed Aug 23, 2018

View reviewed changes

dbtsai reviewed Aug 24, 2018

View reviewed changes

squito reviewed Aug 27, 2018

View reviewed changes

		return self._localProperties.get(key, None)


		def _load_from_socket(port, auth_secret):

		import org.apache.spark.security.SocketAuthHelper
		import org.apache.spark.util._

[SPARK-25095][PySpark] Python support for BarrierTaskContext #22085

[SPARK-25095][PySpark] Python support for BarrierTaskContext #22085

Uh oh!

Conversation

jiangxb1987 commented Aug 12, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

How was this patch tested?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SparkQA commented Aug 12, 2018

Uh oh!

mengxr left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HyukjinKwon Aug 14, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HyukjinKwon Aug 15, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HyukjinKwon commented Aug 13, 2018

Uh oh!

SparkQA commented Aug 13, 2018

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mengxr Aug 14, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SparkQA commented Aug 14, 2018

Uh oh!

mengxr commented Aug 15, 2018

Uh oh!

HyukjinKwon commented Aug 16, 2018

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jiangxb1987 commented Aug 12, 2018 •

edited

Loading

HyukjinKwon Aug 14, 2018 •

edited

Loading

HyukjinKwon Aug 15, 2018 •

edited

Loading

mengxr Aug 14, 2018 •

edited

Loading