torchx
torchx copied to clipboard
Upgrade TorchX to Support Kubeflow Pipelines v2
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.0for 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
ContainerOpwith@dsl.componentdecorator pattern container_from_appnow returnsPipelineTaskinstead ofContainerOp- 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
- Update to
kfp>=2.8.0 - Use KFP v2 patterns (see updated examples)
- Replace
.after()with proper task dependencies or use the KFP v2.after()method - Use ``set_display_name()` instead of container property manipulation
Test plan: