A submit on making a customized IPython Jupyter Magic command

Jupyter Notebooks are commonplace in knowledge science. They permit a mix of “repeat, analysis, loop” (REPL) code writing and documentation in a single place. They’re mostly used for evaluation and brainstorming functions, but additionally, extra contentiously, some choose notebooks to scripts to run manufacturing code (however we gained’t deal with that right here).

Invariably, the code written in notebooks can be repetitive ultimately akin to establishing a database connection, displaying an output, saving outcomes, interacting with an inside platform instrument, and so on. It’s finest to retailer this code as capabilities and/or modules to make them reusable and extra simply handle them.

Nonetheless, the pocket book expertise just isn’t all the time improved once you do this. As an illustration, you continue to must import and name these capabilities all through your pocket book, which doesn’t change the pocket book expertise a lot in any respect. So what’s the reply to augmenting the pocket book improvement expertise itself? IPython Jupyter Magic Instructions.

IPython Jupyter Magic instructions (e.g. strains in pocket book cells beginning with % or %%) can enhance a pocket book cell, or line, to switch its conduct. Many can be found by default, together with %timeit to measure the execution time of the cell and %bash to execute shell instructions, and others are supplied by extensions akin to %sql to write down SQL queries straight in a cell of your pocket book.

On this submit, we’ll present how your group can flip any utility operate(s) into reusable IPython Jupyter magics for a greater pocket book expertise. For example, we’ll use Hamilton, an open supply library we created, to encourage the creation of a magic that facilitates higher improvement ergonomics for utilizing it. You needn’t know what Hamilton is to know this submit.

Be aware. These days, there are various flavors of notebooks (Jupyter, VSCode, Databricks, and so on.), however they’re all constructed on high of IPython. Due to this fact, the Magics developed ought to be reusable throughout environments.

IPython Jupyter Magics (which we’ll shorten to simply Magics) are bits of code that may be dynamically loaded into your notebooks. They arrive in two flavors, line and cell magics.

Line magic, because it suggests, operates on a single line. That’s, it solely takes in as enter what’s specified on the identical line. They’re denoted by a single % in entrance of the command.

# will solely time the primary line
%time print("whats up")
print("world")

Cell magic, because it suggests, takes all the contents of a cell. They’re denoted by a double `%%`in entrance of the command.

# will time all the cell
%%timeit
print("whats up")
print("world")

Jupyter comes with a number of in-built magic instructions. You possibly can consider them as “command line” instruments which have entry to all the pocket book context. This enables them to work together with the pocket book output (e.g., printing outcomes, displaying a PNG, rendering HTML), but additionally to switch the state of present variables and write to different code and markdown cells!

That is nice for creating inside tooling as a result of it will possibly summary away and conceal from the consumer pointless complexities, making the expertise “magical”. It is a highly effective instrument to develop your individual “platform efforts”, particularly for MLOps and LLMOps functions, as you possibly can conceal what’s being built-in with from having to be uncovered within the pocket book. It subsequently additionally implies that notebooks don’t must be up to date if this abstracted code modifications under-the-hood, since it will possibly all be hidden in a python dependency improve.

Magic instructions have the potential to make your workflow less complicated and sooner. For instance, in the event you choose to develop in a pocket book earlier than shifting your code to a Python module, this will contain error-prone slicing & pasting. For this objective, the magic %%writefile my_module.py will straight create a file and replica your cell content material into it.

On the alternative hand, you would possibly choose creating in my_module.py in your IDE after which load it right into a pocket book to check your capabilities. This often includes restarting the pocket book kernel to refresh the module imports, which may be tedious. In that case, %autoreload will routinely reload each module imports earlier than each cell execution, eradicating this friction level!

Within the submit How well-structured ought to your knowledge code be?, it’s argued that standardization/centralization/“platform” efforts, ought to change the form of the “transfer shortly vs. built-to-last” trade-off curve for the higher. A concrete tactic to alter this trade-off is to implement higher tooling. Higher tooling ought to make what was once complicated, less complicated and accessible. Which is precisely what you possibly can obtain with your individual customized Magic instructions, which interprets to much less of a trade-off to be made.

