def _run_pipeline(
self,
pipeline: tfx_pipeline.Pipeline,
exit_handler: Optional[base_node.BaseNode] = None) -> None:
"""Trigger the pipeline execution with a specific job ID."""
# Ensure cleanup regardless of whether pipeline succeeds or fails.
self.addCleanup(self._delete_pipeline_output,
pipeline.pipeline_info.pipeline_name)
config = kubeflow_v2_dag_runner.KubeflowV2DagRunnerConfig(
default_image=self.container_image)
executing_kubeflow_v2_dag_runner = kubeflow_v2_dag_runner.KubeflowV2DagRunner(
config=config, output_filename='pipeline.json')
if exit_handler:
executing_kubeflow_v2_dag_runner.set_exit_handler(exit_handler)
_ = executing_kubeflow_v2_dag_runner.run(pipeline, write_out=True)
job_id = pipeline.pipeline_info.pipeline_name
job = pipeline_jobs.PipelineJob(
template_path='pipeline.json',
job_id=job_id,
display_name=pipeline.pipeline_info.pipeline_name)
job.run(sync=False)
job.wait_for_resource_creation()
vertex_client_utils.poll_job_status(job_id, _MAX_JOB_EXECUTION_TIME,
_POLLING_INTERVAL_IN_SECONDS)
def test_pipeline_failure_raises(self, mock_load_yaml_and_json, sync):
aiplatform.init(
project=_TEST_PROJECT,
staging_bucket=_TEST_GCS_BUCKET_NAME,
location=_TEST_LOCATION,
credentials=_TEST_CREDENTIALS,
)
job = pipeline_jobs.PipelineJob(
display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME,
template_path=_TEST_TEMPLATE_PATH,
job_id=_TEST_PIPELINE_JOB_ID,
parameter_values=_TEST_PIPELINE_PARAMETER_VALUES,
enable_caching=True,
)
with pytest.raises(RuntimeError):
job.run(
service_account=_TEST_SERVICE_ACCOUNT,
network=_TEST_NETWORK,
sync=sync,
)
if not sync:
job.wait()
def test_done_method_pipeline_service(
self,
mock_pipeline_service_create,
mock_pipeline_service_get,
job_spec,
mock_load_yaml_and_json,
):
aiplatform.init(
project=_TEST_PROJECT,
staging_bucket=_TEST_GCS_BUCKET_NAME,
location=_TEST_LOCATION,
credentials=_TEST_CREDENTIALS,
)
job = pipeline_jobs.PipelineJob(
display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME,
template_path=_TEST_TEMPLATE_PATH,
job_id=_TEST_PIPELINE_JOB_ID,
parameter_values=_TEST_PIPELINE_PARAMETER_VALUES,
enable_caching=True,
)
job.submit(service_account=_TEST_SERVICE_ACCOUNT,
network=_TEST_NETWORK)
assert job.done() is False
job.wait()
assert job.done() is True
def test_run_call_pipeline_service_create(
self, mock_pipeline_service_create, mock_pipeline_service_get, sync,
):
aiplatform.init(
project=_TEST_PROJECT,
staging_bucket=_TEST_GCS_BUCKET_NAME,
location=_TEST_LOCATION,
credentials=_TEST_CREDENTIALS,
)
job = pipeline_jobs.PipelineJob(
display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME,
template_path=_TEST_TEMPLATE_PATH,
job_id=_TEST_PIPELINE_JOB_ID,
parameter_values=_TEST_PIPELINE_PARAMETER_VALUES,
enable_caching=True,
)
job.run(
service_account=_TEST_SERVICE_ACCOUNT, network=_TEST_NETWORK, sync=sync,
)
if not sync:
job.wait()
expected_runtime_config_dict = {
"gcs_output_directory": _TEST_GCS_BUCKET_NAME,
"parameters": {"name_param": {"stringValue": "hello"}},
}
runtime_config = gca_pipeline_job_v1beta1.PipelineJob.RuntimeConfig()._pb
json_format.ParseDict(expected_runtime_config_dict, runtime_config)
# Construct expected request
expected_gapic_pipeline_job = gca_pipeline_job_v1beta1.PipelineJob(
display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME,
name=_TEST_PIPELINE_JOB_NAME,
pipeline_spec={
"components": {},
"pipelineInfo": _TEST_PIPELINE_JOB_SPEC["pipelineSpec"]["pipelineInfo"],
"root": _TEST_PIPELINE_JOB_SPEC["pipelineSpec"]["root"],
},
runtime_config=runtime_config,
service_account=_TEST_SERVICE_ACCOUNT,
network=_TEST_NETWORK,
)
mock_pipeline_service_create.assert_called_once_with(
parent=_TEST_PARENT,
pipeline_job=expected_gapic_pipeline_job,
pipeline_job_id=_TEST_PIPELINE_JOB_ID,
)
mock_pipeline_service_get.assert_called_with(name=_TEST_PIPELINE_JOB_NAME)
assert job._gca_resource == make_pipeline_job(
gca_pipeline_state_v1beta1.PipelineState.PIPELINE_STATE_SUCCEEDED
)
def test_cancel_pipeline_job_without_running(
self, mock_pipeline_service_cancel, mock_load_json,
):
aiplatform.init(
project=_TEST_PROJECT,
staging_bucket=_TEST_GCS_BUCKET_NAME,
credentials=_TEST_CREDENTIALS,
)
job = pipeline_jobs.PipelineJob(
display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME,
template_path=_TEST_TEMPLATE_PATH,
job_id=_TEST_PIPELINE_JOB_ID,
)
with pytest.raises(RuntimeError) as e:
job.cancel()
assert e.match(regexp=r"PipelineJob resource has not been created")
def test_list_pipeline_job(self, mock_pipeline_service_list, mock_load_json):
aiplatform.init(
project=_TEST_PROJECT,
staging_bucket=_TEST_GCS_BUCKET_NAME,
credentials=_TEST_CREDENTIALS,
)
job = pipeline_jobs.PipelineJob(
display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME,
template_path=_TEST_TEMPLATE_PATH,
job_id=_TEST_PIPELINE_JOB_ID,
)
job.run()
job.list()
mock_pipeline_service_list.assert_called_once_with(
request={"parent": _TEST_PARENT, "filter": None}
)
def test_cancel_pipeline_job(self, mock_pipeline_service_cancel, mock_load_json):
aiplatform.init(
project=_TEST_PROJECT,
staging_bucket=_TEST_GCS_BUCKET_NAME,
credentials=_TEST_CREDENTIALS,
)
job = pipeline_jobs.PipelineJob(
display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME,
template_path=_TEST_TEMPLATE_PATH,
job_id=_TEST_PIPELINE_JOB_ID,
)
job.run()
job.cancel()
mock_pipeline_service_cancel.assert_called_once_with(
name=_TEST_PIPELINE_JOB_NAME
)
def create_run(self) -> None:
"""Runs a pipeline in Vertex Pipelines."""
self._prepare_vertex()
pipeline_name = self.flags_dict[labels.PIPELINE_NAME]
# In Vertex AI, runtime parameter string value is parsed from the server,
# so client directly sends Dict[str, str] value.
unparsed_runtime_parameters = self.flags_dict[labels.RUNTIME_PARAMETER]
job = pipeline_jobs.PipelineJob(
display_name=pipeline_name,
template_path=self._get_pipeline_definition_path(pipeline_name),
parameter_values=unparsed_runtime_parameters)
# TODO(b/198114641): Delete pytype exception after upgrading source code
# to aiplatform>=1.3.
job.run(sync=False) # pytype: disable=attribute-error
job.wait_for_resource_creation() # pytype: disable=attribute-error
click.echo('Run created for pipeline: ' + pipeline_name)
self._print_run(job)
def test_submit_call_pipeline_service_pipeline_job_create_legacy(
self,
mock_pipeline_service_create,
mock_pipeline_service_get,
job_spec_json,
mock_load_json,
):
aiplatform.init(
project=_TEST_PROJECT,
staging_bucket=_TEST_GCS_BUCKET_NAME,
location=_TEST_LOCATION,
credentials=_TEST_CREDENTIALS,
)
job = pipeline_jobs.PipelineJob(
display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME,
template_path=_TEST_TEMPLATE_PATH,
job_id=_TEST_PIPELINE_JOB_ID,
parameter_values=_TEST_PIPELINE_PARAMETER_VALUES_LEGACY,
enable_caching=True,
)
job.submit(service_account=_TEST_SERVICE_ACCOUNT, network=_TEST_NETWORK)
expected_runtime_config_dict = {
"parameters": {"string_param": {"stringValue": "hello"}},
"gcsOutputDirectory": _TEST_GCS_BUCKET_NAME,
}
runtime_config = gca_pipeline_job_v1.PipelineJob.RuntimeConfig()._pb
json_format.ParseDict(expected_runtime_config_dict, runtime_config)
pipeline_spec = job_spec_json.get("pipelineSpec") or job_spec_json
# Construct expected request
expected_gapic_pipeline_job = gca_pipeline_job_v1.PipelineJob(
display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME,
pipeline_spec={
"components": {},
"pipelineInfo": pipeline_spec["pipelineInfo"],
"root": pipeline_spec["root"],
"schemaVersion": "2.0.0",
},
runtime_config=runtime_config,
service_account=_TEST_SERVICE_ACCOUNT,
network=_TEST_NETWORK,
)
mock_pipeline_service_create.assert_called_once_with(
parent=_TEST_PARENT,
pipeline_job=expected_gapic_pipeline_job,
pipeline_job_id=_TEST_PIPELINE_JOB_ID,
)
assert not mock_pipeline_service_get.called
job.wait()
mock_pipeline_service_get.assert_called_with(
name=_TEST_PIPELINE_JOB_NAME, retry=base._DEFAULT_RETRY
)
assert job._gca_resource == make_pipeline_job(
gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED
)
def test_run_call_pipeline_service_create_with_timeout(
self,
mock_pipeline_service_create,
mock_pipeline_service_get,
job_spec,
mock_load_yaml_and_json,
sync,
):
aiplatform.init(
project=_TEST_PROJECT,
staging_bucket=_TEST_GCS_BUCKET_NAME,
location=_TEST_LOCATION,
credentials=_TEST_CREDENTIALS,
)
job = pipeline_jobs.PipelineJob(
display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME,
template_path=_TEST_TEMPLATE_PATH,
job_id=_TEST_PIPELINE_JOB_ID,
parameter_values=_TEST_PIPELINE_PARAMETER_VALUES,
enable_caching=True,
)
job.run(
service_account=_TEST_SERVICE_ACCOUNT,
network=_TEST_NETWORK,
sync=sync,
create_request_timeout=180.0,
)
if not sync:
job.wait()
expected_runtime_config_dict = {
"gcsOutputDirectory": _TEST_GCS_BUCKET_NAME,
"parameterValues": _TEST_PIPELINE_PARAMETER_VALUES,
}
runtime_config = gca_pipeline_job_v1.PipelineJob.RuntimeConfig()._pb
json_format.ParseDict(expected_runtime_config_dict, runtime_config)
job_spec = yaml.safe_load(job_spec)
pipeline_spec = job_spec.get("pipelineSpec") or job_spec
# Construct expected request
expected_gapic_pipeline_job = gca_pipeline_job_v1.PipelineJob(
display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME,
pipeline_spec={
"components": {},
"pipelineInfo": pipeline_spec["pipelineInfo"],
"root": pipeline_spec["root"],
"schemaVersion": "2.1.0",
},
runtime_config=runtime_config,
service_account=_TEST_SERVICE_ACCOUNT,
network=_TEST_NETWORK,
)
mock_pipeline_service_create.assert_called_once_with(
parent=_TEST_PARENT,
pipeline_job=expected_gapic_pipeline_job,
pipeline_job_id=_TEST_PIPELINE_JOB_ID,
timeout=180.0,
)
