fix(components): Fixed the GCP - ML Engine - Batch predict component. F…

…ixes #4430 (#4432) * Components - Fixed the GCP - ML Engine - Batch predict component Fixes #4430 The issue was introduced in #3850. That PR has added a new parameter in the middle of the the create_job function signature which can cause breaking changes as the parameter ordering changes. * Fixed test
kubeflow · Sep 3, 2020 · 2beddd0 · 2beddd0
1 parent bec0592
commit 2beddd0
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 4 deletions.
diff --git a/...nents/gcp/container/component_sdk/python/kfp_component/google/ml_engine/_batch_predict.py b/...nents/gcp/container/component_sdk/python/kfp_component/google/ml_engine/_batch_predict.py
@@ -67,7 +67,13 @@ def batch_predict(project_id, model_path, input_paths, input_data_format,
     job = {
         'predictionInput': prediction_input
     }
-    create_job(project_id, job, job_id_prefix, wait_interval, job_id_output_path=job_id_output_path)
+    create_job(
+        project_id=project_id,
+        job=job,
+        job_id_prefix=job_id_prefix,
+        wait_interval=wait_interval,
+        job_id_output_path=job_id_output_path,
+    )
 
 def _is_model_name(name):
     return re.match(r'/projects/[^/]+/models/[^/]+$', name)

diff --git a/components/gcp/container/component_sdk/python/kfp_component/google/ml_engine/_train.py b/components/gcp/container/component_sdk/python/kfp_component/google/ml_engine/_train.py
@@ -82,4 +82,11 @@ def train(project_id, job_id_output_path, python_module=None, package_uris=None,
     job = {
         'trainingInput': training_input
     }
-    return create_job(project_id, job, job_id_prefix, job_id, wait_interval, job_id_output_path=job_id_output_path)
+    return create_job(
+        project_id=project_id,
+        job=job,
+        job_id_prefix=job_id_prefix,
+        job_id=job_id,
+        wait_interval=wait_interval,
+        job_id_output_path=job_id_output_path,
+    )
diff --git a/components/gcp/container/component_sdk/python/tests/google/ml_engine/test__train.py b/components/gcp/container/component_sdk/python/tests/google/ml_engine/test__train.py
@@ -39,7 +39,7 @@ def test_train_succeed(self, mock_create_job):
             job_id_output_path='/tmp/kfp/output/ml_engine/job_id.txt',
         )
 
-        mock_create_job.assert_called_with('proj-1', {
+        mock_create_job.assert_called_with(project_id='proj-1', job={
             'trainingInput': {
                 'pythonModule': 'mock.module',
                 'packageUris': ['gs://test/package'],
@@ -55,4 +55,4 @@ def test_train_succeed(self, mock_create_job):
                     'imageUri': 'debian:latest'
                 }
             }
-        }, 'job-', 'job-1', 30, job_id_output_path='/tmp/kfp/output/ml_engine/job_id.txt')
+        }, job_id_prefix='job-', job_id='job-1', wait_interval=30, job_id_output_path='/tmp/kfp/output/ml_engine/job_id.txt')