[BUG] LambdaOp may fail to deserialize if the processing function is declared in another module

[willb]

Describe the bug

By default, Cloudpickle serializes functions and modules by reference. Deserializing workflows that contain user functions declared in other modules can fail if the Python source files for those modules are not available. This means that serializing a workflow and publishing it in a Docker container requires either (1) that all user source files are published in the same paths in the container as they were when the workflow was serialized or (2) that user modules are not serialized by reference.

Steps/Code to reproduce bug

Assume that identity.py contains the identity function, called identity.

with open("identity.py", "w") as of:
    of.write("""
def identity(col):
    return col

""")

Then run this code:

import nvtabular as nvt
import identity

wf = nvt.Workflow(
    ["col_a"] >> nvt.ops.LambdaOp(identity.identity)
)

wf.save("identity-workflow")

This serialized workflow will fail to deserialize if identity.py is not available in the same location, e.g., in a Docker container or on another host.

import sys
import os

del sys.modules["identity"]
os.unlink("identity.py")

try:
    wf_ref = nvt.Workflow.load("identity-workflow")
except ModuleNotFoundError as mnfe:
    print("Failed to load workflow")
    print(str(mnfe))

Expected behavior

Serialized workflows should be largely self-contained and resilient to missing source files containing user transformation code.

Environment details (please complete the following information):

• Environment location: Docker
• Method of NVTabular install: Docker
  • If method of install is [Docker], provide docker pull & docker run commands used:

docker run -v /home/willb/devel/nvtabular-sandbox:/workspace/host --runtime=nvidia --rm -it -p 8888:8888 -p 8797:8787 -p 8796:8786 --ipc=host --cap-add SYS_NICE nvcr.io/nvidia/merlin/merlin-training:22.03 /bin/bash

Additional context

This notebook demonstrates the minimal reproducer as well as a workaround.

I expect to submit a patch shortly that will enhance Workflow.save in two ways:

1. to allow clients to provide explicit direction that certain modules should be serialized by value, and
2. to allow clients to request optional automatic identification of such modules.

[nv-alaiacano]

Closing because this was resolved.