aws-samples / aws-samples/sagemaker-ssh-helper
[Feature] SageMaker job as Studio kernel
- Dominant language
- Python
- Stars
- 261
- Forks
- 34
- PR merge metrics
- No merged PRs in 30d
Description
Lately I work mainly in SageMaker Studio, and I'd really like to be able to debug / interact with a running job using the same UI.
## Solution idea
Create a [custom Studio kernel image](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-byoi-create.html) using an IPython [extension](https://ipython.readthedocs.io/en/stable/config/extensions/index.html) and/or [custom magic](https://ipython.readthedocs.io/en/stable/config/custommagics.html) through which users can connect to a running SSH Helper job and run notebook cells on that instead of the Studio app.
The user experience would be something like [using EMR clusters in Studio](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-notebooks-emr-cluster-connect.html):
- One-time up-front job to build/register the custom "SageMakerSSH" image (maybe?)
- User launches their SSH-helper-enabled job from "normal" notebook A and fetches the managed instance ID `mi-1234567890abcdef0`
- User opens / switches to a notebook with SageMakerSSH kernel and runs something like
- `%load_ext sagemaker_ssh_helper.notebook` to initialize the IPython extension
- `%sagemaker_ssh connect mi-1234567890abcdef0` to connect to the instance
- From here on out, cells should run on the connected instance rather than the local Studio app unless a `%%local` cell magic is used: Same as how [SageMaker Studio SparkMagic kernel works](https://github.com/aws-samples/sagemaker-studio-emr/blob/3a19797e49c6e112767152d58db9e2e1f1dbe997/interactive_notebooks/sparkmagic-example.ipynb)
- Probably some kind of `%sagemaker_ssh disconnect` command would also be useful
Since the `sagemaker_ssh_helper` library is pip-installable, it *might* even be possible to get this working with default (e.g. `Data Science 3.0`) kernels? I'm not sure - assume it depends how much hacking is possible during IPython extension load vs what needs setting up in advance.
## Why this route
To my knowledge, JupyterLab is a bit more fragmented in support for remote kernels than IDEs like VSCode/PyCharm/etc. It seems like there **are [ways](https://stackoverflow.com/questions/29037211/how-do-i-add-a-kernel-on-a-remote-machine-in-ipython-jupyter-notebook)** to set up SSH kernels, but it's also a tricky topic to navigate because so many pages online are talking about *"accessing your remotely-running Jupyter server"* instead. Investigating the [Jupyter standard kernel spec paths](https://jupyter-client.readthedocs.io/en/stable/kernels.html#kernel-specs), I see `/opt/conda/envs/studio/share/jupyter/kernels` exists but contains only a single `python3` kernel which doesn't appear in Studio UI. It looks like there's a custom `sagemaker_nb2kg` Python library that manages kernels, but no obvious integration points there for alternative kernel sources besides the studio "Apps" system - and sufficiently internal/complex that patching it seems like a bad idea.
...So it looks like directly registering the remote instance as a kernel in JupyterLab would be a non-starter.
If the magic-based approach works, it might also be possible to use with other existing kernel images (as mentioned above) and even inline in the same notebook after a training job is kicked off. Hopefully it would also enable toggling over to a new job/instance without having to run CLI commands to change the installed Jupyter kernels.
Contributor guide
Assessment
This issue has not been assessed yet.