transformWithState might be causing 'Module not found'

[AndriusVitkausk]

Hi,

So i've been playing around with transformWithState over the last few weeks and have got the code to work exactly as desired.

Wrote some integration tests that execute fine within the workspace in serverless and everything, the issue is however is once it's deployed and those same integration tests try to run from the Azure CI pipeline in serverless through databricks connect, I keep getting Module not found on the location where the StatefulProcessor instance is located.

Multiple other integration tests use the same directory and load everything fine across drivers and executors (use quite a bit of foreachBatch functionality), but it seems like this might be a different case. Been trying to provision the lib as a wheel, as a editable install but nothing is working so far. All modules are found in order to initialise, but I have a feeling once the workers try to deserialise the class, they don't respect the wheel attached (the job class and the stateful processor are both defined in the same .py file.

Anyone encountered this?

[DoTA]

This smells like a client/server module-resolution mismatch that's specific to Databricks Connect + serverless, rather than a wheel packaging problem per se.

A few things worth checking:

1. Where the StatefulProcessor class lives. Since it's defined in the same .py file as the job entrypoint, how that file gets executed matters a lot for cloudpickle. If your CI pipeline invokes it as a script or via pytest in a way that loads it as __main__ (or under a synthetic test-runner module name), cloudpickle will try to pickle the class by reference using that module name - and that name won't exist as an importable module on the remote serverless session, hence "module not found" once the workers try to deserialize it. When you run the same code interactively in a workspace notebook, Databricks' notebook execution model handles this differently, which is why it "just works" there.

Fix: move the StatefulProcessor into its own module inside the wheel's package (not co-located with the pipeline entrypoint / test file), so it has a stable, importable dotted path independent of how the entrypoint is invoked.

2. How the wheel actually gets onto the serverless session. "Attached" isn't the same as "installed in that Connect session's remote environment." For serverless + Databricks Connect, cluster-style library installs don't apply - you need to explicitly ship the dependency into the remote session, e.g. spark.addArtifact("<path-to-wheel>", pyfile=True) (or the equivalent environment/requirements spec for the serverless environment) before you build/run the streaming query in your CI job. If that step is missing or happens after the session is created, the driver can resolve the class locally (because of your editable install) but the remote executors can't.

3. Environment parity. Serverless is stricter about client/server Python & dependency version parity than classic clusters. Worth double-checking the databricks-connect client version and serverless environment version used in the Azure CI pipeline match what you use interactively.

If you can share how the CI job actually launches the tests (plain pytest? a wrapper script?) and how/where addArtifact (or equivalent) is called relative to session creation, that would help narrow down which of the two it is.