[SPARK-29582][PYSPARK] Support `TaskContext.get()` in a barrier task from Python side #26239

ConeyLiu · 2019-10-24T02:53:05Z

What changes were proposed in this pull request?

Add support of TaskContext.get() in a barrier task from Python side, this makes it easier to migrate legacy user code to barrier execution mode.

Why are the changes needed?

In Spark Core, there is a TaskContext object which is a singleton. We set a task context instance which can be TaskContext or BarrierTaskContext before the task function startup, and unset it to none after the function end. So we can both get TaskContext and BarrierTaskContext with the object. However we can only get the BarrierTaskContext with BarrierTaskContext, we will get None if we get it by TaskContext.get in a barrier stage.

This is useful when people switch from normal code to barrier code, and only need a little update.

Does this PR introduce any user-facing change?

Yes.
Previously:

def func(iterator):
    task_context = TaskContext.get() . # this could be None.
    barrier_task_context = BarrierTaskContext.get() # get the BarrierTaskContext instance
    ...

rdd.barrier().mapPartitions(func)

Proposed:

def func(iterator):
    task_context = TaskContext.get() . # this could also get the BarrierTaskContext instance which is same as barrier_task_context
    barrier_task_context = BarrierTaskContext.get() # get the BarrierTaskContext instance
    ...

rdd.barrier().mapPartitions(func)

How was this patch tested?

New UT tests.

ConeyLiu · 2019-10-24T02:54:21Z

Hi, @cloud-fan @HyukjinKwon, could you help to review this? Thanks a lot.

jiangxb1987 · 2019-10-24T13:10:10Z

This is useful when people switch from normal code to barrier code

Do you have any use cases that people want to reuse their production code and migrate on to barrier execution mode ?

ConeyLiu · 2019-10-24T13:53:51Z

Yeah, we use rdd.mappartitionsWithIndex to startup a distributed dl model training, which is smilier to horovod.spark. Now we want to switch to barrier mode. And we want to reduce the code difference between scala and python. I think this patch is also useful for other people. 在 2019年10月24日，下午9:12，Jiang Xingbo <notifications@github.com<mailto:notifications@github.com>> 写道： This is useful when people switch from normal code to barrier code Do you have any use cases that people want to reuse their production code and migrate on to barrier execution mode ? — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub<#26239?email_source=notifications&email_token=ADBEWSE5BQPFGSB6HUI72K3QQGNKJA5CNFSM4JEN6GRKYY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOECE625I#issuecomment-545910133>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/ADBEWSDQTLZZPSDEYTRVC5TQQGNKJANCNFSM4JEN6GRA>.

BryanCutler · 2019-10-24T18:59:57Z

ok to test

SparkQA · 2019-10-24T19:18:12Z

Test build #112617 has finished for PR 26239 at commit c8fe7cc.

This patch fails PySpark unit tests.
This patch merges cleanly.
This patch adds no public classes.

python/pyspark/tests/test_taskcontext.py

SparkQA · 2019-10-25T02:01:23Z

Test build #112638 has finished for PR 26239 at commit e3e5872.

This patch fails PySpark unit tests.
This patch merges cleanly.
This patch adds no public classes.

SparkQA · 2019-10-25T08:37:59Z

Test build #112656 has finished for PR 26239 at commit a55fb6b.

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

python/pyspark/taskcontext.py

python/pyspark/tests/test_taskcontext.py

HyukjinKwon · 2019-10-28T10:16:53Z

python/pyspark/worker.py

+
+        # reset task context to None
+        TaskContext._setTaskContext(None)
+        BarrierTaskContext._setTaskContext(None)


Hm, what happens if it fails with exceptions in the middle of execution in this worker?

Is it really needed? We always set the global TaskContext and never reset it previouslly.

Hm, what happens if it fails with exceptions in the middle of execution in this worker?

If exceptions occured, the worker will be closed with sys.exit(-1).

Is it really needed? We always set the global TaskContext and never reset it previouslly.

Previously:

val rdd = ... val barriered = rdd.barrier().mapPartitions(...) barriered.mapPartitions(...) # here the BarrierTaskContext still existed.

This code is just a guard program, it shouldn't increase extra overhead or behavior change.

This code is just a guard program, it shouldn't increase extra overhead or behavior change.

I guess that's only when the worker is reused. Can you clarify it with comments here?

thanks for reviewing. updated.

HyukjinKwon · 2019-10-28T10:22:44Z

Looks making sense to me otherwise

python/pyspark/worker.py

SparkQA · 2019-10-29T04:56:20Z

Test build #112818 has finished for PR 26239 at commit dc89212.

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

jiangxb1987 · 2019-10-29T22:20:23Z

Please also update the previous and proposed behaviors in the PR description, like:

Previously:
rdd.barrier().mapPartitions {.....}
>> ......
Proposed:
rdd.barrier().mapPartitions {.....}
>> ......

jiangxb1987 · 2019-10-29T22:21:57Z

python/pyspark/taskcontext.py

