Skip to content

Instantly share code, notes, and snippets.

@namitkewat

namitkewat/AWS StepFunction to Graphviz.ipynb

Created July 11, 2025 12:32
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save namitkewat/effea48cc39803fc8300b342b90e1441 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save namitkewat/effea48cc39803fc8300b342b90e1441 to your computer and use it in GitHub Desktop.
Download ZIP
Visualizing AWS Step Functions with Python and Graphviz
Display the source blob
Display the rendered blob
Raw
AWS StepFunction to Graphviz.ipynb
{
"cells": [
{
"cell_type": "markdown",
"id": "28ca94d5-eb20-489e-b8de-b873f4503dde",
"metadata": {},
"source": [
"# Visualizing AWS Step Functions with Python and Graphviz\n",
"\n",
"AWS Step Functions are powerful for orchestrating serverless workflows, but their JSON definitions can be difficult to visualize. This notebook demonstrates a Python function that converts a Step Function's Amazon States Language (ASL) JSON into a DOT language graph, which can then be rendered into a clear, intuitive diagram using Graphviz.\n",
"\n",
"## 1. Initial Setup\n",
"\n",
"First, we import the necessary libraries. We'll use json to parse the state machine definition and graphviz to render the final diagram. Image from IPython.display is used to show the graph within the notebook environment."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "024019f4-e488-48c5-9f44-ba4120f4465e",
"metadata": {},
"outputs": [],
"source": [
"from IPython.display import Image\n",
"import graphviz"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "e56459cc-f538-4e22-bdb9-6869acaa833a",
"metadata": {},
"outputs": [],
"source": [
"import json"
]
},
{
"cell_type": "markdown",
"id": "6647424f-f78a-45c9-a486-568cb68e175a",
"metadata": {},
"source": [
"## 2. The Conversion Logic\n",
"This is the core function that traverses the Step Function JSON. It identifies each state type (`Task`, `Choice`, `Parallel`, etc.) and maps it to a corresponding Graphviz node shape and style. It also handles the transitions between states, including conditional logic from `Choice` states and branches from `Parallel` states, to create the connecting edges in the graph."
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "3cfb023d-ed1f-48f6-aa37-8d616d8eccdc",
"metadata": {},
"outputs": [],
"source": [
"def stepfunction_to_graphviz(state_machine_json, graph_name=\"StateMachine\", rankdir=\"LR\"):\n",
" \"\"\"\n",
" Converts AWS Step Functions JSON definition to Graphviz DOT format.\n",
" \n",
" Args:\n",
" state_machine_json: Step Function JSON definition (as string or dict)\n",
" graph_name: Name for the Graphviz graph\n",
" rankdir: Graph direction (e.g., \"LR\" for Left-to-Right, \"TB\" for Top-to-Bottom, \"BT\" for Bottom to Top, \"RL\" for Right to Left)\n",
" \n",
" Returns:\n",
" Graphviz DOT format string\n",
" \"\"\"\n",
" # Ensure the input is a dictionary\n",
" if isinstance(state_machine_json, str):\n",
" data = json.loads(state_machine_json)\n",
" else:\n",
" data = state_machine_json\n",
" \n",
" start_state = data[\"StartAt\"]\n",
" states = data[\"States\"]\n",
" \n",
" # Helper to generate human-readable labels for Choice state conditions\n",
" def get_condition_label(condition):\n",
" if \"Variable\" in condition:\n",
" var = condition[\"Variable\"]\n",
" # Find first operator key\n",
" operators = [k for k in condition if k not in \n",
" (\"Variable\", \"Next\", \"And\", \"Or\", \"Not\")]\n",
" if operators:\n",
" op = operators[0]\n",
" val = condition[op]\n",
" return f\"{var} {op} {json.dumps(val)}\"\n",
" if \"And\" in condition:\n",
" return \"AND Condition\"\n",
" if \"Or\" in condition:\n",
" return \"OR Condition\"\n",
" if \"Not\" in condition:\n",
" return \"NOT Condition\"\n",
" return \"Condition\"\n",
"\n",
" # Initialize the DOT language string with graph attributes\n",
" dot_lines = [\n",
" f'digraph {graph_name} {{',\n",
" f' rankdir=\"{rankdir}\";',\n",
" ' node [shape=rectangle, fontname=\"Arial\", style=\"rounded\"];',\n",
" ' start [shape=point, width=0.15];',\n",
" f' start -> \"{start_state}\";'\n",
" ]\n",
" \n",
" # First pass: define all nodes with their specific shapes and styles\n",
" for name, state in states.items():\n",
" state_type = state[\"Type\"]\n",
" node_attrs = []\n",
" \n",
" # Set shape based on state type\n",
" if state_type == \"Choice\":\n",
" node_attrs.append('shape=diamond')\n",
" elif state_type == \"Wait\":\n",
" node_attrs.append('shape=ellipse')\n",
" elif state_type == \"Parallel\":\n",
" node_attrs.append('shape=box')\n",
" node_attrs.append('style=\"filled,bold\"')\n",
" node_attrs.append('fillcolor=\"#d5e8d4\"')\n",
" elif state_type == \"Succeed\":\n",
" node_attrs.append('shape=doublecircle')\n",
" node_attrs.append('color=\"#82b366\"')\n",
" node_attrs.append('penwidth=2')\n",
" elif state_type == \"Fail\":\n",
" node_attrs.append('shape=doublecircle')\n",
" node_attrs.append('color=\"#ff9999\"')\n",
" node_attrs.append('penwidth=2')\n",
" elif state_type == \"Pass\":\n",
" node_attrs.append('style=\"rounded,dashed\"')\n",
" \n",
" # Mark end states\n",
" if state.get(\"End\", False) and state_type not in (\"Succeed\", \"Fail\"):\n",
" node_attrs.append('peripheries=2')\n",
" \n",
" # Format node attributes\n",
" attrs = \", \".join(node_attrs)\n",
" dot_lines.append(f' \"{name}\" [{attrs}];')\n",
"\n",
" # Second pass: define all edges (transitions) between nodes\n",
" for name, state in states.items():\n",
" state_type = state[\"Type\"]\n",
" \n",
" # Handle Choice state branches\n",
" if state_type == \"Choice\":\n",
" for choice in state[\"Choices\"]:\n",
" label = get_condition_label(choice)\n",
" dot_lines.append(\n",
" f' \"{name}\" -> \"{choice[\"Next\"]}\" [label=\"{label}\"];'\n",
" )\n",
" if \"Default\" in state:\n",
" dot_lines.append(\n",
" f' \"{name}\" -> \"{state[\"Default\"]}\" [label=\"Default\"];'\n",
" )\n",
" \n",
" # Handle normal transitions\n",
" elif \"Next\" in state:\n",
" dot_lines.append(f' \"{name}\" -> \"{state[\"Next\"]}\";')\n",
" \n",
" # Handle Parallel branches\n",
" elif state_type == \"Parallel\":\n",
" for i, branch in enumerate(state.get(\"Branches\", [])):\n",
" branch_start = branch[\"StartAt\"]\n",
" dot_lines.append(\n",
" f' \"{name}\" -> \"{branch_start}\" '\n",
" f'[label=\"Branch {i+1}\", style=\"dashed\"];'\n",
" )\n",
" if \"Next\" in state:\n",
" dot_lines.append(f' \"{name}\" -> \"{state[\"Next\"]}\";')\n",
" \n",
" dot_lines.append(\"}\")\n",
" return \"\\n\".join(dot_lines)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7d4c9f53-f2c7-4533-9cfb-9145d0c3140a",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "markdown",
"id": "b04a5480-a0c0-4229-b299-66553b3531f7",
"metadata": {},
"source": [
"## 3. Example Usage\n",
"Here is a sample Step Function definition from the AWS documentation. We'll pass this JSON string to our function to generate the DOT graph definition.\n",
"\n",
"\n",
"https://docs.aws.amazon.com/step-functions/latest/dg/concepts-amazon-states-language.html"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "fd570771-7bb2-4725-b750-06c7bb213147",
"metadata": {},
"outputs": [],
"source": [
"aws_state_machine_json_sample = '''{\n",
" \"Comment\": \"An example of the Amazon States Language using a choice state.\",\n",
" \"QueryLanguage\": \"JSONata\",\n",
" \"StartAt\": \"FirstState\",\n",
" \"States\": {\n",
" \"FirstState\": {\n",
" \"Type\": \"Task\",\n",
" \"Assign\": {\n",
" \"foo\" : \"{% $states.input.foo_input %}\" \n",
" },\n",
" \"Resource\": \"arn:aws:lambda:region:123456789012:function:FUNCTION_NAME\",\n",
" \"Next\": \"ChoiceState\"\n",
" },\n",
" \"ChoiceState\": {\n",
" \"Type\": \"Choice\",\n",
" \"Default\": \"DefaultState\",\n",
" \"Choices\": [\n",
" {\n",
" \"Next\": \"FirstMatchState\",\n",
" \"Condition\": \"{% $foo = 1 %}\"\n",
" },\n",
" {\n",
" \"Next\": \"SecondMatchState\",\n",
" \"Condition\": \"{% $foo = 2 %}\"\n",
" }\n",
" ]\n",
" },\n",
" \"FirstMatchState\": {\n",
" \"Type\" : \"Task\",\n",
" \"Resource\": \"arn:aws:lambda:region:123456789012:function:OnFirstMatch\",\n",
" \"Next\": \"NextState\"\n",
" },\n",
"\n",
" \"SecondMatchState\": {\n",
" \"Type\" : \"Task\",\n",
" \"Resource\": \"arn:aws:lambda:region:123456789012:function:OnSecondMatch\",\n",
" \"Next\": \"NextState\"\n",
" },\n",
"\n",
" \"DefaultState\": {\n",
" \"Type\": \"Fail\",\n",
" \"Error\": \"DefaultStateError\",\n",
" \"Cause\": \"No Matches!\"\n",
" },\n",
"\n",
" \"NextState\": {\n",
" \"Type\": \"Task\",\n",
" \"Resource\": \"arn:aws:lambda:region:123456789012:function:FUNCTION_NAME\",\n",
" \"End\": true\n",
" }\n",
" }\n",
"}'''"
]
},
{
"cell_type": "markdown",
"id": "c5a352bf-db5f-4e2f-99a1-160865f61238",
"metadata": {},
"source": [
"## 4. Rendering the Graph\n",
"Finally, we create a Graphviz `Source` object from the generated DOT string and display it. The result is a clear, easy-to-read flowchart of the state machine's logic."
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "a602ffea-d8e2-4bf7-a879-a02c6f1c270f",
"metadata": {},
"outputs": [
{
"data": {
"image/svg+xml": [
"<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
"<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
" \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
"<!-- Generated by graphviz version 2.44.0 (0)\n",
" -->\n",
"<!-- Title: StateMachine Pages: 1 -->\n",
"<svg width=\"678pt\" height=\"242pt\"\n",
" viewBox=\"0.00 0.00 678.25 242.14\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
"<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 238.14)\">\n",
"<title>StateMachine</title>\n",
"<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-238.14 674.25,-238.14 674.25,4 -4,4\"/>\n",
"<!-- start -->\n",
"<g id=\"node1\" class=\"node\">\n",
"<title>start</title>\n",
"<ellipse fill=\"black\" stroke=\"black\" cx=\"5.4\" cy=\"-162.14\" rx=\"5.4\" ry=\"5.4\"/>\n",
"</g>\n",
"<!-- FirstState -->\n",
"<g id=\"node2\" class=\"node\">\n",
"<title>FirstState</title>\n",
"<path fill=\"none\" stroke=\"black\" d=\"M109.8,-180.14C109.8,-180.14 59.8,-180.14 59.8,-180.14 53.8,-180.14 47.8,-174.14 47.8,-168.14 47.8,-168.14 47.8,-156.14 47.8,-156.14 47.8,-150.14 53.8,-144.14 59.8,-144.14 59.8,-144.14 109.8,-144.14 109.8,-144.14 115.8,-144.14 121.8,-150.14 121.8,-156.14 121.8,-156.14 121.8,-168.14 121.8,-168.14 121.8,-174.14 115.8,-180.14 109.8,-180.14\"/>\n",
"<text text-anchor=\"middle\" x=\"84.8\" y=\"-158.44\" font-family=\"Arial\" font-size=\"14.00\">FirstState</text>\n",
"</g>\n",
"<!-- start&#45;&gt;FirstState -->\n",
"<g id=\"edge1\" class=\"edge\">\n",
"<title>start&#45;&gt;FirstState</title>\n",
"<path fill=\"none\" stroke=\"black\" d=\"M11.01,-162.14C16.59,-162.14 26.6,-162.14 37.45,-162.14\"/>\n",
"<polygon fill=\"black\" stroke=\"black\" points=\"37.54,-165.64 47.54,-162.14 37.54,-158.64 37.54,-165.64\"/>\n",
"</g>\n",
"<!-- ChoiceState -->\n",
"<g id=\"node3\" class=\"node\">\n",
"<title>ChoiceState</title>\n",
"<path fill=\"none\" stroke=\"black\" d=\"M229.8,-177.6C229.8,-177.6 170.3,-164.69 170.3,-164.69 164.44,-163.42 164.44,-160.87 170.3,-159.6 170.3,-159.6 229.8,-146.69 229.8,-146.69 235.66,-145.42 247.39,-145.42 253.25,-146.69 253.25,-146.69 312.75,-159.6 312.75,-159.6 318.61,-160.87 318.61,-163.42 312.75,-164.69 312.75,-164.69 253.25,-177.6 253.25,-177.6 247.39,-178.87 235.66,-178.87 229.8,-177.6\"/>\n",
"<text text-anchor=\"middle\" x=\"241.52\" y=\"-158.44\" font-family=\"Arial\" font-size=\"14.00\">ChoiceState</text>\n",
"</g>\n",
"<!-- FirstState&#45;&gt;ChoiceState -->\n",
"<g id=\"edge2\" class=\"edge\">\n",
"<title>FirstState&#45;&gt;ChoiceState</title>\n",
"<path fill=\"none\" stroke=\"black\" d=\"M122.08,-162.14C130.21,-162.14 139.2,-162.14 148.48,-162.14\"/>\n",
"<polygon fill=\"black\" stroke=\"black\" points=\"148.64,-165.64 158.63,-162.14 148.63,-158.64 148.64,-165.64\"/>\n",
"</g>\n",
"<!-- FirstMatchState -->\n",
"<g id=\"node4\" class=\"node\">\n",
"<title>FirstMatchState</title>\n",
"<path fill=\"none\" stroke=\"black\" d=\"M527.75,-234.14C527.75,-234.14 439.75,-234.14 439.75,-234.14 433.75,-234.14 427.75,-228.14 427.75,-222.14 427.75,-222.14 427.75,-210.14 427.75,-210.14 427.75,-204.14 433.75,-198.14 439.75,-198.14 439.75,-198.14 527.75,-198.14 527.75,-198.14 533.75,-198.14 539.75,-204.14 539.75,-210.14 539.75,-210.14 539.75,-222.14 539.75,-222.14 539.75,-228.14 533.75,-234.14 527.75,-234.14\"/>\n",
"<text text-anchor=\"middle\" x=\"483.75\" y=\"-212.44\" font-family=\"Arial\" font-size=\"14.00\">FirstMatchState</text>\n",
"</g>\n",
"<!-- ChoiceState&#45;&gt;FirstMatchState -->\n",
"<g id=\"edge3\" class=\"edge\">\n",
"<title>ChoiceState&#45;&gt;FirstMatchState</title>\n",
"<path fill=\"none\" stroke=\"black\" d=\"M283,-171.24C319.85,-179.53 374.88,-191.9 417.86,-201.56\"/>\n",
"<polygon fill=\"black\" stroke=\"black\" points=\"417.15,-204.99 427.67,-203.76 418.68,-198.16 417.15,-204.99\"/>\n",
"<text text-anchor=\"middle\" x=\"371.25\" y=\"-199.94\" font-family=\"Times-Roman\" font-size=\"14.00\">Condition</text>\n",
"</g>\n",
"<!-- SecondMatchState -->\n",
"<g id=\"node5\" class=\"node\">\n",
"<title>SecondMatchState</title>\n",
"<path fill=\"none\" stroke=\"black\" d=\"M537.25,-180.14C537.25,-180.14 430.25,-180.14 430.25,-180.14 424.25,-180.14 418.25,-174.14 418.25,-168.14 418.25,-168.14 418.25,-156.14 418.25,-156.14 418.25,-150.14 424.25,-144.14 430.25,-144.14 430.25,-144.14 537.25,-144.14 537.25,-144.14 543.25,-144.14 549.25,-150.14 549.25,-156.14 549.25,-156.14 549.25,-168.14 549.25,-168.14 549.25,-174.14 543.25,-180.14 537.25,-180.14\"/>\n",
"<text text-anchor=\"middle\" x=\"483.75\" y=\"-158.44\" font-family=\"Arial\" font-size=\"14.00\">SecondMatchState</text>\n",
"</g>\n",
"<!-- ChoiceState&#45;&gt;SecondMatchState -->\n",
"<g id=\"edge4\" class=\"edge\">\n",
"<title>ChoiceState&#45;&gt;SecondMatchState</title>\n",
"<path fill=\"none\" stroke=\"black\" d=\"M324.52,-162.14C351.52,-162.14 381.39,-162.14 407.83,-162.14\"/>\n",
"<polygon fill=\"black\" stroke=\"black\" points=\"408.16,-165.64 418.16,-162.14 408.16,-158.64 408.16,-165.64\"/>\n",
"<text text-anchor=\"middle\" x=\"371.25\" y=\"-165.94\" font-family=\"Times-Roman\" font-size=\"14.00\">Condition</text>\n",
"</g>\n",
"<!-- DefaultState -->\n",
"<g id=\"node6\" class=\"node\">\n",
"<title>DefaultState</title>\n",
"<ellipse fill=\"none\" stroke=\"#ff9999\" stroke-width=\"2\" cx=\"483.75\" cy=\"-63.14\" rx=\"59.28\" ry=\"59.28\"/>\n",
"<ellipse fill=\"none\" stroke=\"#ff9999\" stroke-width=\"2\" cx=\"483.75\" cy=\"-63.14\" rx=\"63.29\" ry=\"63.29\"/>\n",
"<text text-anchor=\"middle\" x=\"483.75\" y=\"-59.44\" font-family=\"Arial\" font-size=\"14.00\">DefaultState</text>\n",
"</g>\n",
"<!-- ChoiceState&#45;&gt;DefaultState -->\n",
"<g id=\"edge5\" class=\"edge\">\n",
"<title>ChoiceState&#45;&gt;DefaultState</title>\n",
"<path fill=\"none\" stroke=\"black\" d=\"M271.02,-150.4C306.45,-135.8 368.03,-110.42 415.69,-90.78\"/>\n",
"<polygon fill=\"black\" stroke=\"black\" points=\"417.11,-93.98 425.02,-86.94 414.44,-87.51 417.11,-93.98\"/>\n",
"<text text-anchor=\"middle\" x=\"371.25\" y=\"-122.94\" font-family=\"Times-Roman\" font-size=\"14.00\">Default</text>\n",
"</g>\n",
"<!-- NextState -->\n",
"<g id=\"node7\" class=\"node\">\n",
"<title>NextState</title>\n",
"<path fill=\"none\" stroke=\"black\" d=\"M654.25,-206.14C654.25,-206.14 602.25,-206.14 602.25,-206.14 596.25,-206.14 590.25,-200.14 590.25,-194.14 590.25,-194.14 590.25,-182.14 590.25,-182.14 590.25,-176.14 596.25,-170.14 602.25,-170.14 602.25,-170.14 654.25,-170.14 654.25,-170.14 660.25,-170.14 666.25,-176.14 666.25,-182.14 666.25,-182.14 666.25,-194.14 666.25,-194.14 666.25,-200.14 660.25,-206.14 654.25,-206.14\"/>\n",
"<path fill=\"none\" stroke=\"black\" d=\"M658.25,-210.14C658.25,-210.14 598.25,-210.14 598.25,-210.14 592.25,-210.14 586.25,-204.14 586.25,-198.14 586.25,-198.14 586.25,-178.14 586.25,-178.14 586.25,-172.14 592.25,-166.14 598.25,-166.14 598.25,-166.14 658.25,-166.14 658.25,-166.14 664.25,-166.14 670.25,-172.14 670.25,-178.14 670.25,-178.14 670.25,-198.14 670.25,-198.14 670.25,-204.14 664.25,-210.14 658.25,-210.14\"/>\n",
"<text text-anchor=\"middle\" x=\"628.25\" y=\"-184.44\" font-family=\"Arial\" font-size=\"14.00\">NextState</text>\n",
"</g>\n",
"<!-- FirstMatchState&#45;&gt;NextState -->\n",
"<g id=\"edge6\" class=\"edge\">\n",
"<title>FirstMatchState&#45;&gt;NextState</title>\n",
"<path fill=\"none\" stroke=\"black\" d=\"M539.84,-205.32C551.72,-202.99 564.22,-200.53 575.94,-198.23\"/>\n",
"<polygon fill=\"black\" stroke=\"black\" points=\"576.87,-201.61 586,-196.25 575.52,-194.74 576.87,-201.61\"/>\n",
"</g>\n",
"<!-- SecondMatchState&#45;&gt;NextState -->\n",
"<g id=\"edge7\" class=\"edge\">\n",
"<title>SecondMatchState&#45;&gt;NextState</title>\n",
"<path fill=\"none\" stroke=\"black\" d=\"M549.33,-173.93C558.33,-175.57 567.47,-177.24 576.17,-178.83\"/>\n",
"<polygon fill=\"black\" stroke=\"black\" points=\"575.69,-182.29 586.15,-180.65 576.94,-175.41 575.69,-182.29\"/>\n",
"</g>\n",
"</g>\n",
"</svg>\n"
],
"text/plain": [
"<graphviz.sources.Source at 0x7f9534308f20>"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"dot_output = stepfunction_to_graphviz(aws_state_machine_json_sample)\n",
"graph = graphviz.Source(dot_output)\n",
"display(graph)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "40508259-5bbd-44c9-aa2d-2c3352a9f8dd",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.7"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment