Polyaxon v3 is coming →

Run an existing TFX pipeline in a Polyaxon job

Wrap a local TensorFlow Extended pipeline in one Polyaxon job while preserving its own metadata and artifacts.

September 15, 2024by Polyaxon

Run an existing TensorFlow Extended (TFX) pipeline inside a Polyaxon job to keep its execution alongside your other ML workloads. TFX defines the validation, transformation, training, and evaluation steps; Polyaxon schedules the container and provides its run context.

The workflow below wraps a TFX local runner in one job. It does not translate a TFX graph into native Polyaxon DAG operations or provide per-component distributed scheduling.

A TFX local pipeline executes inside one Polyaxon job while its metadata and artifacts are persisted separately

Prepare the pipeline image

Start with a pipeline that already runs successfully using TFX's local pipeline workflow. Package its Python entry point, custom components, and compatible TFX, TensorFlow, and Python dependencies in an image. Pin those dependencies together rather than independently upgrading packages at container startup.

Your entry point should construct the pipeline and execute it with the appropriate runner. TFX's LocalDagRunner runs components in dependency order inside that execution environment. Components that require external services still need those services and credentials.

Schedule it as a job

The following is a job wrapper, not a complete TFX pipeline. Replace the example image with your published image, containing /app/local_runner.py and its dependencies. Configure storage and credentials before submitting it.

version: 1.1
kind: component
name: tfx-pipeline
run:
  kind: job
  container:
    image: registry.example.com/ml/tfx-pipeline:validated
    workingDir: /app
    command: [python, local_runner.py]

Save the wrapper as tfx-pipeline.yaml, select your Polyaxon project and queue, and submit it using the job workflow:

polyaxon run -f tfx-pipeline.yaml

Size the job for the pipeline's actual resource use. A single wrapper does not give each TFX component its own Polyaxon CPU, memory, or GPU allocation.

Preserve artifacts and metadata

Configure the TFX pipeline root and ML Metadata store explicitly. Container-local paths disappear when the workload is removed; use suitable durable storage and Polyaxon connections for the resources your pipeline needs.

Isolate outputs from concurrent runs, retain the pipeline's metadata needed for reuse or inspection, and avoid sharing a local SQLite database across independent writers. Record the dataset version, image revision, pipeline identifier, and output locations in Polyaxon tracking. TFX metadata and Polyaxon tracking are separate systems; logging one does not automatically synchronize the other.

Choose an orchestration boundary

Use this single-job approach for an existing local pipeline whose resource and failure behavior fits one workload. If your TFX pipeline needs a different production orchestrator, follow TFX's supported runner configuration and observe that external execution explicitly.

Alternatively, implement the stages as separate Polyaxon components connected in a DAG. That is an explicit application design, not an automatic conversion. For a validated SavedModel ready for inference, continue with TensorFlow Serving.