def cleanup(self, pipeline=None):
"""Cleans up cached states for the given pipeline. Noop if the given
pipeline is absent from the environment. Cleans up for all pipelines
if no pipeline is specified."""
if pipeline:
from apache_beam.runners.interactive import background_caching_job as bcj
bcj.attempt_to_cancel_background_caching_job(pipeline)
bcj.attempt_to_stop_test_stream_service(pipeline)
cache_manager = self.get_cache_manager(pipeline)
if cache_manager:
cache_manager.cleanup()
else:
for _, job in self._background_caching_jobs.items():
if job:
job.cancel()
for _, controller in self._test_stream_service_controllers.items():
if controller:
controller.stop()
for _, cache_manager in self._cache_managers.items():
if cache_manager:
cache_manager.cleanup()
self.evict_recording_manager(pipeline)
self.evict_background_caching_job(pipeline)
self.evict_test_stream_service_controller(pipeline)
self.evict_computed_pcollections(pipeline)
self.evict_cached_source_signature(pipeline)
self.evict_pipeline_result(pipeline)
def test_stop_a_running_test_stream_service(self):
pipeline = _build_an_empty_stream_pipeline()
test_stream_service = TestStreamServiceController(reader=None)
test_stream_service.start()
ie.current_env().set_test_stream_service_controller(
pipeline, test_stream_service)
bcj.attempt_to_stop_test_stream_service(pipeline)
self.assertFalse(bcj.is_a_test_stream_service_running(pipeline))
def cleanup_pipeline(self, pipeline):
from apache_beam.runners.interactive import background_caching_job as bcj
bcj.attempt_to_cancel_background_caching_job(pipeline)
bcj.attempt_to_stop_test_stream_service(pipeline)
cache_manager = self.get_cache_manager(pipeline)
# Recording manager performs cache manager cleanup during eviction, so we
# don't need to clean it up here.
if cache_manager and self.get_recording_manager(pipeline) is None:
cache_manager.cleanup()
self.clusters.cleanup(pipeline)
def evict_captured_data():
"""Evicts all deterministic replayable data that have been captured by
Interactive Beam. In future PCollection evaluation/visualization and pipeline
runs, Interactive Beam will capture fresh data."""
if ie.current_env().options.enable_capture_replay:
_LOGGER.info(
'You have requested Interactive Beam to evict all captured '
'data that could be deterministically replayed among multiple '
'pipeline runs.')
ie.current_env().track_user_pipelines()
for user_pipeline in ie.current_env().tracked_user_pipelines:
bcj.attempt_to_cancel_background_caching_job(user_pipeline)
bcj.attempt_to_stop_test_stream_service(user_pipeline)
ie.current_env().cleanup()
def test_noop_when_no_test_stream_service_running(self, _mocked_stop):
pipeline = _build_an_empty_stream_pipeline()
self.assertFalse(bcj.is_a_test_stream_service_running(pipeline))
bcj.attempt_to_stop_test_stream_service(pipeline)
_mocked_stop.assert_not_called()
