[beam] branch master updated: Fix BCJ termination conditions.

[pa...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

pabloem pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new 6bc39bc  Fix BCJ termination conditions.
     new 31af8b1  Merge pull request #12702 from [BEAM-10603] Fix BCJ termination conditions.
6bc39bc is described below

commit 6bc39bc52f517000b9e8b59c56ede57d49e11ef3
Author: Sam Rohde <sr...@google.com>
AuthorDate: Thu Aug 27 11:19:39 2020 -0700

    Fix BCJ termination conditions.
    
    Change-Id: I24eb5a94e1bdafa1833ede0d04c9e2b824417f87
---
 .../apache_beam/runners/interactive/background_caching_job.py     | 8 ++++----
 .../runners/interactive/options/capture_control_test.py           | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/sdks/python/apache_beam/runners/interactive/background_caching_job.py b/sdks/python/apache_beam/runners/interactive/background_caching_job.py
index 9346adb..ec1c3c4 100644
--- a/sdks/python/apache_beam/runners/interactive/background_caching_job.py
+++ b/sdks/python/apache_beam/runners/interactive/background_caching_job.py
@@ -87,11 +87,10 @@ class BackgroundCachingJob(object):
     return any([l.is_triggered() for l in self._limiters])
 
   def is_done(self):
-    is_terminated = self._pipeline_result.state is PipelineState.DONE
+    is_terminated = self._pipeline_result.state in (
+        PipelineState.DONE, PipelineState.CANCELLED)
     is_triggered = self._should_end_condition_checker()
-    is_cancelling = (
-        self._pipeline_result.state in (
-            PipelineState.CANCELLED, PipelineState.CANCELLING))
+    is_cancelling = self._pipeline_result.state is PipelineState.CANCELLING
     return is_terminated or (is_triggered and is_cancelling)
 
   def is_running(self):
@@ -103,6 +102,7 @@ class BackgroundCachingJob(object):
     if not PipelineState.is_terminal(self._pipeline_result.state):
       try:
         self._pipeline_result.cancel()
+        self._pipeline_result.wait_until_finish()
       except NotImplementedError:
         # Ignore the cancel invocation if it is never implemented by the runner.
         pass
diff --git a/sdks/python/apache_beam/runners/interactive/options/capture_control_test.py b/sdks/python/apache_beam/runners/interactive/options/capture_control_test.py
index f68ee85..86561a9 100644
--- a/sdks/python/apache_beam/runners/interactive/options/capture_control_test.py
+++ b/sdks/python/apache_beam/runners/interactive/options/capture_control_test.py
@@ -172,7 +172,7 @@ class CaptureControlTest(unittest.TestCase):
 
     limiter = FakeLimiter()
     background_caching_job = bcj.BackgroundCachingJob(
-        runner.PipelineResult(runner.PipelineState.CANCELLED),
+        runner.PipelineResult(runner.PipelineState.CANCELLING),
         limiters=[limiter])
     ie.current_env().set_background_caching_job(p, background_caching_job)