A pattern that shows up repeatedly when people move a working Structured Streaming pipeline from "runs fine in a workspace notebook" to "runs in CI via Databricks Connect against serverless" is a ModuleNotFoundError that only appears in the second context - and it usually isn't a packaging bug so much as a difference in how Python module resolution works across the two execution paths.

The symptom

Code using transformWithState (or any API that needs to serialize a custom class - a StatefulProcessor, a UDF class, etc.) works correctly when run interactively in a notebook attached to serverless compute. The exact same code, run from a CI pipeline through Databricks Connect against the same serverless environment, fails once the workers try to deserialize the class, with an error pointing at the module where that class is defined - even though the wheel containing it was built and "attached" to the job.

Why this happens

Two things differ between "interactive notebook" and "CI via Databricks Connect" that matter a lot for cloudpickle:

1. How the class gets pickled. cloudpickle can pickle a class two ways: by value (it embeds the class's actual bytecode/definition in the pickle) or by reference (it just records the module path and expects the remote side to import it). Which one happens depends on how the module containing the class was loaded. Notebook execution has its own synthetic module handling that often ends up pickling by value. A class defined in a plain .py file that's executed as a script, or loaded as __main__ by a test runner (pytest running a test file directly, a CI wrapper script, etc.), is far more likely to get pickled by reference using that script's module name - a name that has no meaning on the remote serverless session.

2. How the dependency actually reaches the remote session. "Attached" and "installed in the remote Connect session" are not the same thing. For classic clusters you'd rely on cluster-scoped library installs; for serverless + Databricks Connect there is no equivalent implicit step - you need to explicitly ship the dependency into the session, e.g. via spark.addArtifact(<wheel_path>, pyfile=True), or an environment/requirements spec for the serverless environment, before you build or run the query. If that call happens after the session/query is already created, or doesn't happen at all in the CI path (because the interactive path had the package installed locally via %pip install and never needed it), the driver can resolve the class fine locally while the remote executors can't.

The fix that actually holds up

- Move classes that need to survive a round-trip through cloudpickle (StatefulProcessor, custom UDF classes, etc.) into their own module inside the package that gets built into the wheel - not co-located with whatever script or test file is the actual entrypoint. That gives them a stable, importable dotted path regardless of how the entrypoint itself gets invoked.

- In the CI job, explicitly call spark.addArtifact(..., pyfile=True) (or the serverless environment/requirements equivalent) before constructing the streaming query, every time, rather than relying on whatever happened to already be on the local Python path.

- Keep the Databricks Connect client version and the serverless environment version in lockstep between what you use interactively and what CI uses - serverless is stricter about client/server parity than classic clusters, and a mismatch here can produce failures that look identical to the module-resolution issue above but have a different root cause.

None of this is unique to transformWithState specifically - it applies to anything that needs a custom class to survive serialization across the Connect boundary - but stateful streaming processors are where people run into it first, because they're usually the first custom class shape in a pipeline that has to be reconstructed on the executor at microbatch/checkpoint-recovery time rather than just called once on the driver.

If you've hit a variant of this that isn't covered by the two causes above, I'd be curious to hear it - the client/server split in Databricks Connect + serverless is still a fairly new mental model for a lot of people coming from classic clusters, and it wouldn't surprise me if there's a third failure mode I haven't run into yet.