For these unfamiliar with Hamilton we level readers to the numerous TDS articles with it (e.g. origin story, Manufacturing immediate engineering, Simplifying Airflow DAG Creation & Upkeep, Tidy Manufacturing Pandas, and so on.) in addition to https://www.tryhamilton.dev/.

Hamilton is an open supply instrument that we created at Sew Repair again in 2019. Hamilton helps knowledge scientists and engineers outline testable, modular, self-documenting dataflows, that encode lineage and metadata. Hamilton achieves these traits partly by requiring Python capabilities to be curated into modules.

Nonetheless, the everyday Jupyter pocket book utilization sample results in code residing within the pocket book and nowhere else posing a developer ergonomics problem:

How can we allow somebody to create Python modules simply and shortly from a pocket book, whereas additionally bettering the event expertise?

The Hamilton developer loop seems like the next:

Take a minute to learn this loop. The loop reveals that anytime a code change is made, the consumer would wish to not solely re-import the Python module, but additionally recreate the Driver object as effectively. Since notebooks permit cell execution in any order, it will possibly turn out to be tough for the consumer to trace which model is loaded for every module and what’s at the moment loaded in a Driver. This burden lies on the consumer and would possibly require restarting the kernel, which might lose different computations (fortunately, Hamilton may be set as much as execute complicated dataflows and resume the place you left-off…), which is lower than ultimate.

Right here’s how we might enhance this loop utilizing Magics:

1. Create a “momentary” Python module from the capabilities outlined in a cell, and import this new module straight within the pocket book.
2. Robotically visualize the directed acyclic graph (DAG) outlined by the capabilities to cut back visualization boilerplate code.
3. Rebuild all Hamilton Drivers discovered within the pocket book with up to date modules, saving the consumer time to have to recollect to manually recreate drivers to choose up the change.

We want a command that appears like this:

%%cell_to_module -m my_module --display --rebuild-drivers
def my_func(some_input: str) -> str:
"""Some logic"""
return ...

And trigger the next conduct after working the cell:

• Create a module with the title my_module within the pocket book.
• Show the DAG constructed by the capabilities inside the cell.
• Rebuild any downstream drivers that used my_module in different cells, saving the consumer having to re-run these cells.

As you possibly can see it is a non-trivial Magic command, since we’re adjusting output of the cell and state of the pocket book.

Right here we lay out step-by-step find out how to create a Magic Command. To keep away from solely exhibiting a trivial “whats up world” instance, we’ll clarify how we constructed Hamilton’s %%cell_to_module magic as effectively.

Create a brand new Python module the place we’ll write the magic code and a jupyter pocket book to attempt it. The title of this module (i.e., `.py` file) would be the title of the extension you’ll must load.

If Jupyter pocket book is put in, you may have all of the required Python dependencies already. Then, add libraries you’ll need in your Magic, in our case Hamilton (`pip set up sf-hamilton[visualization]`).

To outline a easy Magic command, you should use capabilities or objects (see these docs). For extra complicated Magics the place state is required, you’ll need the category method. We’ll use the category primarily based method right here. To begin we have to import IPython modules/capabilities after which outline a category that inherits magic.Magics. Every methodology adorned with @cell_magic or @line_magic defines a brand new magic, and the category can home nonetheless lots of them.

To begin, your code ought to appear to be this at a excessive stage:

# my_magic.py
from IPython.core import magic
from IPython.core.magic_arguments import argument, magic_arguments, parse_argstring
@magic.magics_class
class MyMagic(magic.Magics):
"""Customized class you write"""
@magic_arguments() # must be on high to allow parsing
@argument(...)   
@magic.cell_magic
def a_cell_magic_command(self, line, cell):
...
@magic_arguments() # must be on high to allow parsing
@argument(...)   
@magic.line_magic
def a_line_magic_command(self, line):
...

