Improvements to legate.jupyter

[bryevdv]

This PR fixes the legate.jupyter functionality and also make some improvements:

• stub typings added so that mypy can be employed

• driver code is re-used to:

  • ensure consistency with standard legate invocations
  • reduce code (semi-)duplication
  • take advantage of existing extensive unit tests

• config values for %legate_info are passed in kernel spec metadata, reducing filesystem operations

• kernel name is included in the kernelspec env, so that:

  • legate.jupyter can create multiple differently named kernel configurations
  • %legate_info can report on the exact running kernel

  

Questions

1. Currently the --legate-dir option is removed. Is this option needed? At present, legate.jupyter always uses the legate install for itself, which seems appropriate, especially now that installs always happen into python environments. This includes pointing the spec argv at a _legion_kernel.py module in the install, instead of actually copying files around.

2. Currently --json (renamed --config) is a no-op. Is it actually needed? Just running python -m legate.jupyter uses exactly the same defaults as legate. These can be overridden with the standard legate command line options. Do we need to allow configuration of kernels from a file as well? If so I would propose splitting into two sub-commmands, one for using a file and one for using cmdline args, to avoid complicated config merges.

3. I don't see any reason we couldn't easily support any other standard driver options (e.g. for profiling) easily now. Would it be useful for users to be able to install a debugging or profiling kernel alongside others?

4. Do we want a real entry point? e.g. users run legate-jupyter instead of python -m legate.jupyter

5. Currently have to explicitly run %load_ext legate.jupyter first before %legate_info was this avoided somehow before?

6. %legate_info could report the computed legion_python invocation. Would that be useful for debugging? Or perhaps %legate info and %legate driver ? or similar

Todo

Once questions are resolved I intend to implement any appropriate changes in this PR, and then add a unit test suite.

I would also like to push some improvements down into legion's version of this code, but I do think it will result in much less friction if legate's implementation remains self-contained and decoupled from legion (apart from the ultimate legion_python invocation)

Notes:

The changes under legate.driver are to facilitate increased code re-use, e.g the exact command line args can be re-used between different command line programs.

Updated code never needs to create/install a fake spec, so there is no code to remove specs. If users give an existing name they get actionable error message:

dev38 ❯ python -m legate.jupyter  --name leg3 -v       
ERROR: kernel spec 'leg3' already exists. Remove it by running 'jupyter kernelspec uninstall 'leg3', or choose a new kernel name.

Verbose mode has multiple levels of detail. -v gives basic info:

dev38 ❯ python -m legate.jupyter  --name leg4 -v
Wrote kernel spec file leg4/kernel.json

Jupyter kernel spec leg4 (leg4) has been installed

But --vv dumps the full spec that is installed

dev38 ❯ python -m legate.jupyter  --name leg5 -vv
Wrote kernel spec file leg5/kernel.json


{
  "argv": [
    "/home/bryan/work/legate.core/_skbuild/linux-x86_64-3.8/cmake-build/_deps/legion-build/bin/legion_python",
    "-ll:py",
    "1",
    "-lg:local",
    "0",
    "-ll:cpu",
    "4",
    "-ll:util",
    "2",
    "-ll:csize",
    "4000",
    "-level",
    "openmp=5,",
    "-lg:eager_alloc_percentage",
    "50",
    "/home/bryan/work/legate.core/legate/jupyter/_legion_kernel.py",
    "-f",
    "{connection_file}"
  ],
  "display_name": "leg5",
  "env": {
    "GASNET_MPI_THREAD": "MPI_THREAD_MULTIPLE",
    "LD_LIBRARY_PATH": "/home/bryan/work/legate.core/_skbuild/linux-x86_64-3.8/cmake-build/_deps/legion-build/lib:/home/bryan/work/legate.core/build/lib",
    "LEGATE_MAX_DIM": "4",
    "LEGATE_MAX_FIELDS": "256",
    "NCCL_LAUNCH_MODE": "PARALLEL",
    "PYTHONDONTWRITEBYTECODE": "1",
    "PYTHONPATH": "/home/bryan/work/legate.core/_skbuild/linux-x86_64-3.8/cmake-build/_deps/legion-src/bindings/python/build/lib:/home/bryan/work/legate.core/_skbuild/linux-x86_64-3.8/cmake-build/_deps/legion-src/jupyter_notebook:/home/bryan/work/legate.core/legate",
    "REALM_BACKTRACE": "1",
    "__LEGATE_JUPYTER_KERNEL_SPEC__": "leg5"
  },
  "interrupt_mode": "signal",
  "language": "python",
  "metadata": {
    "legate": {
      "argv": [
        "--name",
        "leg5",
        "-vv"
      ],
      "core": {
        "cpus": 4,
        "gpus": 0,
        "ompthreads": 4,
        "openmp": 0,
        "utility": 2
      },
      "memory": {
        "eager_alloc": 50,
        "fbmem": 4000,
        "numamem": 0,
        "regmem": 0,
        "sysmem": 4000,
        "zcmem": 32
      },
      "multi_node": {
        "launcher": "none",
        "launcher_extra": [],
        "nodes": 1,
        "not_control_replicable": false,
        "ranks_per_node": 1
      }
    }
  }
}

Jupyter kernel spec leg5 (leg5) has been installed

