def create_pipeline(ctx: Context, engine: str, pipeline_path: str,
package_path: Optional[str],
build_target_image: Optional[str],
build_base_image: Optional[str],
skaffold_cmd: Optional[str], endpoint: Optional[str],
iap_client_id: Optional[str], namespace: str,
build_image: bool) -> None:
"""Command definition to create a pipeline."""
# TODO(b/179847638): Delete checks for deprecated flags.
_check_deprecated_image_build_flags(build_target_image, skaffold_cmd,
package_path)
if build_base_image is not None and not build_image:
sys.exit(
'--build-base-image used without --build-image. You have to use '
'--build-image flag to build a container image for the pipeline.')
# TODO(b/142358865): Add support for container building for Airflow and Beam
# runners when they support container executors.
click.echo('Creating pipeline')
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.PIPELINE_DSL_PATH] = pipeline_path
ctx.flags_dict[labels.BASE_IMAGE] = build_base_image
ctx.flags_dict[labels.ENDPOINT] = endpoint
ctx.flags_dict[labels.IAP_CLIENT_ID] = iap_client_id
ctx.flags_dict[labels.NAMESPACE] = namespace
ctx.flags_dict[labels.BUILD_IMAGE] = build_image
handler_factory.create_handler(ctx.flags_dict).create_pipeline()
def testCreateHandlerOther(self):
self.flags_dict[labels.ENGINE_FLAG] = 'flink'
with self.assertRaises(Exception) as err:
handler_factory.create_handler(self.flags_dict)
self.assertEqual(
str(err.exception), 'Engine {} is not supported.'.format(
self.flags_dict[labels.ENGINE_FLAG]))
def test_create_handler_other(self):
self.flags_dict[labels.ENGINE_FLAG] = 'beam'
with self.assertRaises(RuntimeError) as err:
handler_factory.create_handler(self.flags_dict)
self.assertEqual(
str(err.exception), 'Engine {} is not supported.'.format(
self.flags_dict[labels.ENGINE_FLAG]))
def compile_pipeline(ctx: Context, engine: Text, pipeline_path: Text,
package_path: Text) -> None:
"""Command definition to compile a pipeline."""
click.echo('Compiling pipeline')
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.PIPELINE_DSL_PATH] = pipeline_path
ctx.flags_dict[labels.PIPELINE_PACKAGE_PATH] = package_path
handler_factory.create_handler(ctx.flags_dict).compile_pipeline()
def test_create_handler_auto(self):
self.flags_dict[labels.ENGINE_FLAG] = 'auto'
with self.assertRaises(Exception) as err:
handler_factory.create_handler(self.flags_dict)
self.assertEqual(
str(err.exception),
'Orchestrator {} missing in the environment.'.format(
self.flags_dict[labels.ENGINE_FLAG]))
def get_run(ctx: Context, engine: Text, pipeline_name: Text,
run_id: Text) -> None:
"""Command definition to stop a run."""
click.echo('Retrieving run status.')
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.RUN_ID] = run_id
ctx.flags_dict[labels.PIPELINE_NAME] = pipeline_name
handler_factory.create_handler(ctx.flags_dict).get_run()
def list_pipelines(ctx: Context, engine: Text, endpoint: Text,
iap_client_id: Text, namespace: Text) -> None:
"""Command definition to list pipelines."""
click.echo('Listing all pipelines')
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.ENDPOINT] = endpoint
ctx.flags_dict[labels.IAP_CLIENT_ID] = iap_client_id
ctx.flags_dict[labels.NAMESPACE] = namespace
handler_factory.create_handler(ctx.flags_dict).list_pipelines()
def compile_pipeline(ctx: Context, engine: str, pipeline_path: str,
package_path: str) -> None:
"""Command definition to compile a pipeline."""
# TODO(b/179847638): Delete checks for deprecated flags.
_check_deprecated_image_build_flags(pipeline_package_path=package_path)
click.echo('Compiling pipeline')
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.PIPELINE_DSL_PATH] = pipeline_path
handler_factory.create_handler(ctx.flags_dict).compile_pipeline()
def delete_pipeline(ctx: Context, engine: str, pipeline_name: str,
endpoint: str, iap_client_id: str, namespace: str) -> None:
"""Command definition to delete a pipeline."""
click.echo('Deleting pipeline')
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.PIPELINE_NAME] = pipeline_name
ctx.flags_dict[labels.ENDPOINT] = endpoint
ctx.flags_dict[labels.IAP_CLIENT_ID] = iap_client_id
ctx.flags_dict[labels.NAMESPACE] = namespace
handler_factory.create_handler(ctx.flags_dict).delete_pipeline()
def create_run(ctx: Context, engine: Text, pipeline_name: Text, endpoint: Text,
iap_client_id: Text, namespace: Text) -> None:
"""Command definition to create a pipeline run."""
click.echo('Creating a run for pipeline: ' + pipeline_name)
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.PIPELINE_NAME] = pipeline_name
ctx.flags_dict[labels.ENDPOINT] = endpoint
ctx.flags_dict[labels.IAP_CLIENT_ID] = iap_client_id
ctx.flags_dict[labels.NAMESPACE] = namespace
handler_factory.create_handler(ctx.flags_dict).create_run()
def delete_run(ctx: Context, engine: Text, run_id: Text, endpoint: Text,
iap_client_id: Text, namespace: Text) -> None:
"""Command definition to delete a run."""
click.echo('Deleting run.')
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.RUN_ID] = run_id
ctx.flags_dict[labels.ENDPOINT] = endpoint
ctx.flags_dict[labels.IAP_CLIENT_ID] = iap_client_id
ctx.flags_dict[labels.NAMESPACE] = namespace
handler_factory.create_handler(ctx.flags_dict).delete_run()
def list_runs(ctx: Context, engine: str, pipeline_name: str, endpoint: str,
iap_client_id: str, namespace: str) -> None:
"""Command definition to list all runs of a pipeline."""
click.echo('Listing all runs of pipeline: ' + pipeline_name)
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.PIPELINE_NAME] = pipeline_name
ctx.flags_dict[labels.ENDPOINT] = endpoint
ctx.flags_dict[labels.IAP_CLIENT_ID] = iap_client_id
ctx.flags_dict[labels.NAMESPACE] = namespace
handler_factory.create_handler(ctx.flags_dict).list_runs()
def terminate_run(ctx: Context, engine: str, run_id: str, endpoint: str,
iap_client_id: str, namespace: str) -> None:
"""Command definition to stop a run."""
click.echo('Terminating run.')
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.RUN_ID] = run_id
ctx.flags_dict[labels.ENDPOINT] = endpoint
ctx.flags_dict[labels.IAP_CLIENT_ID] = iap_client_id
ctx.flags_dict[labels.NAMESPACE] = namespace
handler_factory.create_handler(ctx.flags_dict).terminate_run()
def get_schema(ctx: Context, engine: Text, pipeline_name: Text, endpoint: Text,
iap_client_id: Text, namespace: Text) -> None:
"""Command definition to infer latest schema."""
click.echo('Getting latest schema.')
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.PIPELINE_NAME] = pipeline_name
ctx.flags_dict[labels.ENDPOINT] = endpoint
ctx.flags_dict[labels.IAP_CLIENT_ID] = iap_client_id
ctx.flags_dict[labels.NAMESPACE] = namespace
handler_factory.create_handler(ctx.flags_dict).get_schema()
def get_run(ctx: Context, engine: Text, pipeline_name: Text, run_id: Text,
endpoint: Text, iap_client_id: Text, namespace: Text) -> None:
"""Command definition to stop a run."""
click.echo('Retrieving run status.')
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.RUN_ID] = run_id
ctx.flags_dict[labels.PIPELINE_NAME] = pipeline_name
ctx.flags_dict[labels.ENDPOINT] = endpoint
ctx.flags_dict[labels.IAP_CLIENT_ID] = iap_client_id
ctx.flags_dict[labels.NAMESPACE] = namespace
handler_factory.create_handler(ctx.flags_dict).get_run()
def testCreateHandlerAirflow(self):
try:
from tfx.tools.cli.handler import airflow_handler # pylint: disable=g-import-not-at-top
except ImportError:
self.skipTest('Airflow is not available.')
self.flags_dict[labels.ENGINE_FLAG] = 'airflow'
with mock.patch.object(airflow_handler,
'AirflowHandler',
autospec=True) as mock_airflow_handler:
handler_factory.create_handler(self.flags_dict)
mock_airflow_handler.assert_called_once_with(self.flags_dict)
def create_pipeline(ctx: Context, engine: Text, pipeline_path: Text,
package_path: Text, endpoint: Text, iap_client_id: Text,
namespace: Text) -> None:
"""Command definition to create a pipeline."""
click.echo('Creating pipeline')
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.PIPELINE_DSL_PATH] = pipeline_path
ctx.flags_dict[labels.PIPELINE_PACKAGE_PATH] = package_path
ctx.flags_dict[labels.ENDPOINT] = endpoint
ctx.flags_dict[labels.IAP_CLIENT_ID] = iap_client_id
ctx.flags_dict[labels.NAMESPACE] = namespace
handler_factory.create_handler(ctx.flags_dict).create_pipeline()
def create_run(ctx: Context, engine: str, pipeline_name: str, endpoint: str,
iap_client_id: str, namespace: str, project: str,
region: str) -> None:
"""Command definition to create a pipeline run."""
click.echo('Creating a run for pipeline: ' + pipeline_name)
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.PIPELINE_NAME] = pipeline_name
ctx.flags_dict[labels.ENDPOINT] = endpoint
ctx.flags_dict[labels.IAP_CLIENT_ID] = iap_client_id
ctx.flags_dict[labels.NAMESPACE] = namespace
ctx.flags_dict[labels.GCP_PROJECT_ID] = project
ctx.flags_dict[labels.GCP_REGION] = region
handler_factory.create_handler(ctx.flags_dict).create_run()
def update_pipeline(ctx: Context, engine: str, pipeline_path: str,
package_path: Optional[str], skaffold_cmd: Optional[str],
endpoint: Optional[str], iap_client_id: Optional[str],
namespace: str, build_image: bool) -> None:
"""Command definition to update a pipeline."""
# TODO(b/179847638): Delete checks for deprecated flags.
_check_deprecated_image_build_flags(None, skaffold_cmd, package_path)
click.echo('Updating pipeline')
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.PIPELINE_DSL_PATH] = pipeline_path
ctx.flags_dict[labels.ENDPOINT] = endpoint
ctx.flags_dict[labels.IAP_CLIENT_ID] = iap_client_id
ctx.flags_dict[labels.NAMESPACE] = namespace
ctx.flags_dict[labels.BUILD_IMAGE] = build_image
handler_factory.create_handler(ctx.flags_dict).update_pipeline()
def create_pipeline(ctx: Context, engine: Text, pipeline_path: Text,
package_path: Text, build_target_image: Text,
skaffold_cmd: Text, endpoint: Text, iap_client_id: Text,
namespace: Text) -> None:
"""Command definition to create a pipeline."""
# TODO(b/142358865): Add support for container building for Airflow and Beam
# runners when they support container executors.
click.echo('Creating pipeline')
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.PIPELINE_DSL_PATH] = pipeline_path
ctx.flags_dict[labels.PIPELINE_PACKAGE_PATH] = package_path
ctx.flags_dict[labels.TARGET_IMAGE] = build_target_image
ctx.flags_dict[labels.SKAFFOLD_CMD] = skaffold_cmd
ctx.flags_dict[labels.ENDPOINT] = endpoint
ctx.flags_dict[labels.IAP_CLIENT_ID] = iap_client_id
ctx.flags_dict[labels.NAMESPACE] = namespace
handler_factory.create_handler(ctx.flags_dict).create_pipeline()
def testCreateHandlerKubeflow(self):
flags_dict = {
labels.ENGINE_FLAG: 'kubeflow',
labels.ENDPOINT: 'dummyEndpoint',
labels.IAP_CLIENT_ID: 'dummyID',
labels.NAMESPACE: 'kubeflow',
}
from tfx.tools.cli.handler import kubeflow_handler # pylint: disable=g-import-not-at-top
self.assertIsInstance(handler_factory.create_handler(flags_dict),
kubeflow_handler.KubeflowHandler)
def create_run(ctx: Context, engine: str, pipeline_name: str, endpoint: str,
iap_client_id: str, namespace: str, project: str, region: str,
runtime_parameter: Iterable[str]) -> None:
"""Command definition to create a pipeline run."""
click.echo('Creating a run for pipeline: ' + pipeline_name)
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.PIPELINE_NAME] = pipeline_name
ctx.flags_dict[labels.ENDPOINT] = endpoint
ctx.flags_dict[labels.IAP_CLIENT_ID] = iap_client_id
ctx.flags_dict[labels.NAMESPACE] = namespace
ctx.flags_dict[labels.GCP_PROJECT_ID] = project
ctx.flags_dict[labels.GCP_REGION] = region
ctx.flags_dict[labels.RUNTIME_PARAMETER] = _parse_runtime_parameters(
runtime_parameter)
handler = handler_factory.create_handler(ctx.flags_dict)
if (ctx.flags_dict[labels.ENGINE_FLAG]
not in (labels.KUBEFLOW_ENGINE, labels.AIRFLOW_ENGINE,
labels.VERTEX_ENGINE)) and runtime_parameter:
raise NotImplementedError(
'Currently runtime parameter is only supported in kubeflow, vertex, '
'and airflow.')
handler.create_run()
def get_schema(ctx: Context, engine: Text, pipeline_name: Text) -> None:
"""Command definition to infer latest schema."""
click.echo('Getting latest schema.')
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.PIPELINE_NAME] = pipeline_name
handler_factory.create_handler(ctx.flags_dict).get_schema()
def update_pipeline(ctx: Context, engine: Text, pipeline_path: Text) -> None:
"""Command definition to update a pipeline."""
click.echo('Updating pipeline')
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.PIPELINE_DSL_PATH] = pipeline_path
handler_factory.create_handler(ctx.flags_dict).update_pipeline()
def delete_pipeline(ctx: Context, pipeline_name: Text, engine: Text) -> None:
"""Command definition to delete a pipeline."""
click.echo('Deleting pipeline')
ctx.flags_dict[labels.ENGINE_FLAG] = engine
ctx.flags_dict[labels.PIPELINE_NAME] = pipeline_name
handler_factory.create_handler(ctx.flags_dict).delete_pipeline()
def list_pipelines(ctx: Context, engine: Text) -> None:
"""Command definition to list pipelines."""
click.echo('Listing all pipelines')
ctx.flags_dict[labels.ENGINE_FLAG] = engine
handler_factory.create_handler(ctx.flags_dict).list_pipelines()
def test_create_handler_kubeflow(self):
self.flags_dict[labels.ENGINE_FLAG] = 'kubeflow'
self.assertIsInstance(
handler_factory.create_handler(self.flags_dict),
kubeflow_handler.KubeflowHandler)
def testCreateHandlerBeam(self):
self.flags_dict[labels.ENGINE_FLAG] = 'beam'
self.assertIsInstance(handler_factory.create_handler(self.flags_dict),
beam_handler.BeamHandler)
def testCreateHandlerAirflow(self, mock_airflow_handler):
self.flags_dict[labels.ENGINE_FLAG] = 'airflow'
handler_factory.create_handler(self.flags_dict)
mock_airflow_handler.assert_called_once_with(self.flags_dict)
def testCreateHandlerAirflow(self):
self.flags_dict[labels.ENGINE_FLAG] = 'airflow'
self.assertIsInstance(handler_factory.create_handler(self.flags_dict),
airflow_handler.AirflowHandler)
