DLT pipelines - sporadic ModuleNotFoundError

jtrousdale-lyb
New Contributor III

When we run DLT pipelines (which we deploy via DABs), we get a sporadic issue when attempting to install our bundle's wheel file.

First, in every DLT pipeline, we as a first step a script that looks like the following

import subprocess as sp
from importlib.util import find_spec as importlib_find_spec
from pathlib import Path

from databricks.sdk.runtime import dbutils, spark

if not importlib_find_spec('my_package'):
    artifact_path = spark.conf.get('bundle.artifactPath')

    if not artifact_path:
        msg = 'Bundle artifact path not found in Spark configuration.'
        raise ValueError(msg)

    wheel_path = next(
        (Path(artifact_path) / '.internal').glob(
            'my_package-*-py3-none-any.whl',
        ),
    )
    sp.check_call(  # noqa: S603
        [  # noqa: S607
            'pip',
            'install',
            str(wheel_path),
        ],
    )

    dbutils.library.restartPython()

Where we pass bundle.artifactPath in from the bundle in the pipeline options.

Some large percentage of the time, this works perfectly. From this point, we can run code which has available both our package (here called "my_package") as well as the 3rd party dependencies that are implied onto the .whl file by our project dependencies in the pyproject.toml.

However, every now and then, during the DLT pipeline "Initializing" stage, this will fail, with an error like the following

pyspark.errors.exceptions.captured.AnalysisException: Traceback (most recent call last):
  File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-6584cd7b-1c47-4509-891d-918d38e18089/lib/python3.12/site-packages/my_package/base_pipeline_config.py", line 293, in _fcn
    return source.load()
           ^^^^^^^^^^^^^

pyspark.errors.exceptions.captured.AnalysisException: [PYTHON_DATA_SOURCE_ERROR] Failed to create Python data source instance: Traceback (most recent call last):
  File "/databricks/spark/python/pyspark/serializers.py", line 192, in _read_with_length
    return self.loads(obj)
           ^^^^^^^^^^^^^^^
  File "/databricks/spark/python/pyspark/serializers.py", line 617, in loads
    return cloudpickle.loads(obj, encoding=encoding)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ModuleNotFoundError: No module named 'my_package'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
pyspark.serializers.SerializationError: Caused by Traceback (most recent call last):
  File "/databricks/spark/python/pyspark/serializers.py", line 192, in _read_with_length
    return self.loads(obj)
           ^^^^^^^^^^^^^^^
  File "/databricks/spark/python/pyspark/serializers.py", line 617, in loads
    return cloudpickle.loads(obj, encoding=encoding)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ModuleNotFoundError: No module named 'my_package'

 SQLSTATE: 38000

Usually this ends up working after a re-run, though sometimes it can take several re-runs before it finally succeeds. And then sometimes it will work the first time several times in a row.

We'd really like to resolve this in a way that we can continue installing our wheel source code along with its external dependencies without this flakiness necessitating so many retries.

We thought about also trying to update sys.path, but this doesn't handle installing 3rd party dependencies. It seems like somehow the Python runtime gets out of sync.