aws-samples / aws-samples/sagemaker-ssh-helper

[Feature] SageMaker job as Studio kernel

Open
#15 2 comments 0 reactions 0 assignees View on GitHub
enhancement
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.