chore: fix pipeline job system tests

sararob · copybara-github · commit 032aa21daf60 · 2023-09-26T05:52:15.000-07:00
PiperOrigin-RevId: 568511863
diff --git a/setup.py b/setup.py
@@ -148,6 +148,7 @@
         "torch; python_version<'3.8'",
         "xgboost",
         "xgboost_ray",
+        "requests-toolbelt < 1.0.0",
     ]
 )
 
diff --git a/tests/system/aiplatform/test_pipeline_job.py b/tests/system/aiplatform/test_pipeline_job.py
@@ -29,7 +29,7 @@ class TestPipelineJob(e2e_base.TestEndToEnd):
     _temp_prefix = "tmpvrtxsdk-e2e"
 
     def test_add_pipeline_job_to_experiment(self, shared_state):
-        from kfp import components
+        from kfp import dsl
 
         # Components:
         def train(
@@ -39,7 +39,7 @@ def train(
             print(f"number_of_epochs={number_of_epochs}")
             print(f"learning_rate={learning_rate}")
 
-        train_op = components.create_component_from_func(train)
+        train_op = dsl.component(train)
 
         # Pipeline:
         def training_pipeline(number_of_epochs: int = 10):
