This post is also available as an interactive notebook.
Consider the following problem: you’d like to enable users to automatically extract a function from a Jupyter notebook and publish it as a service. Actually serializing a closure from a function in a given environment is not difficult, given the
cloudpickle module. But merely constructing a serialized closure isn’t enough, since in general this function may require other modules to be available to run in another context.
Therefore, we need some way to identify the modules required by a function (and, ultimately, the packages that provide these modules). Since engineering time is limited, it’s probably better to have an optimistic-but-incomplete (or unsound) estimate and allow users to override it (by supplying additional dependencies when they publish a function) than it is to have a sound and conservative module list.1
modulefinder module in the Python standard library might initially seem like an attractive option, but it is unsuitable because it operates at the level of scripts. In order to use
modulefinder on a single function from a notebook, we’d either have an imprecise module list (due to running the whole notebook) or we’d need to essentially duplicate a lot of its effort in order to slice backwards from the function invocation so we could extract a suitably pruned script.
Fortunately, you can interrogate nearly any property of any object in a Python program, including functions. If we could inspect the captured variables in a closure, we could identify the ones that are functions and figure out which modules they were declared in. That would look something like this:
inspect module provides a friendly interface to inspecting object metadata. In the above function, we’re constructing a worklist of all of the captured variables in a given function’s closure. We’re then constructing a set of all of the modules directly or transitively referred to by those captured variables, whether these are modules referred to directly, modules declaring functions referred to by captured variables, or modules declaring other values referred to by captured variables (e.g., native functions). Note that we add any functions we find to the worklist (although we don’t handle
eval or other techniques), so we’ll capture at least some of the transitive call graph in this case.
This approach seems to work pretty sensibly on simple examples:
It also works on itself, which is a relief:
While these initial experiments are promising, we shouldn’t expect that a simple approach will cover everything we might want to do. Let’s look at a (slightly) more involved example to see if it breaks down.
We’ll use the k-means clustering implementation from scikit-learn to optimize some cluster centers in a model object. We’ll then capture that model object in a closure and analyze it to see what we might need to import to run it elsewhere.
So far, so good. Let’s say we want to publish this simple model as a lighter-weight service (without a scikit-learn dependency). We can get that by reimplementing the
predict method from the k-means model:
What do we get if we analyze the second method?
This is a problem! We’d expect that
norm would be a captured variable in the body of
km_predict_two (and thus that
numpy.linalg would be listed in its module frontier), but that isn’t the case. We can inspect the closure variables:
We can see the cluster centers as well as the
min function and the
enumerate type. But
norm isn’t in the list. Let’s dive deeper. We can use the
dis module (and some functionality that was introduced in Python 3.4) to inspect the Python bytecode for a given function:
Ah ha! The body of our list comprehension, which contains the call to
norm, is a separate code object that has been stored in a constant. Let’s look at the constants for our function:
We can see the code object in the constant list and use
dis to disassemble it as well:
Once we’ve done so, we can see that the list comprehension has loaded
norm from a global, which we can then resolve and inspect:
Nested functions and lambda expressions
We can see a similar problem if we look at a function with local definitions (note that there is no need for the nesting in this example other than to expose a limitation of our technique):
In this case, we can see that Python compiles these nested functions in essentially the same way it compiles the bodies of list comprehensions:
But we can inspect the nested function just as we did the list comprehension. Let’s do that but just look at the load instructions; we’ll see that we’ve loaded
norm as we’d expect.
Predictably, we can also see a similar problem if we analyze a function with lambda expressions:
Let’s look at what happens when we import modules inside the function we’re analyzing. Because of the semantics of
import in Python, the module dependency list for this function will depend on whether or not we’ve already imported
sys in the global namespace. If we have, we’ll get a reasonable module list; if we haven’t, we’ll get an empty module list. (If you’re running this code in a notebook, you can try it out by restarting the kernel, re-executing the cell with the definition of
module_frontier, and then executing this cell.)
We can see this more clearly by importing a module in a function’s scope that we haven’t imported into the global namespace:
json is an unbound variable (since it isn’t bound in the enclosing environment of the closure). If it were bound in the global namespace, however, the
json we’re referring to in
example_six would be captured as a global variable:
Obviously, we’d like to return the same module-dependency results for functions that import modules locally independently of whether those modules have been imported into the global namespace. We can look at the bytecode for this function to see what instructions might be relevant:
To address the cases that the closure-inspecting approach misses, we can inspect the bytecode of each function. (We could also inspect abstract syntax trees, using the
ast module, but in general it’s easier to do this sort of work with a lower-level, more regular representation. ASTs have more cases to treat than bytecode.)
Python bytecode is stack-based, meaning that each instruction may take one or more arguments from the stack (in addition to explicit arguments encoded in the instruction). For a more involved analysis, we’d probably want to convert Python bytecode to a representation with explicit operands (like three-address code; see section 2 of Vallée-Rai et al. for a reasonable approach), but let’s see how far we can get by just operating on bytecode.
Identifying interesting bytecodes
We know from the problem cases we examined earlier that we need to worry about a few different kinds of bytecode instructions to find some of the modules that our
inspect-based approach missed:
LOAD_CONSTinstructions that load code objects (e.g., list comprehension bodies, lambda expressions, or nested functions);
LOAD_instructions that might do the same; and
IMPORT_NAMEinstructions that import a module into a function’s namespace.
Let’s extend our technique to also inspect relevant bytecodes. We’ll see how far we can get by just looking at bytecodes in isolation (without modeling the stack or value flow). First, we’ll identify the “interesting” bytecodes and return the modules, functions, or code blocks that they implicate:
Now we can make a revised version of our
module_frontier function. This starts with the same basic approach as the initial function but it also:
- processes the bytecode for each code block transitively referred to in each function,
- processes any modules explicitly imported in code.
As you can see, this new approach produces sensible results for all of our examples, including the ones that had confounded the closure-variable approach.
This is ongoing work and I hope to cover refinements to and extensions of this technique in future posts; as I mentioned at the beginning of the post, the ultimate goal is a tool to publish functions in notebook cells as self-contained services. It has been fun to learn about the power (and limitations) of the
inspect module – as well as a little more about how Python compiles code blocks and nested functions.
Thanks to my friend and colleague Erik Erlandson for suggesting improvements to the presentation of this post.
Sound program analyses present conservative overapproximations of program behavior. Consider a may-alias analysis, which determines if two reference variables may refer to the same location in memory. Precise may-alias analysis is undecidable, but certain kinds of imprecision are acceptable. Often we’re interested in sound analyses to support verification or semantics-preserving program transformations, so false positives are acceptable but false negatives are not. Put another way, the worst that can come of spuriously identifying a pair of variables as potentially-aliasing is that we’d miss an opportunity to optimize our program; the worst that can come of not identifying a pair of potentially-aliasing variables as such is a program tranformation that introduces a behavior change. By contrast, unsound analyses are imprecise but not conservative: both false positives and false negatives are possible. These analyses can still be useful for program understanding (e.g., in linters or static bug detectors) even if they are not sufficient to support safe program transformations. ↩