torchx icon indicating copy to clipboard operation
torchx copied to clipboard

Upgrade TorchX to Support Kubeflow Pipelines v2

Open chethanuk opened this issue 7 months ago • 0 comments

KFP v1 has been deprecated for a while now, and we really needed to get onto the newer version of the KFP SDK, so we are migrating KFP v1 to the v2 API

Key Changes

Dependencies

  • kfp: 1.8.22>=2.8.0
  • Added kfp-kubernetes>=1.4.0 for K8s-specific features
  • protobuf: 3.20.3>=4.21.0 (KFP v2 requirement)

Core Adapter (adapter.py)

  • Complete refactor to use KFP v2 APIs
  • Replaced ContainerOp with @dsl.component decorator pattern
  • container_from_app now returns PipelineTask instead of ContainerOp
  • Enhanced GPU/accelerator configuration support
  • Improved resource limits and env var handling

Version Handling

  • Changed from blocking KFP v2 to showing deprecation warning for v1 users
  • Warning: "KFP version 1.x.x is deprecated! Please upgrade to kfp version 2.x.x"

Examples & Tests

  • Updated all existing pipeline examples for KFP v2 patterns
  • Added new examples showcasing v2 features (retry policies, caching, PVC management)
  • Comprehensive test coverage with new integration tests

Breaking Changes

# 1. Import Changes: Components using TorchX KFP adapter need to update imports:
# Old (KFP v1)
from kfp import dsl
from kfp.compiler import Compiler

# New (KFP v2)  
from kfp import compiler, dsl

# 2. Component Creation: The adapter now returns PipelineTask instead of ContainerOp
# Old
task = container_from_app(app)  # Returns ContainerOp
task.container.set_tty()

# New
task = container_from_app(app)  # Returns PipelineTask
task.set_display_name("My Task")

# 3. Pipeline Compilation: Must use keyword arguments
# Old
compiler.Compiler().compile(pipeline, "pipeline.yaml")

# New
compiler.Compiler().compile(pipeline_func=pipeline, package_path="pipeline.yaml")

Migration

  1. Update to kfp>=2.8.0
  2. Use KFP v2 patterns (see updated examples)
  3. Replace .after() with proper task dependencies or use the KFP v2 .after() method
  4. Use ``set_display_name()` instead of container property manipulation

Test plan: image

chethanuk avatar May 25 '25 19:05 chethanuk