Failed to import custom package in managed notebook

Problem

When you schedule a notebook file with the executor, and you have some custom packages installed in user space, you could come across failed import error from the execution results like below: 

import apache_beam
dir(apache_beam)
---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
/tmp/ipykernel_162/2335761653.py in <module>
----> 1 import apache_beam
      2 dir(apache_beam)

ModuleNotFoundError: No module named 'apache_beam'

Environment

  • Vertex AI Workbench
  • Notebook files with the executor
  • Custom Python packages

Solution

  1. Remove --user option by using: 
    pip install apache-beam[gcp]
    Instead of: 
    pip install apache-beam[gcp] --user

Cause

When you write notebook code to run in the executor, the code will run in a tenant project separate from your managed notebooks instance's project. Hence, installing python package to user directory causes import path issues from both custom container and installing the package within your notebook file's code.