+            A RuntimeError will raise if it is not in a barrier stage.
        """
+        if not isinstance(cls._taskContext, BarrierTaskContext):
+            raise RuntimeError('''It is not in a barrier stage''')


This is not supported behavior, so please raise Exception instead of RuntimeError

Yeah, I think Exception is consistent at least.

thanks for reviewing. updated

SparkQA · 2019-10-30T13:52:45Z

Test build #112925 has finished for PR 26239 at commit 04a2b0e.

This patch fails PySpark unit tests.
This patch merges cleanly.
This patch adds no public classes.

SparkQA · 2019-10-30T14:33:59Z

Test build #112927 has finished for PR 26239 at commit 384e1d2.

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

jiangxb1987 · 2019-10-30T17:44:35Z

LGTM

HyukjinKwon · 2019-10-31T01:48:07Z

python/pyspark/taskcontext.py

        running tasks.

        .. note:: Must be called on the worker, not the driver. Returns None if not initialized.
+            A RuntimeError will raise if it is not in a barrier stage.


Suggested change

A RuntimeError will raise if it is not in a barrier stage.

An exception will raise if it is not in a barrier stage.

python/pyspark/tests/test_taskcontext.py

SparkQA · 2019-10-31T02:51:38Z

Test build #112986 has finished for PR 26239 at commit 7d0f0d2.

This patch fails PySpark unit tests.
This patch merges cleanly.
This patch adds no public classes.

SparkQA · 2019-10-31T03:51:29Z

Test build #112989 has finished for PR 26239 at commit a3d14c8.

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

HyukjinKwon · 2019-10-31T04:09:24Z

Merged to master.

ConeyLiu · 2019-10-31T04:52:09Z

thanks all !

Unify the behavior of pyspark.TaskContext with spark core

c8fe7cc

update

e3e5872

ConeyLiu commented Oct 25, 2019

View reviewed changes

python/pyspark/tests/test_taskcontext.py Show resolved Hide resolved

dongjoon-hyun added the PYSPARK label Oct 25, 2019

fix UT errors

a55fb6b

HyukjinKwon reviewed Oct 28, 2019

View reviewed changes

python/pyspark/taskcontext.py Outdated Show resolved Hide resolved

python/pyspark/taskcontext.py Outdated Show resolved Hide resolved

python/pyspark/tests/test_taskcontext.py Outdated Show resolved Hide resolved

HyukjinKwon reviewed Oct 28, 2019

View reviewed changes

cloud-fan reviewed Oct 28, 2019

View reviewed changes

python/pyspark/worker.py Outdated Show resolved Hide resolved

address comments

dc89212

jiangxb1987 changed the title ~~[SPARK-29582][PYSPARK] Unify the behavior of pyspark.TaskContext with spark core~~ [SPARK-29582][PYSPARK] Support TaskContext.get() in a barrier task from Python side Oct 29, 2019

jiangxb1987 reviewed Oct 29, 2019

View reviewed changes

address comments

04a2b0e

fix UT errors

384e1d2

HyukjinKwon approved these changes Oct 31, 2019

View reviewed changes

HyukjinKwon reviewed Oct 31, 2019

View reviewed changes

python/pyspark/tests/test_taskcontext.py Outdated Show resolved Hide resolved

HyukjinKwon reviewed Oct 31, 2019

View reviewed changes

python/pyspark/tests/test_taskcontext.py Outdated Show resolved Hide resolved

ConeyLiu added 2 commits October 31, 2019 10:25

address comments

74e3515

typo

7d0f0d2

fix errors

a3d14c8

HyukjinKwon closed this in 1e599e5 Oct 31, 2019

ConeyLiu deleted the barrier_task_context branch October 31, 2019 04:52

	A RuntimeError will raise if it is not in a barrier stage.
	An exception will raise if it is not in a barrier stage.

[SPARK-29582][PYSPARK] Support TaskContext.get() in a barrier task from Python side #26239

[SPARK-29582][PYSPARK] Support TaskContext.get() in a barrier task from Python side #26239

Uh oh!

Conversation

ConeyLiu commented Oct 24, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Uh oh!

ConeyLiu commented Oct 24, 2019

Uh oh!

jiangxb1987 commented Oct 24, 2019

Uh oh!

ConeyLiu commented Oct 24, 2019 via email

Uh oh!

BryanCutler commented Oct 24, 2019

Uh oh!

SparkQA commented Oct 24, 2019

Uh oh!

Uh oh!

SparkQA commented Oct 25, 2019

Uh oh!

SparkQA commented Oct 25, 2019

Uh oh!

Uh oh!

Uh oh!

Uh oh!

HyukjinKwon Oct 28, 2019 • 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!

cloud-fan Oct 28, 2019

Choose a reason for hiding this comment

Uh oh!

ConeyLiu Oct 29, 2019 • 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 Oct 30, 2019

Choose a reason for hiding this comment

Uh oh!

ConeyLiu Oct 30, 2019

Choose a reason for hiding this comment

Uh oh!

HyukjinKwon commented Oct 28, 2019

Uh oh!

Uh oh!

SparkQA commented Oct 29, 2019

Uh oh!

jiangxb1987 commented Oct 29, 2019

Uh oh!

jiangxb1987 Oct 29, 2019

Choose a reason for hiding this comment

Uh oh!

HyukjinKwon Oct 30, 2019

Choose a reason for hiding this comment

Uh oh!

ConeyLiu Oct 30, 2019

Choose a reason for hiding this comment

Uh oh!

SparkQA commented Oct 30, 2019

Uh oh!

SparkQA commented Oct 30, 2019

Uh oh!

jiangxb1987 commented Oct 30, 2019

Uh oh!

HyukjinKwon Oct 31, 2019

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

SparkQA commented Oct 31, 2019

Uh oh!

SparkQA commented Oct 31, 2019

Uh oh!

HyukjinKwon commented Oct 31, 2019

Uh oh!

ConeyLiu commented Oct 31, 2019

[SPARK-29582][PYSPARK] Support `TaskContext.get()` in a barrier task from Python side #26239

[SPARK-29582][PYSPARK] Support `TaskContext.get()` in a barrier task from Python side #26239

ConeyLiu commented Oct 24, 2019 •

edited

Loading

HyukjinKwon Oct 28, 2019 •

edited

Loading

ConeyLiu Oct 29, 2019 •

edited

Loading