For stateful magic, it may be helpful so as to add an __init__() methodology (i.e. constructor). It isn’t wanted in our case.

By inheriting from `magic.Magics`, this class has entry to a number of essential fields together with self.shell, which is the IPython InteractiveShell that underlies the pocket book. Utilizing it lets you pull and introspect variables loaded within the energetic Jupyter pocket book.

Our Hamilton Magic Command will begin off trying like:

from IPython.core import magic
from IPython.core.magic_arguments import argument, magic_arguments, parse_argstring
@magic.magics_class
class HamiltonMagics(magic.Magics):
"""Magics to facilitate Hamilton improvement in Jupyter notebooks"""
@magic_arguments()  # wanted on high to allow parsing
@arguments(...)
@magics.cell_magic
def cell_to_module(self, line, cell):
...

Subsequent, we specify what arguments can be handed and find out how to parse them. For every argument, add a @argument, and add a @magic_arguments() decorator on high. They observe an analogous sample to argparse arguments in the event you’re acquainted, however they aren’t fairly as absolutely featured. Throughout the operate, you might want to name the parse_argstring() operate. It receives the operate itself to learn directions from the decorators, and `line` (the one with % or %%) which accommodates the arguments values.

Our command would begin to appear to be this:

@magic_arguments() # must be on high to allow parsing
# flag, lengthy kind, default worth, assist string.
@argument("-a", "--argument", default="some_value", assist="Some non-compulsory line argument")   
@magic.cell_magic
def a_cell_magic_command(self, line, cell):
args = parse_argstring(self.a_cell_magic_command, line)
if args.argument:
# do stuff -- place your utility capabilities right here

Be aware, for required arguments, there isn’t a facility in magic_arguments() for that, so you might want to manually examine within the physique of the operate for correctness, and so on.

Persevering with our dissection of the Hamilton Magic instance, the strategy on the category now seems like the next; we use many non-compulsory arguments:

@magic_arguments()  # wanted on high to allow parsing
@argument(
"-m", "--module_name", assist="Module title to offer. Default is jupyter_module."
)  # key phrase / non-compulsory arg
@argument(
"-c", "--config", assist="JSON config, or variable title containing config to make use of."
)  # key phrase / non-compulsory arg
@argument(
"-r", "--rebuild-drivers", motion="store_true", assist="Flag to rebuild drivers"
)  # Flag / non-compulsory arg
@argument(
"-d", "--display", motion="store_true", assist="Flag to visualise dataflow."
)  # Flag / non-compulsory arg
@argument(
"-v", "--verbosity", kind=int, default=1, assist="0 to cover. 1 is regular, default"
)  # key phrase / non-compulsory arg
@magics.cell_magic
def cell_to_module(self, line, cell):
"""Execute the cell and dynamically create a Python module from its content material.
A Hamilton Driver is routinely instantiated with that module for variable `{MODULE_NAME}_dr`.
> %%cell_to_module -m MODULE_NAME --display --rebuild-drivers
Sort in ?%%cell_to_module to see the arugments to this magic.
"""
# specify find out how to parse by passing 
args = parse_argstring(self.cell_to_module, line)  
# now use args for logic ...

Be aware, the additional arguments to @argument are helpful for when somebody makes use of ? to question what the magic does. I.e. ?%%cell_to_module will present documentation.

Now that we now have parsed the arguments, we will implement the logic of the magic command. There aren’t any specific constraints right here and you may write any Python code. Skipping a generic instance (you may have sufficient to get began from the prior step), let’s dig into our Hamilton Magic instance. For it, we wish to use the arguments to find out the specified conduct for the command:

1. Create the Python module with module_name.
2. If — rebuild-driver, rebuild the drivers, passing in verbosity.
3. If — config is current, get that prepared.
4. If — show, show the DAG.

See feedback within the code for explanations:

