workspace notebook UX

[d4l3k]

Description

We should add some notebook specific integrations to make working with workspace and launching remote jobs first class. This builds upon the workspace support tracked by #333.

Motivation/Background

Currently there's no specific TorchX integrations for running from within notebooks. It's possible but it's not as fleshed out as it could be.

Detailed Proposal

Jupyter Custom Magics

We want to add a custom magic to allow adding files to the workspace.

https://ipython.readthedocs.io/en/stable/config/custommagics.html#defining-custom-magics

from torchx.notebook import register_magics, get_workspace

register_magics()

%%workspacefile train.py

print("train Hello world!")

from torchx.components.utils import python
from torchx.runner import get_runner

app = python(m="train")
app_id = runner.run(app, scheduler="local_docker", workspace=get_workspace())
print(app_id)
status = runner.wait(app_id)
print(status)

Alternatives

This can already be accomplished by writing out a file and directly calling docker build etc. That's a lot more work on the user and requires having an on disk project so the notebook isn't fully self contained.

Additional context/links

Workspace/canary tracking #333