-
Notifications
You must be signed in to change notification settings - Fork 4.2k
/
interactive_runner.py
348 lines (296 loc) · 13.9 KB
/
interactive_runner.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
#
# 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.
#
"""A runner that allows running of Beam pipelines interactively.
This module is experimental. No backwards-compatibility guarantees.
"""
# pytype: skip-file
import logging
from typing import Optional
import apache_beam as beam
from apache_beam import runners
from apache_beam.options.pipeline_options import FlinkRunnerOptions
from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import WorkerOptions
from apache_beam.pipeline import PipelineVisitor
from apache_beam.runners.direct import direct_runner
from apache_beam.runners.interactive import interactive_environment as ie
from apache_beam.runners.interactive import pipeline_instrument as inst
from apache_beam.runners.interactive import background_caching_job
from apache_beam.runners.interactive.dataproc.types import ClusterMetadata
from apache_beam.runners.interactive.display import pipeline_graph
from apache_beam.runners.interactive.options import capture_control
from apache_beam.runners.interactive.utils import to_element_list
from apache_beam.runners.interactive.utils import watch_sources
from apache_beam.testing.test_stream_service import TestStreamServiceController
# size of PCollection samples cached.
SAMPLE_SIZE = 8
_LOGGER = logging.getLogger(__name__)
class InteractiveRunner(runners.PipelineRunner):
"""An interactive runner for Beam Python pipelines.
Allows interactively building and running Beam Python pipelines.
"""
def __init__(
self,
underlying_runner=None,
render_option=None,
skip_display=True,
force_compute=True,
blocking=True):
"""Constructor of InteractiveRunner.
Args:
underlying_runner: (runner.PipelineRunner)
render_option: (str) this parameter decides how the pipeline graph is
rendered. See display.pipeline_graph_renderer for available options.
skip_display: (bool) whether to skip display operations when running the
pipeline. Useful if running large pipelines when display is not
needed.
force_compute: (bool) whether sequential pipeline runs can use cached data
of PCollections computed from the previous runs including show API
invocation from interactive_beam module. If True, always run the whole
pipeline and compute data for PCollections forcefully. If False, use
available data and run minimum pipeline fragment to only compute data
not available.
blocking: (bool) whether the pipeline run should be blocking or not.
"""
self._underlying_runner = (
underlying_runner or direct_runner.DirectRunner())
self._render_option = render_option
self._in_session = False
self._skip_display = skip_display
self._force_compute = force_compute
self._blocking = blocking
def is_fnapi_compatible(self):
# TODO(https://github.com/apache/beam/issues/19937):
# return self._underlying_runner.is_fnapi_compatible()
return False
def set_render_option(self, render_option):
"""Sets the rendering option.
Args:
render_option: (str) this parameter decides how the pipeline graph is
rendered. See display.pipeline_graph_renderer for available options.
"""
self._render_option = render_option
def start_session(self):
"""Start the session that keeps back-end managers and workers alive.
"""
if self._in_session:
return
enter = getattr(self._underlying_runner, '__enter__', None)
if enter is not None:
_LOGGER.info('Starting session.')
self._in_session = True
enter()
else:
_LOGGER.error('Keep alive not supported.')
def end_session(self):
"""End the session that keeps backend managers and workers alive.
"""
if not self._in_session:
return
exit = getattr(self._underlying_runner, '__exit__', None)
if exit is not None:
self._in_session = False
_LOGGER.info('Ending session.')
exit(None, None, None)
def apply(self, transform, pvalueish, options):
# TODO(qinyeli, BEAM-646): Remove runner interception of apply.
return self._underlying_runner.apply(transform, pvalueish, options)
def run_pipeline(self, pipeline, options):
if not ie.current_env().options.enable_recording_replay:
capture_control.evict_captured_data()
if self._force_compute:
ie.current_env().evict_computed_pcollections()
# Make sure that sources without a user reference are still cached.
watch_sources(pipeline)
user_pipeline = ie.current_env().user_pipeline(pipeline)
from apache_beam.runners.portability.flink_runner import FlinkRunner
if isinstance(self._underlying_runner, FlinkRunner):
self.configure_for_flink(user_pipeline, options)
pipeline_instrument = inst.build_pipeline_instrument(pipeline, options)
# The user_pipeline analyzed might be None if the pipeline given has nothing
# to be cached and tracing back to the user defined pipeline is impossible.
# When it's None, there is no need to cache including the background
# caching job and no result to track since no background caching job is
# started at all.
if user_pipeline:
# Should use the underlying runner and run asynchronously.
background_caching_job.attempt_to_run_background_caching_job(
self._underlying_runner, user_pipeline, options)
if (background_caching_job.has_source_to_cache(user_pipeline) and
not background_caching_job.is_a_test_stream_service_running(
user_pipeline)):
streaming_cache_manager = ie.current_env().get_cache_manager(
user_pipeline)
# Only make the server if it doesn't exist already.
if (streaming_cache_manager and
not ie.current_env().get_test_stream_service_controller(
user_pipeline)):
def exception_handler(e):
_LOGGER.error(str(e))
return True
test_stream_service = TestStreamServiceController(
streaming_cache_manager, exception_handler=exception_handler)
test_stream_service.start()
ie.current_env().set_test_stream_service_controller(
user_pipeline, test_stream_service)
pipeline_to_execute = beam.pipeline.Pipeline.from_runner_api(
pipeline_instrument.instrumented_pipeline_proto(),
self._underlying_runner,
options)
if ie.current_env().get_test_stream_service_controller(user_pipeline):
endpoint = ie.current_env().get_test_stream_service_controller(
user_pipeline).endpoint
# TODO: make the StreamingCacheManager and TestStreamServiceController
# constructed when the InteractiveEnvironment is imported.
class TestStreamVisitor(PipelineVisitor):
def visit_transform(self, transform_node):
from apache_beam.testing.test_stream import TestStream
if (isinstance(transform_node.transform, TestStream) and
not transform_node.transform._events):
transform_node.transform._endpoint = endpoint
pipeline_to_execute.visit(TestStreamVisitor())
if not self._skip_display:
a_pipeline_graph = pipeline_graph.PipelineGraph(
pipeline_instrument.original_pipeline_proto,
render_option=self._render_option)
a_pipeline_graph.display_graph()
main_job_result = PipelineResult(
pipeline_to_execute.run(), pipeline_instrument)
# In addition to this pipeline result setting, redundant result setting from
# outer scopes are also recommended since the user_pipeline might not be
# available from within this scope.
if user_pipeline:
ie.current_env().set_pipeline_result(user_pipeline, main_job_result)
if self._blocking:
main_job_result.wait_until_finish()
if main_job_result.state is beam.runners.runner.PipelineState.DONE:
# pylint: disable=bad-option-value
ie.current_env().mark_pcollection_computed(
pipeline_instrument.cached_pcolls)
return main_job_result
def configure_for_flink(
self, user_pipeline: beam.Pipeline, options: PipelineOptions) -> None:
"""Configures the pipeline options for running a job with Flink.
When running with a FlinkRunner, a job server started from an uber jar
(locally built or remotely downloaded) hosting the beam_job_api will
communicate with the Flink cluster located at the given flink_master in the
pipeline options.
"""
clusters = ie.current_env().clusters
if clusters.pipelines.get(user_pipeline, None):
# Noop for a known pipeline using a known Dataproc cluster.
return
flink_master = self._strip_protocol_if_any(
options.view_as(FlinkRunnerOptions).flink_master)
cluster_metadata = clusters.default_cluster_metadata
if flink_master == '[auto]':
# Try to create/reuse a cluster when no flink_master is given.
project_id = options.view_as(GoogleCloudOptions).project
region = options.view_as(GoogleCloudOptions).region or 'us-central1'
if project_id:
if clusters.default_cluster_metadata:
# Reuse the cluster name from default in case of a known cluster.
cluster_metadata = ClusterMetadata(
project_id=project_id,
region=region,
cluster_name=clusters.default_cluster_metadata.cluster_name)
else:
# Generate the metadata with a new unique cluster name.
cluster_metadata = ClusterMetadata(
project_id=project_id, region=region)
# Add additional configurations.
self._worker_options_to_cluster_metadata(options, cluster_metadata)
# else use the default cluster metadata.
elif flink_master in clusters.master_urls:
cluster_metadata = clusters.cluster_metadata(flink_master)
else: # Noop if a self-hosted Flink is in use.
return
if not cluster_metadata:
return # Not even a default cluster to create/reuse, run Flink locally.
dcm = clusters.create(cluster_metadata)
# Side effects associated with the user_pipeline.
clusters.pipelines[user_pipeline] = dcm
dcm.pipelines.add(user_pipeline)
self._configure_flink_options(
options,
clusters.DATAPROC_FLINK_VERSION,
dcm.cluster_metadata.master_url)
def _strip_protocol_if_any(self, flink_master: Optional[str]):
if flink_master:
parts = flink_master.split('://')
if len(parts) > 1:
return parts[1]
return flink_master
def _worker_options_to_cluster_metadata(
self, options: PipelineOptions, cluster_metadata: ClusterMetadata):
worker_options = options.view_as(WorkerOptions)
if worker_options.subnetwork:
cluster_metadata.subnetwork = worker_options.subnetwork
if worker_options.num_workers:
cluster_metadata.num_workers = worker_options.num_workers
if worker_options.machine_type:
cluster_metadata.machine_type = worker_options.machine_type
def _configure_flink_options(
self, options: PipelineOptions, flink_version: str, master_url: str):
flink_options = options.view_as(FlinkRunnerOptions)
flink_options.flink_version = flink_version
# flink_options.flink_job_server_jar will be populated by the
# apache_beam.utils.subprocess_server.JavaJarServer.path_to_beam_jar,
# do not populate it explicitly.
flink_options.flink_master = master_url
class PipelineResult(beam.runners.runner.PipelineResult):
"""Provides access to information about a pipeline."""
def __init__(self, underlying_result, pipeline_instrument):
"""Constructor of PipelineResult.
Args:
underlying_result: (PipelineResult) the result returned by the underlying
runner running the pipeline.
pipeline_instrument: (PipelineInstrument) pipeline instrument describing
the pipeline being executed with interactivity applied and related
metadata including where the interactivity-backing cache lies.
"""
super().__init__(underlying_result.state)
self._underlying_result = underlying_result
self._pipeline_instrument = pipeline_instrument
@property
def state(self):
return self._underlying_result.state
def wait_until_finish(self):
self._underlying_result.wait_until_finish()
def get(self, pcoll, include_window_info=False):
"""Materializes the PCollection into a list.
If include_window_info is True, then returns the elements as
WindowedValues. Otherwise, return the element as itself.
"""
return list(self.read(pcoll, include_window_info))
def read(self, pcoll, include_window_info=False):
"""Reads the PCollection one element at a time from cache.
If include_window_info is True, then returns the elements as
WindowedValues. Otherwise, return the element as itself.
"""
key = self._pipeline_instrument.cache_key(pcoll)
cache_manager = ie.current_env().get_cache_manager(
self._pipeline_instrument.user_pipeline)
if key and cache_manager.exists('full', key):
coder = cache_manager.load_pcoder('full', key)
reader, _ = cache_manager.read('full', key)
return to_element_list(reader, coder, include_window_info)
else:
raise ValueError('PCollection not available, please run the pipeline.')
def cancel(self):
self._underlying_result.cancel()