# we're within the bowels of def cell_to_module(self, line, cell):
# and we take away an indentation for readability
...
# specify find out how to parse by passing this methodology to the operate 
args = parse_argstring(self.cell_to_module, line)
# we set a default worth, else use the handed in worth 
# for the module title.
if args.module_name is None:
module_name = "jupyter_module"
else:
module_name = args.module_name
# we decide whether or not the configuration is a variable
# within the pocket book setting
# or if it is a JSON string that must be parsed.
display_config = {}
if args.config:
if args.config in self.shell.user_ns:
display_config = self.shell.user_ns[args.config]
else:
if args.config.startswith("'") or args.config.startswith('"'):
# strip quotes if current
args.config = args.config[1:-1]
attempt:
display_config = json.masses(args.config)
besides json.JSONDecodeError:
print("Didn't parse config as JSON. "
"Please guarantee it is a legitimate JSON string:")
print(args.config)
# we create the python module (utilizing a customized operate)
module_object = create_module(cell, module_name)
# shell.push() assign a variable within the pocket book. 
# The dictionary keys are the variable title
self.shell.push({module_name: module_object})
# Be aware: self.shell.user_ns is a dict of all variables within the pocket book
# -- we go that down by way of self.shell.
if args.rebuild_drivers:
# rebuild drivers that use this module (customized operate)
rebuilt_drivers = rebuild_drivers(
self.shell, module_name, module_object, 
verbosity=args.verbosity
)
self.shell.user_ns.replace(rebuilt_drivers)
# create a driver to show issues for each cell with %%cell_to_module
dr = (
driver.Builder()
.with_modules(module_object)
.with_config(display_config)
.construct()
)
self.shell.push({f"{module_name}_dr": dr})
if args.show:
# return will go to the output cell. 
# To show a number of components, use IPython.show.show(
# print("whats up"), dr.display_all_functions(), ... )
return dr.display_all_functions()

Discover how we use self.shell. This enables us to replace and inject variables into the pocket book. The values returned by the operate can be used as “cell output” (the place you see printed values).

Lastly, we have to inform IPython and the pocket book concerning the Magic Command. Our module the place our Magic is outlined will need to have the next operate to register our Magic class, and have the ability to load our extension. If doing something stateful, that is the place you’d instantiate it.

Discover that the argument `ipython` right here is similar InteractiveShell out there by self.shell within the class methodology we outlined.

def load_ipython_extension(ipython: InteractiveShell):
"""
Any module file that outline a operate named `load_ipython_extension`
may be loaded by way of `%load_ext module.path` or be configured to be
autoloaded by IPython at startup time.
"""
ipython.register_magics(MyMagic)
ipython.register_magics(HamiltonMagics)

See the total Hamilton Magic Command right here.

To load your magic within the pocket book, attempt the next:

%load_ext my_magic

within the case of our Hamilton Magic we’d load it by way of:

%load_ext hamilton.plugins.jupyter_magic

Whilst you’re creating, use this to reload your up to date magic with out having to restart the pocket book kernel.

%reload_ext my_magic

You possibly can then invoke the magic instructions outlined on a per line or cell foundation. So for the Hamilton one we’d now have the ability to do:

%%?cell_to_module

Right here’s an instance use of it, with it injecting the visualization:

In an actual world use case, you’ll probably model and package deal your magic possible right into a library, you can then handle simply in python environments as required. With the Hamilton Magic Command it’s packaged into the Hamilton library, and thus to get it, one want solely set up sf-hamilton and loading the magic command would turn out to be accessible within the pocket book.

On this submit we confirmed you the steps required to create and cargo your individual IPython Jupyter Magic Command. Hopefully you’re now pondering of the frequent cells/job/actions that you just carry out in a pocket book setting, which might be enhanced/simplified/and even eliminated with the addition of a easy Magic!

To display a real-life instance, we motivated and confirmed the internals of a Hamilton Magic Command to point out a command that was constructed to enhance the developer expertise inside a Jupyter pocket book, by augmenting the output and altering inside state.

We hope that this submit helps you recover from the hump and create one thing extra ergonomic and helpful for you and your groups’ Jupyter Pocket book expertise.