[eddy16112]

I just tried the new jupyter branch, it works for me. Nice work! Regarding the questions:

Currently the --legate-dir option is removed. Is this option needed?

I think we do not need it. I use it because I do not know the legate installation directory with the old makefile system.

This includes pointing the spec argv at a _legion_kernel.py module in the install, instead of actually copying files around.

Are we able to find this file in both editable and non-editable installation? I copied this file into kernel directory because I do not know how to find this file.

Currently --json (renamed --config) is a no-op. Is it actually needed?

If without --json, what if we run jupyter with non-default legate ? Just passing cmdline args? The reason I use json file is because Legion and Legate share the same jupyter support (install_jupyter.py), but they do not have the same cmdline args, so in the json file, people can specify the cmdline args and arg values. Apparently, we are not use Legion's install_jupyter.py anymore, so I am OK with removing the --json.

I don't see any reason we couldn't easily support any other standard driver options (e.g. for profiling) easily now. Would it be useful for users to be able to install a debugging or profiling kernel alongside others?

I am OK with installing a profiling kernel.

Do we want a real entry point? e.g. users run legate-jupyter instead of python -m legate.jupyter

I think python -m legate.jupyter may not work if we are not in the root directory of legate, so I use legate -m legate.jupyter before.

Currently have to explicitly run %load_ext legate.jupyter first before %legate_info was this avoided somehow before?

No, I do not see how to run %load_ext legate.jupyter automatically.

[bryevdv]

I think we do not need it. I use it because I do not know the legate installation directory with the old makefile system.

Brilliant, I will cross it off the list

Are we able to find this file in both editable and non-editable installation?

Thanks for this, I had been testing editable. It looks like some work may be needed for non-editable (still investigating)

If without --json, what if we run jupyter with non-default legate

The workflow I had imagined is that users would switch to whatever legate environment they want, and then install jupyter extensions (once each) and run jupyter, out of that environment. To me that keeps things mentally separated in a nice way, but I can believe there are other use cases or needs I am not aware of. I would definitely like feedback on this point especially cc @manopapad @ipdemes

I think python -m legate.jupyter may not work if we are not in the root directory of legate, so I use legate -m legate.jupyter before.

Hrmm you may be right, I am still seeing if I can make python -m work outside the repo. FWIW the README states to use that

Please install Legate, then run the following command to install the IPython
kernel:

python -m legate.jupyter --json=legate_jupyter.json

I'll make sure to update the README regardless.

EDIT: I actually can't get legate -m legate.jupyter in or out of the repo. Some more invesigation is needed. It might be that an entry point turns out to be the most robust option

I am OK with installing a profiling kernel.

👍 I will leave this for a future PR

No, I do not see how to run %load_ext legate.jupyter automatically.

OK great just double checking. I think if we published this as a separate package then there might be some conda mechanism, but we can save that until/unless we split the jupyter extension off in the future.

[bryevdv]

@eddy16112 I think there is currently an issue with legions -m support, do you agree? The tmp.py just prints sys.argv

test38 ❯ python tmp.py 
['tmp.py']

test38 ❯ legate tmp.py 
['tmp.py']

test38 ❯ python -m tmp   
['/home/bryan/work/tmp.py']

test38 ❯ legate -m tmp
WARNING: Disabling control replication for interactive run
No module named tmp

[manopapad]

Not much to add over Wei's comments.

%legate_info could report the computed legion_python invocation. Would that be useful for debugging? Or perhaps %legate info and %legate driver ? or similar

I don't see why not.

The workflow I had imagined is that users would switch to whatever legate environment they want, and then install jupyter extensions (once each) and run jupyter, out of that environment. To me that keeps things mentally separated in a nice way, but I can believe there are other use cases or needs I am not aware of. I would definitely like feedback on this point especially cc @manopapad @ipdemes

I think this makes sense, but I don't personally use jupyter very often.

[manopapad]

Took a look at the code as it is right now, didn't see any issues. I know we're still waiting on more changes.

[eddy16112]

@bryevdv Have you updated your Legion to the latest control replication branch? There was a bug in Legion python binding regarding calling -m, but I recently fixed it. https://gitlab.com/StanfordLegion/legion/-/commits/master/bindings/python/legion_top.py

Here is my result:

(legate-jupyter) wwu@g0004:/scratch2/wwu/legate.jupyter$ python -m tmp
['/scratch2/wwu/legate.jupyter/tmp.py']
(legate-jupyter) wwu@g0004:/scratch2/wwu/legate.jupyter$ legate -m tmp
WARNING: Disabling control replication for interactive run
['/scratch2/wwu/legate.jupyter/tmp.py']

[bryevdv]

@eddy16112 @manopapad I have resolved the issues [1] by defining a real entry point, so the tool is used like:

legate-jupyter --name legate-kernel

and verified it works with non-editable installs. Does this seem acceptable? If so I will update docs and add tests.

[1] Boils down to legate -m legate being a problem. Maybe one we could fix, but an entry point is both simpler and nicer for users (IMO).

[bryevdv]

@manopapad apologies for the force push, I had been doing merges but somehow must have accidently done a rebase at some point. This is ready to go now.

[manopapad]

Beautiful! Any further comments @eddy16112 ?

[eddy16112]

LGTM, nice job!