Building a Reusable Kubeflow Pipelines Component for TFJob‑Based TensorFlow Training

Problem: Repeating TFJob boilerplate in every pipeline

When you want to run distributed TensorFlow training inside Kubeflow Pipelines, the usual approach is to write a separate Python function that calls kubectl apply -f tfjob.yaml or directly uses the Kubernetes client. Each new experiment then duplicates the TFJob manifest, the image‑pull logic, and the cleanup steps. This makes pipelines hard to maintain and prone to version drift between the training code and the TFJob CRD.

Thesis: Encapsulate the TFJob creation in a KFP component

By packaging the TFJob manifest inside a containerized Kubeflow Pipelines component, you get a single, version‑controlled step that can be reused across experiments. The component applies the TFJob CR, lets the TFJob operator handle scheduling, GPU allocation, and pod cleanup, and surfaces the job’s status and artifacts back to the pipeline.

Component design

The component image must contain:

• The training code (or a way to pull it from a repository).
• A tool that can apply Kubernetes manifests – typically kubectl or the TFJob operator’s sidecar.
• Any dependencies needed to run the training script (e.g., TensorFlow, CUDA).

The component receives parameters such as the training script path, container image URI, hyper‑parameters, and output locations. It renders a TFJob manifest (using a simple Jinja2 template or plain string substitution) and applies it with kubectl apply -f -. After the apply, the component waits for the TFJob to reach a terminal state (Succeeded or Failed) by polling the TFJob status, then exits.

Example component definition (Python SDK)

from kfp import dsl
from kfp.dsl import component

@component(
    base_image='your-registry/tfjob-component:1.0.0',
    packages_to_install=['jinja2']
)
def tfjob_train_component(
    project_id: str,
    region: str,
    tfjob_image: str,
    script_uri: str,
    hyperparams: dict,
    output_path: str,
    tfjob_name: str = dsl.OutputPath(str),
):
    import json, os, time, jinja2, subprocess
    
    # Render TFJob manifest
    template_str = """
apiVersion: kubeflow.org/v1
kind: TFJob
metadata:
  name: {{ name }}
spec:
  tfReplicaSpecs:
    Worker:
      replicas: 1
      template:
        spec:
          containers:
          - name: tensorflow
            image: {{ image }}
            command: ["python", "{{ script }}"]
            args: {{ hyperparams | tojson }}
            volumeMounts:
            - name: output
              mountPath: /output
          volumes:
          - name: output
            emptyDir: {}
    """
    rendered = jinja2.Template(template_str).render(
        name=tfjob_name,
        image=tfjob_image,
        script=os.path.basename(script_uri),
        hyperparams=hyperparams
    )
    manifest_path = '/tmp/tfjob.yaml'
    with open(manifest_path, 'w') as f:
        f.write(rendered)
    
    # Apply the manifest
    subprocess.check_call(['kubectl', 'apply', '-f', manifest_path])
    
    # Wait for completion
    while True:
        result = subprocess.check_output([
            'kubectl', 'get', 'tfjob', tfjob_name, '-o', 'json'
        ], text=True)
        status = json.loads(result).get('status', {})
        if status.get('succeeded') == 1:
            break
        if status.get('failed') == 1:
            raise RuntimeError('TFJob failed')
        time.sleep(10)
    
    # Copy output artifact (example: model saved to /output/model.h5)
    subprocess.check_call([
        'kubectl', 'cp', f'{tfjob_name}:/output/model.h5', output_path
    ])

Place this function in a file, e.g., tfjob_component.py, and register it with your Kubeflow Pipelines instance using the SDK or the UI.

Worked example: Running a simple MNIST training job

1. Build the component image (run on a workstation with Docker and kubectl configured for the target cluster):

   docker build -t your-registry/tfjob-component:1.0.0 \
       --build-arg TF_VERSION=2.13.0 \
       -f Dockerfile.component .
   docker push your-registry/tfjob-component:1.0.0
   

2. Define the pipeline that calls the component:

   @dsl.pipeline(name='mnist-tfjob')
   def mnist_pipeline():
       tfjob_train_component(
           project_id='my-gcp-project',
           region='us-central1',
           tfjob_image='your-registry/tensorflow-mnist:latest',
           script_uri='gs://my-bucket/train_mnist.py',
           hyperparams={'--epochs': '5', '--batch-size': '64'},
           output_path='/tmp/model.h5'
       )
   

3. Compile and submit the pipeline (requires the KFP SDK and appropriate namespace permissions):

   dsl.compile(mnist_pipeline, 'mnist_pipeline.yaml')
   kfp client create-run --experiment-name mnist --job-name mnist-run \
       --pipeline-package-file mnist_pipeline.yaml
   

4. Verify that a TFJob appears:

   kubectl get tfjob -n kubeflow-pipelines
   # Expected output includes a TFJob with status Running/Succeeded
   

5. Check logs from the component pod and the TFJob worker pods if you need to debug:

   kubectl logs -l job-name=mnist-run-component -c tfjob-component
   kubectl logs -l tfjob-name=mnist-run-tfjob
   

Trade‑offs and limitations

• Image maintenance: The component image must bundle both the training dependencies and a manifest‑application tool (kubectl or the TFJob operator sidecar). If you upgrade TensorFlow or the TFJob CRD, you need to rebuild and push a new image.
• Permission scope: The component runs with the service account of the pipeline pod. It needs rbac.authorization.k8s.io/v1 permissions to create, get, and delete TFJob objects in the target namespace. Over‑privileged accounts increase blast‑radius risk.
• Debugging surface: Failures can stem from manifest validation (caught in the component pod logs) or from the TFJob controller (resource quotas, GPU device plugin issues). You must check both log streams to pinpoint the root cause.

Actionable next steps

Start by extracting the TFJob manifest from one of your existing pipelines into a reusable component as shown above. Parameterize the image URI, script location, and hyper‑parameters so the same component can serve multiple teams. Keep the component image lightweight—multi‑stage builds help separate build‑time dependencies from the runtime image. Finally, enforce a naming convention for TFJob resources (e.g., <pipeline-name>-<run-id>-tfjob) to simplify monitoring and cleanup.