#

# 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 portable "runner" that renders a beam graph.

This runner can either render the graph to a (set of) output path(s), as

designated by (possibly repeated) --render_output, or serve the pipeline as

an interactive graph, if --render_port is set.

In Python, this runner can be passed directly at pipeline construction, e.g.::

with beam.Pipeline(runner=beam.runners.render.RenderRunner(), options=...)

For other languages, start this service a by running::

python -m apache_beam.runners.render --job_port=PORT ...

and then run your pipeline with the PortableRunner setting the job endpoint

to `localhost:PORT`.

If any `--render_output=path.ext` flags are passed, each submitted job will

get written to the given output (overwriting any previously existing file).

If `--render_port` is set to a non-negative value, a local http server will

be started which allows for interactive exploration of the pipeline graph.

As an alternative to starting a job server, a single pipeline can be rendered

by passing a pipeline proto file to `--pipeline_proto`. For example

python -m apache_beam.runners.render \\

--pipeline_proto gs://<staging_location>/pipeline.pb \\

--render_output=/tmp/pipeline.svg

Requires the graphviz dot executable to be available in the path.

"""

import argparse

import base64

import collections

import http.server

import json

import logging

import os

import re

import subprocess

import sys

import tempfile

import threading

import time

import urllib.parse

from google.protobuf import json_format

from google.protobuf import text_format

from apache_beam.options import pipeline_options

from apache_beam.portability.api import beam_runner_api_pb2

from apache_beam.runners import runner

from apache_beam.runners.portability import local_job_service

from apache_beam.runners.portability import local_job_service_main

from apache_beam.runners.portability.fn_api_runner import translations

try:

from apache_beam.io.gcp import gcsio

except ImportError:

gcsio = None # type: ignore

_LOGGER = logging.getLogger(__name__)

# From the Beam site, circa November 2022.

DEFAULT_EDGE_STYLE = 'color="#ff570b"'

DEFAULT_TRANSFORM_STYLE = (

'shape=rect style="rounded, filled" color="#ff570b" fillcolor="#fff6dd"')

DEFAULT_HIGHLIGHT_STYLE = (

'shape=rect style="rounded, filled" color="#ff570b" fillcolor="#ffdb97"')

class RenderOptions(pipeline_options.PipelineOptions):

"""Rendering options."""

@classmethod

def _add_argparse_args(cls, parser):

parser.add_argument(

'--render_port',

type=int,

default=-1,

help='The port at which to serve the graph. '

'If 0, an unused port will be chosen. '

'If -1, the server will not be started.')

parser.add_argument(

'--render_output',

action='append',

help='A path or paths to which to write rendered output. '

'The output type will be deduced from the file extension.')

parser.add_argument(

'--render_leaf_composite_nodes',

action='append',

help='A set of regular expressions for transform names that should '

'not be expanded. For example, one could pass "\bRead.*" to indicate '

'the inner structure of read nodes should not be expanded. '

'If not given, defaults to the top-level nodes if interactively '

'serving the graph and expanding all nodes otherwise.')

parser.add_argument(

'--render_edge_attributes',

default='',

help='Graphviz attributes to add to all edges.')

parser.add_argument(

'--render_node_attributes',

default='',

help='Graphviz attributes to add to all nodes.')

parser.add_argument(

'--render_highlight_attributes',

default='',

help='Graphviz attributes to add to all highlighted nodes.')

parser.add_argument(

'--log_proto',

default=False,

action='store_true',

help='Set to also log input pipeline proto to stdout.')

return parser

class PipelineRenderer:

def __init__(self, pipeline, options):

self.pipeline = pipeline

self.options = options

# Drill down into any uninteresting, top-level transforms that contain

# the whole pipeline (often added by the SDK).

roots = self.pipeline.root_transform_ids

while len(roots) == 1:

root = self.pipeline.components.transforms[roots[0]]

if not root.subtransforms:

break

roots = root.subtransforms

self.roots = roots

# Figure out at what point to stop rendering composite internals.

if options.render_leaf_composite_nodes:

is_leaf = lambda transform_id: any(

re.match(

pattern, self.pipeline.components.transforms[transform_id].

unique_name) for patterns in options.render_leaf_composite_nodes

for pattern in patterns.split(','))

self.leaf_composites = set()

def mark_leaves(transform_ids):

for transform_id in transform_ids:

if is_leaf(transform_id):

self.leaf_composites.add(transform_id)

else:

mark_leaves(

self.pipeline.components.transforms[transform_id].subtransforms)

mark_leaves(self.roots)

elif options.render_port >= 0:

# Start interactive with no unfolding.

self.leaf_composites = set(self.roots)

else:

# For non-interactive, expand fully.

self.leaf_composites = set()

# Useful for attempting graph layout consistency.

self.latest_positions = {}

self.highlighted = []

def update(self, toggle=None):

if toggle:

transform_id = toggle[0]

self.highlighted = [transform_id]

if transform_id in self.leaf_composites:

transform = self.pipeline.components.transforms[transform_id]

if transform.subtransforms:

self.leaf_composites.remove(transform_id)

for subtransform in transform.subtransforms:

self.leaf_composites.add(subtransform)

if transform_id in self.latest_positions:

self.latest_positions[subtransform] = self.latest_positions[

transform_id]

else:

self.leaf_composites.add(transform_id)

def style(self, transform_id):

base = ' '.join(

[DEFAULT_TRANSFORM_STYLE, self.options.render_node_attributes])

if transform_id in self.highlighted:

return ' '.join([

base,

DEFAULT_HIGHLIGHT_STYLE,

self.options.render_highlight_attributes

])

else:

return base

def to_dot(self):

return '\n'.join(self.to_dot_iter())

def to_dot_iter(self):

yield 'digraph G {'

# Defer drawing any edges until the end lest we declare nodes too early.

edges_out = []

for transform_id in self.roots:

yield from self.transform_to_dot(

transform_id, self.pcoll_leaf_consumers(), edges_out)

yield from edges_out

yield '}'

def transform_to_dot(self, transform_id, pcoll_leaf_consumers, edges_out):

transform = self.pipeline.components.transforms[transform_id]

if self.is_leaf(transform_id):

yield self.transform_node(transform_id)

transform_inputs = set(transform.inputs.values())

for name, output in transform.outputs.items():

# For outputs that are also inputs, it's ambiguous whether they are

# consumed as the outputs of this transform, or of the upstream

# transform. Render the latter.

if output in transform_inputs:

continue

output_label = name if len(transform.outputs) > 1 else ''

for consumer, is_side_input in pcoll_leaf_consumers[output]:

# Can't yield this here as the consumer might not be in this cluster.

edge_style = 'dashed' if is_side_input else 'solid'

edge_attributes = ' '.join([

f'label="{output_label}" style={edge_style}',

DEFAULT_EDGE_STYLE,

self.options.render_edge_attributes

])

edges_out.append(

f'"{transform_id}" -> "{consumer}" [{edge_attributes}]')

else:

yield f'subgraph "cluster_{transform_id}" {{'

yield self.transform_attributes(transform_id)

for subtransform in transform.subtransforms:

yield from self.transform_to_dot(

subtransform, pcoll_leaf_consumers, edges_out)

yield '}'

def transform_node(self, transform_id):

return f'"{transform_id}" [{self.transform_attributes(transform_id)}]'

def transform_attributes(self, transform_id):

transform = self.pipeline.components.transforms[transform_id]

local_name = transform.unique_name.split('/')[-1]

if transform_id in self.latest_positions:

pos_str = f'pos="{self.latest_positions[transform_id]}"'

else:

pos_str = ''

return (

f'label="{local_name}" {self.style(transform_id)} '

f'URL="javascript:click(\'{transform_id}\')" {pos_str}')

def pcoll_leaf_consumers_iter(self, transform_id):

transform = self.pipeline.components.transforms[transform_id]

transform_inputs = set(transform.inputs.values())

side_inputs = set(translations.side_inputs(transform).values())

if self.is_leaf(transform_id):

for pcoll in transform.inputs.values():

yield pcoll, (transform_id, pcoll in side_inputs)

for subtransform in transform.subtransforms:

for pcoll, (consumer,

annotation) in self.pcoll_leaf_consumers_iter(subtransform):

if self.is_leaf(transform_id):

if pcoll not in transform_inputs:

yield pcoll, (transform_id, annotation)

else:

yield pcoll, (consumer, annotation)

def pcoll_leaf_consumers(self):

result = collections.defaultdict(list)

for transform_id in self.roots:

for pcoll, consumer_info in self.pcoll_leaf_consumers_iter(transform_id):

result[pcoll].append(consumer_info)

return result

def is_leaf(self, transform_id):

return (

transform_id in self.leaf_composites or

not self.pipeline.components.transforms[transform_id].subtransforms)

def info(self):

if len(self.highlighted) != 1:

return ''

transform_id = self.highlighted[0]

return f'<pre>{self.pipeline.components.transforms[transform_id]}</pre>'

def layout_dot(self):

layout = subprocess.run(['dot', '-Tdot'],

input=self.to_dot().encode('utf-8'),

capture_output=True,

check=True).stdout

# Try to capture the positions for layout consistency.

json_out = json.loads(

subprocess.run(['dot', '-n2', '-Kneato', '-Tjson'],

input=layout,

capture_output=True,

check=True).stdout)

for box in json_out['objects']:

name = box.get('name', None)

if name in self.pipeline.components.transforms:

if 'pos' in box:

self.latest_positions[name] = box['pos']

elif 'bb' in box:

x0, y0, x1, y1 = [float(r) for r in box['bb'].split(',')]

self.latest_positions[name] = f'{(x0+x1)/2},{(y0+y1)/2}'

return layout

def page_callback_data(self, layout):

svg = subprocess.run(['dot', '-Kneato', '-n2', '-Tsvg'],

input=layout,

capture_output=True,

check=True).stdout

cmapx = subprocess.run(['dot', '-Kneato', '-n2', '-Tcmapx'],

input=layout,

capture_output=True,

check=True).stdout

return {

'src': 'data:image/svg+xml;base64,' +

base64.b64encode(svg).decode('utf-8'),

'cmapx': cmapx.decode('utf-8'),

'info': self.info(),

}

def render_data(self):

_LOGGER.info("Re-rendering pipeline...")

layout = self.layout_dot()

if self.options.render_output:

for path in self.options.render_output:

format = os.path.splitext(path)[-1][1:]

result = subprocess.run(

['dot', '-Kneato', '-n2', '-T' + format, '-o', path],

input=layout,

check=False)

if result.returncode:

_LOGGER.error(

"Failed render pipeline as %r: exit %s", path, result.returncode)

else:

_LOGGER.info("Rendered pipeline as %r", path)

return self.page_callback_data(layout)

def render_json(self):

return json.dumps(self.render_data())

def page(self):

data = self.render_data()

src = data['src']

cmapx = data['cmapx']

return """

