google / google/dpsynth

Lazy-load TensorFlow for TFRecord-specific pipeline paths

Open
#151 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
29
Forks
10
Avg merge
2d 2h
Merged PRs (30d)
36

Description

## Summary

This PR removes eager TensorFlow initialization from shared pipeline import paths.

Previously, importing non-TFRecord pipeline functionality could transitively import TensorFlow:

```text
dpsynth.data_generation
→ creating_data_recorder_converter
→ tfrecord_descriptor
→ tensorflow
```

As a result, workflows that do not use TFRecord, including the SWIFT pipeline test, still required TensorFlow to initialize successfully during module import.

This PR defers TFRecord-specific imports until `DataFormat.TFRECORD` is selected and imports TensorFlow only inside TFRecord-specific I/O paths.

## Changes

* Lazily import `tfrecord_descriptor` when `DataFormat.TFRECORD` is selected.
* Move TensorFlow imports into TFRecord-specific branches in pipeline I/O.
* Preserve TensorFlow type annotations using the existing `from __future__ import annotations` support and `TYPE_CHECKING` imports without requiring TensorFlow at module import time.
* Add regression coverage to verify that importing non-TFRecord pipeline modules does not load TensorFlow.
* Preserve existing public APIs and dependency extras.

## Why

Python executes module-level imports when importing a module. The previous import structure therefore initialized TensorFlow during test collection:

```text
pytest collection
→ import SWIFT test
→ import shared DPSynth modules
→ import TFRecord implementation
→ import TensorFlow
```

This happened before any SWIFT test executed.

The SWIFT workflow uses PipelineDP's `LocalBackend` and a dummy record converter and does not require TFRecord support. However, it could still be blocked if TensorFlow failed to initialize in the environment.

The change makes TensorFlow initialization conditional on actually entering a TFRecord code path:

```text
shared pipeline import
→ select data format
├── non-TFRecord → TensorFlow not imported
└── TFRECORD → load TFRecord implementation → import TensorFlow
```

This preserves TFRecord support while preventing unrelated pipeline workflows from depending on successful TensorFlow initialization.

## Verification

Confirmed that importing the shared modules no longer loads TensorFlow:

```bash
python -c "import sys; \
from dpsynth import data_generation; \
from dpsynth.pipeline_transformations import input_output; \
assert 'tensorflow' not in sys.modules"
```

The previously blocked pipeline tests now complete successfully:

```text
11 passed
```

Some unrelated JAX, Beam, `httplib2`, and Pyparsing warnings remain.

Contributor guide

Open the contributing guide

Research direction

Start with dpsynth.data_generation and dpsynth.pipeline_transformations.input_output, then inspect the tfrecord_descriptor import path and its TFRecord-specific I/O branches. Run the provided Python import check to confirm non-TFRecord modules leave TensorFlow out of sys.modules, then run the pipeline tests and verify the existing 11 tests pass.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, tensorflow
Domain
data-engineering, machine-learning
Issue type
Refactor
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.