Skip to content

Thread safety for pipelines in KedroServiceSession #5655

Description

@ankatiyar

Description

Follow up from #5535
The pipeline loading mechanism at the moment is not thread safe for concurrent runs due to the recent optimisation added for selective pipeline loading.

Context

Full context here: #5535 (comment)

The central issue is _ProjectPipelines. It is a process-global singleton with a lazy-loading cache. Its job is to defer the import of user pipeline modules until they are first needed, which is valuable for CLI startup speed. But for a long-running server, this deferred-load mechanism becomes a liability.

_ProjectPipelines holds three pieces of mutable state:

  • _pipelines_module — which module to call register_pipelines() on
  • _is_data_loaded — whether the pipeline dict has been populated
  • _content — the populated dict of pipeline name → Pipeline object
  • _requested_pipelines — a filter hint for which pipelines to load
    KedroServiceSession.run() calls pipelines.set_requested(pipeline_names) at the start of every run. This is an optimisation: tell the loader to only import the modules for the requested pipelines, avoiding the cost of importing everything. set_requested() does this by invalidating the cache when the filter changes:
def set_requested(self, pipeline_names):
    if set(self._requested_pipelines or []) != set(pipeline_names or []):
        self._is_data_loaded = False   # invalidate
        self._content = {}             # clear
    self._requested_pipelines = list(pipeline_names) if pipeline_names else None

None of these three assignments is atomic relative to the others, and the entire operation is not protected by any lock. With two concurrent runs requesting different pipelines, the sequence becomes:


Thread 1: set_requested(["ingest"])      → _is_data_loaded=False, _content={}
Thread 2: set_requested(["train"])       → _is_data_loaded=False, _content={}, filter="train"
Thread 1: pipelines["ingest"]           → _load_data() fires with filter="train"
                                           → loads "train" pipeline only
                                           → Thread 1 gets KeyError for "ingest"

Secondary concern is the _hook_manager is created once and shared across every run.

Possible Implementation

This optimisation works for CLI based workflows but for serving use cases we could just preload all the pipelines.

  • mode=CLI/serving during session creation which allows for this optimisation when in CLI mode but for serving mode will load all pipelines

Possible Alternatives

  • Incrementally build pipelines snapshot:
run(["ingest"])  → cold miss → load ingest → snapshot = {ingest, __default__}
run(["ingest"])  → cache hit → no lock, no import
run(["train"])   → cold miss → load train  → snapshot = {ingest, train, __default__}
run(["train"])   → cache hit → no lock, no import
run(["ingest"])  → cache hit

Metadata

Metadata

Assignees

Labels

Issue: Feature RequestNew feature or improvement to existing feature

Type

Fields

No fields configured for Task.

Projects

Status
In Progress

Relationships

None yet

Development

No branches or pull requests

Issue actions