<html>

<head>

<script>

function click(transform_id) {

var xhttp = new XMLHttpRequest();

xhttp.onreadystatechange = function() {

render_data = JSON.parse(this.responseText);

document.getElementById('image_map_holder').innerHTML =

render_data.cmapx;

document.getElementById('image_tag').src = render_data.src

document.getElementById('info').innerHTML = render_data.info

};

xhttp.open("GET", "render?toggle=" + transform_id, true);

xhttp.send();

}

</script>

</head>

""" + f"""

<body>

Click on a composite transform to expand.

<br>

<img id='image_tag' src='{src}' usemap='#G'>

<hr>

<div id='info'></div>

<div id='image_map_holder'>

{cmapx}

</div>

</body>

</html>

"""

class RenderRunner(runner.PipelineRunner):

# TODO(robertwb): Consider making this a runner wrapper, where live status

# (such as counters, stage completion status, or possibly even PCollection

# samples) queryable and/or displayed. This could evolve into a full Beam

# UI.

def run_pipeline(self, pipeline_object, options):

return self.run_portable_pipeline(pipeline_object.to_runner_api(), options)

def run_portable_pipeline(self, pipeline_proto, options):

render_options = options.view_as(RenderOptions)

if render_options.render_port < 0 and not render_options.render_output:

raise ValueError(

'At least one of --render_port or --render_output must be provided.')

if render_options.log_proto:

_LOGGER.info(pipeline_proto)

renderer = PipelineRenderer(pipeline_proto, render_options)

try:

subprocess.run(['dot', '-V'], capture_output=True, check=True)

except FileNotFoundError as exn:

# If dot is not available, we can at least output the raw .dot files.

dot_files = [

output for output in render_options.render_output

if output.endswith('.dot')

]

for output in dot_files:

with open(output, 'w') as fout:

fout.write(renderer.to_dot())

_LOGGER.info("Wrote pipeline as %s", output)

non_dot_files = set(render_options.render_output) - set(dot_files)

if non_dot_files:

raise RuntimeError(

"Graphviz dot executable not available "

f"for rendering non-dot output files {non_dot_files}") from exn

elif render_options.render_port >= 0:

raise RuntimeError(

"Graphviz dot executable not available for serving") from exn

return RenderPipelineResult(None)

renderer.page()

if render_options.render_port >= 0:

# TODO: If this gets more complex, we could consider taking on a

# framework like Flask as a dependency.

class RequestHandler(http.server.BaseHTTPRequestHandler):

def do_GET(self):

parts = urllib.parse.urlparse(self.path)

args = urllib.parse.parse_qs(parts.query)

renderer.update(**args)

if parts.path == '/':

response = renderer.page()

elif parts.path == '/render':

response = renderer.render_json()

else:

self.send_response(400)

return

self.send_response(200)

self.send_header("Content-type", "text/html")

self.end_headers()

self.wfile.write(response.encode('utf-8'))

server = http.server.HTTPServer(('localhost', render_options.render_port),

RequestHandler)

server_thread = threading.Thread(target=server.serve_forever, daemon=True)

server_thread.start()

print('Serving at http://%s:%s' % server.server_address)

return RenderPipelineResult(server)

else:

return RenderPipelineResult(None)

class RenderPipelineResult(runner.PipelineResult):

def __init__(self, server):

super().__init__(runner.PipelineState.RUNNING)

self.server = server

def wait_until_finish(self, duration=None):

if self.server:

time.sleep(duration or 1e8)

self.server.shutdown()

self._state = runner.PipelineState.DONE

def monitoring_infos(self):

return []

def run(argv):

if argv[0] == __file__:

argv = argv[1:]

parser = argparse.ArgumentParser(

description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)

parser.add_argument(

'--job_port',

type=int,

default=0,

help='port on which to serve the job api')

parser.add_argument(

'--pipeline_proto', help='file containing the beam pipeline definition')

RenderOptions._add_argparse_args(parser)

options = parser.parse_args(argv)

if options.pipeline_proto:

if not options.render_output and options.render_port < 0:

options.render_port = 0

render_one(options)

if options.render_output:

return

run_server(options)

def render_one(options):

if options.pipeline_proto == '-':

content = sys.stdin.buffer.read()

if content[0] == b'{':

ext = '.json'

else:

try:

content.decode('utf-8')

ext = '.textproto'

except UnicodeDecodeError:

ext = '.pb'

else:

if options.pipeline_proto.startswith('gs://'):

if gcsio is None:

raise ImportError('GCS not available; please install apache_beam[gcp]')

open_fn = gcsio.GcsIO().open

else:

open_fn = open

with open_fn(options.pipeline_proto, 'rb') as fin:

content = fin.read()

ext = os.path.splitext(options.pipeline_proto)[-1]

if ext == '.textproto':

pipeline_proto = text_format.Parse(content, beam_runner_api_pb2.Pipeline())

elif ext == '.json':

pipeline_proto = json_format.Parse(content, beam_runner_api_pb2.Pipeline())

else:

pipeline_proto = beam_runner_api_pb2.Pipeline()

pipeline_proto.ParseFromString(content)

RenderRunner().run_portable_pipeline(

pipeline_proto, pipeline_options.PipelineOptions(**vars(options)))

def run_server(options):

class RenderBeamJob(local_job_service.BeamJob):

def _invoke_runner(self):

return RenderRunner().run_portable_pipeline(

self._pipeline_proto,

pipeline_options.PipelineOptions(**vars(options)))

with tempfile.TemporaryDirectory() as staging_dir:

job_servicer = local_job_service.LocalJobServicer(

staging_dir, beam_job_type=RenderBeamJob)

port = job_servicer.start_grpc_server(options.job_port)

try:

local_job_service_main.serve("Listening for beam jobs on port %d." % port)

finally:

job_servicer.stop()

if __name__ == '__main__':

logging.basicConfig()

logging.getLogger().setLevel(logging.INFO)

run(sys.argv)