def run(self, pipeline: pipeline_py.Pipeline) -> None:
"""Runs given logical pipeline locally.
Args:
pipeline: Logical pipeline containing pipeline args and components.
"""
# For CLI, while creating or updating pipeline, pipeline_args are extracted
# and hence we avoid executing the pipeline.
if 'TFX_JSON_EXPORT_PIPELINE_ARGS_PATH' in os.environ:
return
for component in pipeline.components:
# TODO(b/187122662): Pass through pip dependencies as a first-class
# component flag.
if isinstance(component, base_component.BaseComponent):
component._resolve_pip_dependencies( # pylint: disable=protected-access
pipeline.pipeline_info.pipeline_root)
c = compiler.Compiler()
pipeline = c.compile(pipeline)
# Substitute the runtime parameter to be a concrete run_id
runtime_parameter_utils.substitute_runtime_parameter(
pipeline, {
constants.PIPELINE_RUN_ID_PARAMETER_NAME:
datetime.datetime.now().isoformat(),
})
deployment_config = runner_utils.extract_local_deployment_config(pipeline)
connection_config = deployment_config.metadata_connection_config
logging.info('Running pipeline:\n %s', pipeline)
logging.info('Using deployment config:\n %s', deployment_config)
logging.info('Using connection config:\n %s', connection_config)
with telemetry_utils.scoped_labels(
{telemetry_utils.LABEL_TFX_RUNNER: 'local'}):
# Run each component. Note that the pipeline.components list is in
# topological order.
#
# TODO(b/171319478): After IR-based execution is used, used multi-threaded
# execution so that independent components can be run in parallel.
for node in pipeline.nodes:
pipeline_node = node.pipeline_node
node_id = pipeline_node.node_info.id
executor_spec = runner_utils.extract_executor_spec(
deployment_config, node_id)
custom_driver_spec = runner_utils.extract_custom_driver_spec(
deployment_config, node_id)
component_launcher = launcher.Launcher(
pipeline_node=pipeline_node,
mlmd_connection=metadata.Metadata(connection_config),
pipeline_info=pipeline.pipeline_info,
pipeline_runtime_spec=pipeline.runtime_spec,
executor_spec=executor_spec,
custom_driver_spec=custom_driver_spec)
logging.info('Component %s is running.', node_id)
component_launcher.launch()
logging.info('Component %s is finished.', node_id)
def run(self,
pipeline: tfx_pipeline.Pipeline,
run_name: Optional[str] = None) -> "airflow.DAG":
"""Deploys given logical pipeline on Airflow.
Args:
pipeline: Logical pipeline containing pipeline args and comps.
run_name: Optional name for the run.
Returns:
An Airflow DAG.
"""
# Only import these when needed.
import airflow # noqa
from zenml.integrations.airflow.orchestrators import airflow_component
# Merge airflow-specific configs with pipeline args
airflow_dag = airflow.DAG(
dag_id=pipeline.pipeline_info.pipeline_name,
**(typing.cast(AirflowPipelineConfig,
self._config).airflow_dag_config),
is_paused_upon_creation=False,
catchup=False, # no backfill
)
if "tmp_dir" not in pipeline.additional_pipeline_args:
tmp_dir = os.path.join(pipeline.pipeline_info.pipeline_root,
".temp", "")
pipeline.additional_pipeline_args["tmp_dir"] = tmp_dir
for component in pipeline.components:
if isinstance(component, base_component.BaseComponent):
component._resolve_pip_dependencies(
pipeline.pipeline_info.pipeline_root)
self._replace_runtime_params(component)
c = compiler.Compiler()
pipeline = c.compile(pipeline)
run_name = run_name or datetime.now().strftime("%d_%h_%y-%H_%M_%S_%f")
# Substitute the runtime parameter to be a concrete run_id
runtime_parameter_utils.substitute_runtime_parameter(
pipeline,
{
"pipeline-run-id": run_name,
},
)
deployment_config = runner_utils.extract_local_deployment_config(
pipeline)
connection_config = deployment_config.metadata_connection_config # type: ignore[attr-defined] # noqa
component_impl_map = {}
for node in pipeline.nodes:
pipeline_node = node.pipeline_node
node_id = pipeline_node.node_info.id
executor_spec = runner_utils.extract_executor_spec(
deployment_config, node_id)
custom_driver_spec = runner_utils.extract_custom_driver_spec(
deployment_config, node_id)
current_airflow_component = airflow_component.AirflowComponent(
parent_dag=airflow_dag,
pipeline_node=pipeline_node,
mlmd_connection=connection_config,
pipeline_info=pipeline.pipeline_info,
pipeline_runtime_spec=pipeline.runtime_spec,
executor_spec=executor_spec,
custom_driver_spec=custom_driver_spec,
)
component_impl_map[node_id] = current_airflow_component
for upstream_node in node.pipeline_node.upstream_nodes:
assert (upstream_node in component_impl_map
), "Components is not in topological order"
current_airflow_component.set_upstream(
component_impl_map[upstream_node])
return airflow_dag
def _extract_deployment_config(
self, pipeline: pipeline_pb2.Pipeline
) -> local_deployment_config_pb2.LocalDeploymentConfig:
"""Extracts the proto.Any pipeline.deployment_config to LocalDeploymentConfig."""
return runner_utils.extract_local_deployment_config(pipeline)
def main(argv):
# Log to the container's stdout so Kubeflow Pipelines UI can display logs to
# the user.
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().setLevel(logging.INFO)
parser = argparse.ArgumentParser()
parser.add_argument('--pipeline_root', type=str, required=True)
parser.add_argument('--metadata_ui_path',
type=str,
required=False,
default='/mlpipeline-ui-metadata.json')
parser.add_argument('--kubeflow_metadata_config', type=str, required=True)
parser.add_argument('--tfx_ir', type=str, required=True)
parser.add_argument('--node_id', type=str, required=True)
# There might be multiple runtime parameters.
# `args.runtime_parameter` should become List[str] by using "append".
parser.add_argument('--runtime_parameter', type=str, action='append')
# TODO(b/196892362): Replace hooking with a more straightforward mechanism.
launcher._register_execution = _register_execution # pylint: disable=protected-access
args = parser.parse_args(argv)
tfx_ir = pipeline_pb2.Pipeline()
json_format.Parse(args.tfx_ir, tfx_ir)
_resolve_runtime_parameters(tfx_ir, args.runtime_parameter)
deployment_config = runner_utils.extract_local_deployment_config(tfx_ir)
kubeflow_metadata_config = kubeflow_pb2.KubeflowMetadataConfig()
json_format.Parse(args.kubeflow_metadata_config, kubeflow_metadata_config)
metadata_connection = metadata.Metadata(
_get_metadata_connection_config(kubeflow_metadata_config))
node_id = args.node_id
# Attach necessary labels to distinguish different runner and DSL.
# TODO(zhitaoli): Pass this from KFP runner side when the same container
# entrypoint can be used by a different runner.
with telemetry_utils.scoped_labels({
telemetry_utils.LABEL_TFX_RUNNER: 'kfp',
}):
custom_executor_operators = {
executable_spec_pb2.ContainerExecutableSpec:
kubernetes_executor_operator.KubernetesExecutorOperator
}
executor_spec = runner_utils.extract_executor_spec(
deployment_config, node_id)
custom_driver_spec = runner_utils.extract_custom_driver_spec(
deployment_config, node_id)
pipeline_node = _get_pipeline_node(tfx_ir, node_id)
component_launcher = launcher.Launcher(
pipeline_node=pipeline_node,
mlmd_connection=metadata_connection,
pipeline_info=tfx_ir.pipeline_info,
pipeline_runtime_spec=tfx_ir.runtime_spec,
executor_spec=executor_spec,
custom_driver_spec=custom_driver_spec,
custom_executor_operators=custom_executor_operators)
logging.info('Component %s is running.', node_id)
execution_info = component_launcher.launch()
logging.info('Component %s is finished.', node_id)
# Dump the UI metadata.
_dump_ui_metadata(pipeline_node, execution_info, args.metadata_ui_path)
def run(
self, pipeline: tfx_pipeline.Pipeline, run_name: Optional[str] = None
) -> None:
"""Runs given logical pipeline locally.
Args:
pipeline: Logical pipeline containing pipeline args and components.
run_name: Optional name for the run.
"""
for component in pipeline.components:
if isinstance(component, base_component.BaseComponent):
component._resolve_pip_dependencies(
pipeline.pipeline_info.pipeline_root
)
c = compiler.Compiler()
pipeline = c.compile(pipeline)
run_name = run_name or datetime.now().strftime("%d_%h_%y-%H_%M_%S_%f")
# Substitute the runtime parameter to be a concrete run_id
runtime_parameter_utils.substitute_runtime_parameter(
pipeline,
{
PIPELINE_RUN_ID_PARAMETER_NAME: run_name,
},
)
deployment_config = runner_utils.extract_local_deployment_config(
pipeline
)
connection_config = deployment_config.metadata_connection_config # type: ignore[attr-defined] # noqa
logger.debug(f"Using deployment config:\n {deployment_config}")
logger.debug(f"Using connection config:\n {connection_config}")
# Run each component. Note that the pipeline.components list is in
# topological order.
for node in pipeline.nodes:
pipeline_node = node.pipeline_node
node_id = pipeline_node.node_info.id
executor_spec = runner_utils.extract_executor_spec(
deployment_config, node_id
)
custom_driver_spec = runner_utils.extract_custom_driver_spec(
deployment_config, node_id
)
component_launcher = launcher.Launcher(
pipeline_node=pipeline_node,
mlmd_connection=metadata.Metadata(connection_config),
pipeline_info=pipeline.pipeline_info,
pipeline_runtime_spec=pipeline.runtime_spec,
executor_spec=executor_spec,
custom_driver_spec=custom_driver_spec,
)
start = time.time()
logger.info(f"Step `{node_id}` has started.")
component_launcher.launch()
end = time.time()
logger.info(
f"Step `{node_id}` has finished"
f" in {format_timedelta_pretty(end - start)}."
)
def main():
# Log to the container's stdout so Kubeflow Pipelines UI can display logs to
# the user.
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().setLevel(logging.INFO)
parser = argparse.ArgumentParser()
parser.add_argument('--pipeline_root', type=str, required=True)
parser.add_argument('--kubeflow_metadata_config', type=str, required=True)
parser.add_argument('--serialized_component', type=str, required=True)
parser.add_argument('--tfx_ir', type=str, required=True)
parser.add_argument('--node_id', type=str, required=True)
launcher._register_execution = _register_execution # pylint: disable=protected-access
args = parser.parse_args()
tfx_ir = pipeline_pb2.Pipeline()
json_format.Parse(args.tfx_ir, tfx_ir)
# Substitute the runtime parameter to be a concrete run_id
runtime_parameter_utils.substitute_runtime_parameter(
tfx_ir, {
constants.PIPELINE_RUN_ID_PARAMETER_NAME:
os.environ['WORKFLOW_ID'],
})
deployment_config = runner_utils.extract_local_deployment_config(tfx_ir)
kubeflow_metadata_config = kubeflow_pb2.KubeflowMetadataConfig()
json_format.Parse(args.kubeflow_metadata_config, kubeflow_metadata_config)
metadata_connection = kubeflow_metadata_adapter.KubeflowMetadataAdapter(
_get_metadata_connection_config(kubeflow_metadata_config))
node_id = args.node_id
# Attach necessary labels to distinguish different runner and DSL.
# TODO(zhitaoli): Pass this from KFP runner side when the same container
# entrypoint can be used by a different runner.
with telemetry_utils.scoped_labels({
telemetry_utils.LABEL_TFX_RUNNER: 'kfp',
}):
custom_executor_operators = {
executable_spec_pb2.ContainerExecutableSpec:
kubernetes_executor_operator.KubernetesExecutorOperator
}
executor_spec = runner_utils.extract_executor_spec(
deployment_config, node_id)
custom_driver_spec = runner_utils.extract_custom_driver_spec(
deployment_config, node_id)
pipeline_node = _get_pipeline_node(tfx_ir, node_id)
component_launcher = launcher.Launcher(
pipeline_node=pipeline_node,
mlmd_connection=metadata_connection,
pipeline_info=tfx_ir.pipeline_info,
pipeline_runtime_spec=tfx_ir.runtime_spec,
executor_spec=executor_spec,
custom_driver_spec=custom_driver_spec,
custom_executor_operators=custom_executor_operators)
logging.info('Component %s is running.', node_id)
execution_info = component_launcher.launch()
logging.info('Component %s is finished.', node_id)
# Dump the UI metadata.
_dump_ui_metadata(pipeline_node, execution_info)
def run_with_ir(
self,
pipeline: pipeline_pb2.Pipeline,
run_options: Optional[pipeline_pb2.RunOptions] = None,
) -> None:
"""Runs given pipeline locally.
Args:
pipeline: Pipeline IR containing pipeline args and components.
run_options: Optional args for the run.
Raises:
ValueError: If run_options is provided, and partial_run_options.from_nodes
and partial_run_options.to_nodes are both empty.
"""
# Substitute the runtime parameter to be a concrete run_id
runtime_parameter_utils.substitute_runtime_parameter(
pipeline, {
constants.PIPELINE_RUN_ID_PARAMETER_NAME:
datetime.datetime.now().isoformat(),
})
deployment_config = runner_utils.extract_local_deployment_config(pipeline)
connection_config = getattr(
deployment_config.metadata_connection_config,
deployment_config.metadata_connection_config.WhichOneof(
'connection_config'))
logging.info('Using deployment config:\n %s', deployment_config)
logging.info('Using connection config:\n %s', connection_config)
if run_options:
logging.info('Using run_options:\n %s', run_options)
pr_opts = run_options.partial_run
partial_run_utils.mark_pipeline(
pipeline,
from_nodes=pr_opts.from_nodes or None,
to_nodes=pr_opts.to_nodes or None,
snapshot_settings=pr_opts.snapshot_settings)
with telemetry_utils.scoped_labels(
{telemetry_utils.LABEL_TFX_RUNNER: 'local'}):
# Run each component. Note that the pipeline.components list is in
# topological order.
#
# TODO(b/171319478): After IR-based execution is used, used multi-threaded
# execution so that independent components can be run in parallel.
for node in pipeline.nodes:
pipeline_node = node.pipeline_node
node_id = pipeline_node.node_info.id
if pipeline_node.execution_options.HasField('skip'):
logging.info('Skipping component %s.', node_id)
continue
executor_spec = runner_utils.extract_executor_spec(
deployment_config, node_id)
custom_driver_spec = runner_utils.extract_custom_driver_spec(
deployment_config, node_id)
component_launcher = launcher.Launcher(
pipeline_node=pipeline_node,
mlmd_connection=metadata.Metadata(connection_config),
pipeline_info=pipeline.pipeline_info,
pipeline_runtime_spec=pipeline.runtime_spec,
executor_spec=executor_spec,
custom_driver_spec=custom_driver_spec)
logging.info('Component %s is running.', node_id)
if pipeline_node.execution_options.run.perform_snapshot:
with metadata.Metadata(connection_config) as mlmd_handle:
partial_run_utils.snapshot(mlmd_handle, pipeline)
component_launcher.launch()
logging.info('Component %s is finished.', node_id)
