Skip to content

feat: file-based runtime dispatch for injected tasks#266

Draft
lbakshi wants to merge 1 commit intoNike-Inc:mainfrom
lbakshi:fix/bake-injected-task-config-for-runtime
Draft

feat: file-based runtime dispatch for injected tasks#266
lbakshi wants to merge 1 commit intoNike-Inc:mainfrom
lbakshi:fix/bake-injected-task-config-for-runtime

Conversation

@lbakshi
Copy link
Copy Markdown

@lbakshi lbakshi commented Apr 2, 2026

Summary

Injected tasks deployed via BRICKFLOW_INJECT_TASKS_DIR fail at runtime because the injection YAML/env vars only exist on the CI agent, not on the Databricks cluster. This PR writes the rendered template code to workspace files that the cluster can read at runtime.

Problem

When Brickflow deploys a job with injected tasks, the job definition on Databricks correctly includes those tasks. But when the cluster re-executes the entrypoint notebook, BRICKFLOW_INJECT_TASKS_DIR is not set, so _inject_tasks_from_yaml() is a no-op and workflow.get_task(task_id) raises TaskNotFoundError.

Solution

Deploy time

  1. serialize_for_runtime() renders the Jinja template and writes the output to _brickflow_injected/<task_name>.py in the project directory
  2. A lightweight JSON pointer (task_name, task_type, file_ref) is stored in brickflow_internal_injection_config — no inline source code
  3. The resolved workspace path is stored in brickflow_internal_injection_file_path using DAB's ${workspace.file_path} interpolation
  4. synth() adds sync.include: ["_brickflow_injected/**"] to the bundle config so DAB uploads the files even when the directory is gitignored

Runtime

  1. If workflow.get_task(t_id) fails, _reconstruct_injected_task() reads brickflow_internal_injection_file_path from the task parameters
  2. Opens the .py file from the workspace filesystem (/Workspace/Users/.../.brickflow_bundles/.../files/_brickflow_injected/<task>.py)
  3. Falls back to inline rendered_code in the JSON for backward compatibility with older deployments
  4. Reconstructs the Task object and executes normally

Why files instead of base_parameters?

The previous approach serialized the full rendered source into base_parameters as a JSON string. This was fragile:

  • Databricks has practical size limits on parameter values (~10KB)
  • Double/triple escaping (Python inside JSON inside YAML)
  • Harder to debug — code invisible in the workspace

With workspace files, the rendered code is inspectable in the Databricks UI, has no size limit, and the deploy→runtime contract is a simple file path.

Files changed

File Change
brickflow/context/context.py Add injection_config and injection_file_path to BrickflowInternalVariables
brickflow/engine/task_executor.py Write rendered code to file in serialize_for_runtime(); read from workspace file in create_task_function_from_config()
brickflow/engine/task.py Add injection_config_json field; emit ${workspace.file_path} parameter
brickflow/engine/workflow.py Thread injection_config_json through _add_task
brickflow/engine/project.py Serialize config at deploy; reconstruct task from workspace file at runtime
brickflow/codegen/databricks_bundle.py Add Sync(include=["_brickflow_injected/**"]) when injected files exist
tests/engine/test_task_injection.py Tests for file writing, workspace file reading, inline fallback, missing file fallback, no-code error

Test plan

  • All 38 task injection tests pass (including new file-based tests)
  • bundle.yml contains lightweight JSON pointer (no inline code)
  • bundle.yml contains sync.include for _brickflow_injected/
  • Files confirmed on Databricks workspace via API
  • End-to-end: injected task runs successfully on sole-blazer cluster
  • Verify with larger templates to confirm no size issues

Pre-review cleanup

  • Remove BRICKFLOW_RUNTIME_GIT_REF override from get_brickflow_libraries() in task.py — temporary hack to point the cluster at this fork for e2e testing

Closes nike-cdea-privacy/sidecar-test-zone-tmp#8

@lbakshi lbakshi changed the title [bugfix] Bake injected task config into base_parameters for runtime d… [bugfix] Bake injected task config into base_parameters for runtime dispatch Apr 2, 2026
@lbakshi lbakshi marked this pull request as ready for review April 2, 2026 18:52
@lbakshi lbakshi marked this pull request as draft April 2, 2026 18:53
@lbakshi
Copy link
Copy Markdown
Author

lbakshi commented Apr 2, 2026

⚠️ Temporary: BRICKFLOW_RUNTIME_GIT_REF override in task.py

To validate the fix end-to-end on a real Databricks cluster, I added a temporary env var override in get_brickflow_libraries() (task.py lines 1285-1287). When BRICKFLOW_RUNTIME_GIT_REF is set at deploy time, it overrides the hardcoded Nike-Inc/brickflow.git@main library URL so the cluster installs from our fork/branch instead.

This was necessary because:

  • The deploy-time half of the fix (baking config into base_parameters) works with our local branch
  • The runtime half (reconstructing the task from baked config) requires the cluster to install our patched brickflow — but the library URL is hardcoded to upstream @main

This must be removed before requesting review. The override is only for e2e testing. Once we confirm the fix works on the cluster, revert the BRICKFLOW_RUNTIME_GIT_REF block in get_brickflow_libraries().

# TEMPORARY — REMOVE BEFORE REVIEW
_override = os.environ.get("BRICKFLOW_RUNTIME_GIT_REF")
if _override:
    bf_lib = PypiTaskLibrary(f"brickflows @ git+{_override}")

At deploy time, rendered template code is written to
_brickflow_injected/<task>.py in the project directory. DAB syncs these
files to the workspace via a sync.include directive (overrides .gitignore).

A lightweight JSON pointer (task_name, task_type, file_ref) is stored in
brickflow_internal_injection_config, and the fully-resolved workspace
path goes into brickflow_internal_injection_file_path — DAB substitutes
${workspace.file_path} at deploy time.

At runtime the cluster reads the code from the workspace file, falling
back to inline rendered_code for backward compatibility.

This replaces the previous approach of stuffing the full rendered source
into base_parameters, which was fragile and size-limited.

Also includes a BRICKFLOW_RUNTIME_GIT_REF env-var override so the
cluster can install brickflow from a fork/branch during e2e testing
(remove before final merge).

Made-with: Cursor
@lbakshi lbakshi force-pushed the fix/bake-injected-task-config-for-runtime branch from c959092 to 1eb522f Compare April 2, 2026 22:08
@lbakshi lbakshi changed the title [bugfix] Bake injected task config into base_parameters for runtime dispatch feat: file-based runtime dispatch for injected tasks Apr 2, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant