Jupyter Plugin for Pants v2.0.0

goals.py

# Copyright 2020 Chartbeat.
# Licensed under the Apache License, Version 2.0.

import logging

from pathlib import PurePath
from textwrap import dedent

from pants.backend.python.util_rules.pex import (
    Pex,
    PexRequest,
    PexRequirements,
)
from pants.backend.python.util_rules.pex_from_targets import PexFromTargetsRequest
from pants.backend.python.util_rules.python_sources import (
    PythonSourceFiles,
    PythonSourceFilesRequest,
)
from pants.base.build_root import BuildRoot
from pants.engine.addresses import Addresses
from pants.engine.fs import CreateDigest, Digest, FileContent, MergeDigests
from pants.engine.goal import Goal, GoalSubsystem
from pants.engine.process import InteractiveProcess, InteractiveRunner
from pants.engine.rules import Get, MultiGet, collect_rules, goal_rule
from pants.engine.target import (
    TransitiveTargets,
    TransitiveTargetsRequest,
)
from pants.option.global_options import GlobalOptions


# from jupyter_stubber-v1.0.0
JUPYTER_PEX_REQUIREMENTS_PY2_PY3 = [
    "ipykernel==4.5.2",
    "ipython==5.8.0",
    "jupyter_console==5.1.0",
    "nbconvert==5.6.1",
    "nbformat==4.4.0",
    "notebook==4.4.1",
    "pyzmq==15.3.0",
    "tornado<=4.1,>=1.1",
]

logger = logging.getLogger(__name__)


class JupyterSubsystem(GoalSubsystem):
    """Launch a Jupyter Notebook app (repl)."""

    name = "jupyter"

class Jupyter(Goal):
    subsystem_cls = JupyterSubsystem


@goal_rule
async def jupyter(
    all_specified_addresses: Addresses,
    build_root: BuildRoot,
    global_options: GlobalOptions,
    interactive_runner: InteractiveRunner,
) -> Jupyter:

    logger.warning("RUNNING JUPYTER NOTEBOOK: use ctrl-c to quit")
    logger.warning("KNOWN ISSUE: best to specify --no-pantsd when using (v2.0.0)")
    logger.warning("KNOWN ISSUE: quitting will not exit with status 0 (v2.0.0)")

    # log inputed targets
    for addr in all_specified_addresses:
        logger.debug(f"targets specified: {addr}")

    # Entry point needed to start Jupyter Notebook
    # TODO make dashboard configurable instead of hard coding it as the monorepo root
    jupyter_dashboard_homepage = build_root.path

    jupyter_launcher_file = dedent(
        """\
        import os
        import sys

        from notebook.notebookapp import main

        os.environ["PYTHONPATH"] = ":".join(sys.path)
        sys.exit(main())
    """
    ).encode()

    launcher_file = FileContent("--jupyter_launcher.py", jupyter_launcher_file)
    launcher_digest = await Get(Digest, CreateDigest([launcher_file]))

    # Preparations for creating requirements pex
    requirements_pex_request = await Get(
        PexRequest,
        PexFromTargetsRequest,
        PexFromTargetsRequest.for_requirements(
            all_specified_addresses, internal_only=True
        ),
    )
    requirements_pex_interpreter_constraints = (
        requirements_pex_request.interpreter_constraints
    )
    requirements_pex_get = Get(
        Pex,
        PexRequest,
        requirements_pex_request,
    )

    logger.debug(
        f"requirements_pex_get.input.output_filename: {requirements_pex_get.input.output_filename}"
    )

    # Preparations for creating jupyter pex
    jupyter_pex_setup = {
        "requirements": JUPYTER_PEX_REQUIREMENTS_PY2_PY3,
    }
    logger.info(f"jupyter_pex_setup: {jupyter_pex_setup}")  # XXX make debug

    jupyter_pex_get = Get(
        Pex,
        PexRequest(
            sources=launcher_digest,
            output_filename="jupyter-notebook.pex",
            requirements=PexRequirements(jupyter_pex_setup.get("requirements")),
            internal_only=True,
            interpreter_constraints=requirements_pex_interpreter_constraints,
            entry_point=PurePath(launcher_file.path).stem,
            additional_args=[
                # "--not-zip-safe",  # TODO determine if we need this. Causes performance hit
                "--pex-path",
                requirements_pex_get.input.output_filename,
            ],
        ),
    )

    transitive_targets = await Get(
        TransitiveTargets, TransitiveTargetsRequest(all_specified_addresses)
    )

    # Preparations for creating digest
    sources_request = Get(
        PythonSourceFiles,
        PythonSourceFilesRequest(transitive_targets.closure, include_files=True),
    )

    # Build Pexs and Digests
    jupyter_pex, requirements_pex, sources = await MultiGet(
        jupyter_pex_get, requirements_pex_get, sources_request
    )
    logger.debug(f"jupyter_pex: {jupyter_pex}")
    logger.debug(f"requirements_pex: {requirements_pex}")
    logger.info(f"sources: {sources}")  # XXX make debug

    merged_digest = await Get(
        Digest,
        MergeDigests(
            (
                jupyter_pex.digest,
                requirements_pex.digest,
                sources.source_files.snapshot.digest,
            )
        ),
    )

    # Add Extra Args and ENV
    jupyter_cmd_line = [jupyter_pex.name, jupyter_dashboard_homepage]
    logger.info(f"jupyter_cmd_line: {jupyter_cmd_line}")

    source_roots = ":".join(sources.source_roots)
    logger.debug(f"source_roots: {source_roots}")
    extra_env = {
        "PEX_EXTRA_SYS_PATH": source_roots,
    }

    # Run Interactive Process
    result = interactive_runner.run(
        InteractiveProcess(
            argv=jupyter_cmd_line,
            env=extra_env,
            input_digest=merged_digest,
            hermetic_env=False,
        )
    )

    # XXX Jupyter Exit Issue
    # Exiting from the above with ctrl-c causes the pantsd to exit with
    # exit_code 1.
    #
    # You will only get to this part if you run with `--no-pantsd`
    #
    # This might be resolved by using newer versions of notebook that can be
    # exited from the web GUI. Or it might be fixed by pantsd handling
    # interrupts differently.

    exit_code = result.exit_code
    logger.debug(f"{jupyter_pex.name} exit code is {exit_code}")
    return Jupyter(exit_code=exit_code)


def rules():
    return collect_rules()

@Eric-Arellano I've incorporated your recommended changes.