Files
deb-python-taskflow/tools/state_graph.py
Joshua Harlow 2fa4af7a24 Split-off the additional retry states from the task states
Split the states that are not task states (but are retry
states) into there own additional set and then use that set
and a new function to validate the transition at other locations
in the code-base.

This makes the transitions that are valid for tasks/retries
easily viewable, more easy to read and understand, and
more correct (instead of being a mix of task + retry atom
transitions and states).

Change-Id: I9515c19daf59a21e581f51e757ece2050f348214
2015-06-18 07:41:25 -07:00

211 lines
7.3 KiB
Python
Executable File

#!/usr/bin/env python
# Copyright (C) 2014 Yahoo! Inc. All Rights Reserved.
#
# Licensed 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.
import optparse
import os
import sys
top_dir = os.path.abspath(os.path.join(os.path.dirname(__file__),
os.pardir))
sys.path.insert(0, top_dir)
# To get this installed you have to do the following:
#
# $ pip install pydot2
import pydot
from taskflow.engines.action_engine import runner
from taskflow.engines.worker_based import protocol
from taskflow import states
from taskflow.types import fsm
# This is just needed to get at the runner builder object (we will not
# actually be running it...).
class DummyRuntime(object):
def __init__(self):
self.analyzer = None
self.completer = None
self.scheduler = None
self.storage = None
def clean_event(name):
name = name.replace("_", " ")
name = name.strip()
return name
def make_machine(start_state, transitions):
machine = fsm.FSM(start_state)
machine.add_state(start_state)
for (start_state, end_state) in transitions:
if start_state not in machine:
machine.add_state(start_state)
if end_state not in machine:
machine.add_state(end_state)
# Make a fake event (not used anyway)...
event = "on_%s" % (end_state)
machine.add_transition(start_state, end_state, event.lower())
return machine
def map_color(internal_states, state):
if state in internal_states:
return 'blue'
if state == states.FAILURE:
return 'red'
if state == states.REVERTED:
return 'darkorange'
if state in (states.SUCCESS, states.COMPLETE):
return 'green'
return None
def main():
parser = optparse.OptionParser()
parser.add_option("-f", "--file", dest="filename",
help="write svg to FILE", metavar="FILE")
parser.add_option("-t", "--tasks", dest="tasks",
action='store_true',
help="use task state transitions",
default=False)
parser.add_option("-r", "--retries", dest="retries",
action='store_true',
help="use retry state transitions",
default=False)
parser.add_option("-e", "--engines", dest="engines",
action='store_true',
help="use engine state transitions",
default=False)
parser.add_option("-w", "--wbe-requests", dest="wbe_requests",
action='store_true',
help="use wbe request transitions",
default=False)
parser.add_option("-j", "--jobs", dest="jobs",
action='store_true',
help="use job transitions",
default=False)
parser.add_option("-T", "--format", dest="format",
help="output in given format",
default='svg')
(options, args) = parser.parse_args()
if options.filename is None:
options.filename = 'states.%s' % options.format
types = [
options.engines,
options.retries,
options.tasks,
options.wbe_requests,
options.jobs,
]
if sum([int(i) for i in types]) > 1:
parser.error("Only one of task/retry/engines/wbe requests/jobs"
" may be specified.")
internal_states = list()
ordering = 'in'
if options.tasks:
source_type = "Tasks"
source = make_machine(states.PENDING,
list(states._ALLOWED_TASK_TRANSITIONS))
elif options.retries:
source_type = "Retries"
source = make_machine(states.PENDING,
list(states._ALLOWED_RETRY_TRANSITIONS))
elif options.engines:
source_type = "Engines"
r = runner.Runner(DummyRuntime(), None)
source, memory = r.builder.build()
internal_states.extend(runner._META_STATES)
ordering = 'out'
elif options.wbe_requests:
source_type = "WBE requests"
source = make_machine(protocol.WAITING,
list(protocol._ALLOWED_TRANSITIONS))
elif options.jobs:
source_type = "Jobs"
source = make_machine(states.UNCLAIMED,
list(states._ALLOWED_JOB_TRANSITIONS))
else:
source_type = "Flow"
source = make_machine(states.PENDING,
list(states._ALLOWED_FLOW_TRANSITIONS))
graph_name = "%s states" % source_type
g = pydot.Dot(graph_name=graph_name, rankdir='LR',
nodesep='0.25', overlap='false',
ranksep="0.5", size="11x8.5",
splines='true', ordering=ordering)
node_attrs = {
'fontsize': '11',
}
nodes = {}
for (start_state, on_event, end_state) in source:
on_event = clean_event(on_event)
if start_state not in nodes:
start_node_attrs = node_attrs.copy()
text_color = map_color(internal_states, start_state)
if text_color:
start_node_attrs['fontcolor'] = text_color
nodes[start_state] = pydot.Node(start_state, **start_node_attrs)
g.add_node(nodes[start_state])
if end_state not in nodes:
end_node_attrs = node_attrs.copy()
text_color = map_color(internal_states, end_state)
if text_color:
end_node_attrs['fontcolor'] = text_color
nodes[end_state] = pydot.Node(end_state, **end_node_attrs)
g.add_node(nodes[end_state])
if options.engines:
edge_attrs = {
'label': on_event,
}
if 'reverted' in on_event:
edge_attrs['fontcolor'] = 'darkorange'
if 'fail' in on_event:
edge_attrs['fontcolor'] = 'red'
if 'success' in on_event:
edge_attrs['fontcolor'] = 'green'
else:
edge_attrs = {}
g.add_edge(pydot.Edge(nodes[start_state], nodes[end_state],
**edge_attrs))
start = pydot.Node("__start__", shape="point", width="0.1",
xlabel='start', fontcolor='green', **node_attrs)
g.add_node(start)
g.add_edge(pydot.Edge(start, nodes[source.start_state], style='dotted'))
print("*" * len(graph_name))
print(graph_name)
print("*" * len(graph_name))
print(source.pformat())
print(g.to_string().strip())
g.write(options.filename, format=options.format)
print("Created %s at '%s'" % (options.format, options.filename))
# To make the svg more pretty use the following:
# $ xsltproc ../diagram-tools/notugly.xsl ./states.svg > pretty-states.svg
# Get diagram-tools from https://github.com/vidarh/diagram-tools.git
if __name__ == '__main__':
main()