(Feature Request) JupyterLab-compatible visualizations

[athewsey]

I've been a long-time advocate of JupyterLab over plain Jupyter, but this is becoming even more relevant with the growing & maturing of SageMaker Studio's features since 2019...

Is any work being done on making the graph renders compatible with JupyterLab?

I see some libraries e.g. ipydagred3 that seem to suggest potentially re-usable patterns, but not sure how much focus there'd need to be on retaining the Step Functions console look & feel for accepting visualizations into this SDK.

I'd really like to be able to visualize in-progress executions from my (these days mostly SMStudio-based) notebook workflows.

[wong-a]

We aren't actively working on JupyterLab support right now, but we could prioritize that depending on customer interest. So far we had three other customers ask about JupyterLab support. Are you using this SDK mostly in SageMaker Studio?
#29
#54

Ideally we want to avoid rewriting the graph rendering code. The actual graph construction is not actually part of this repo. It's sharing a JavaScript library with https://github.com/aws/aws-toolkit-vscode and rendering HTML in Jupyter.

https://github.com/aws/aws-step-functions-data-science-sdk-python/blob/main/src/stepfunctions/workflow/widgets/graph.py

[athewsey]

Yes, I'm mainly using SMStudio - but even if I was working locally I'd usually prefer JLab over Jupyter for the added features...

Understand about wanting to share rendering code, but perhaps this library could provide infrastructure to work around rendering in JupyterLab? E.g. using the existing infrastructure from IPython / IPyWidgets, or even creating a new labextension if necessary...

IPython.display.Javascript() and IPython.display.HTML() seem to both still available in SMStudio, so I think if we displayed the HTML template with a set element_id and then ran JS with that same reference, it might still be possible?

To my limited understanding, the problem is not so much with rendering custom JS at all, but that JS having access back to the Python kernel... Which maybe wouldn't be too much of an issue for these simple graph rendering activities?

I tried to hack around a bit to test, but couldn't resolve the JSLIB_URL of https://do0of8uwbahzz.cloudfront.net/sfn (no such key?) and couldn't figure out where your JS require global was coming from to try and construct whatever sub-URL the code might be fetching... Is this a Jupyter-provided thing?

[athewsey]

PoC Update: This is at least partially possible - see screenshot below of a Workflow graph in SageMaker Studio which also supports zooming so I think the interactivity is working too:

(I was using the light UI theme anyway, and of course that looks better with the theming)

Steps I took were:

• Fetch the JS as text:

JS = requests.get(JSLIB_URL + ".js").content.decode("utf-8")

• Customize the HTML template to exclude the script tag, and the code snippet template to use global sfn rather than expecting an argument (as a hack, because of the way we'll call the script later):

NOJS_HTML_TEMPLATE = """
<link rel="stylesheet" type="text/css" href="$css">
<div id="$element_id" class="workflowgraph">
    $graph_legend_template
    <svg></svg>
    {console_snippet}
</div>
"""
JLAB_WORKFLOW_GRAPH_SCRIPT_TEMPLATE = """
var element = document.getElementById('$element_id')
var options = {
    width: parseFloat(getComputedStyle(element, null).width.replace("px", "")),
    height: 600,
    layout: '$layout',
    resizeHeight: true
};
var definition = $definition;
var elementId = '#$element_id';
var graph = new sfn.StateMachineGraph(definition, elementId, options);
graph.render();
"""

• Now instead of trying to invoke the JS directly via an HTML() call with <script> tag in WorkflowGraphWidget.show(), first render the HTML template and then explicitly invoke the JS library and the script itself

display(HTML(self.template.substitute({
    'element_id': self.element_id,
    'definition': self.json_definition,
    'layout': self.layout,
    'css': CSS_URL,
    'jslib': JSLIB_URL,
    'graph_legend_template': ""
})))
display(Javascript(JS))
display(Javascript(Template(JLAB_WORKFLOW_GRAPH_SCRIPT_TEMPLATE).substitute({
    'element_id': self.element_id,
    'definition': self.json_definition,
    'layout': self.layout,
})))

Limitations

Of course the above is pretty hacky, breaks the nice return value API of show() at the moment, and runs the library + widget code in global scope for now - but I tentatively think there could be workarounds for these.

The main limitation AFAIK for this kind of IPython.display method vs trying to integrate more closely with IPyWidgets library, is that there's no way to pass data or events back from JS to the Python kernel.

I'm pretty new still to this library, but I think if you're already sharing the VSCode rendering engine you probably don't have any such interactions anyway right? Just a display tool, no feedback like selecting nodes and feeding ARNs to Python, etc.

So I think there could still be a relatively quick fix here to make the basic rendering work in JupyterLab/SMStudio, without significantly changing the architecture?