def setUp(self):
super().setUp()
pipeline_root = os.path.join(
os.environ.get('TEST_UNDECLARED_OUTPUTS_DIR', self.get_temp_dir()),
self.id())
# Makes sure multiple connections within a test always connect to the same
# MLMD instance.
metadata_path = os.path.join(pipeline_root, 'metadata', 'metadata.db')
connection_config = metadata.sqlite_metadata_connection_config(
metadata_path)
connection_config.sqlite.SetInParent()
self._mlmd_connection = metadata.Metadata(
connection_config=connection_config)
self._testdata_dir = os.path.join(os.path.dirname(__file__),
'testdata')
# Sets up pipelines
pipeline = pipeline_pb2.Pipeline()
self.load_proto_from_text(
os.path.join(os.path.dirname(__file__), 'testdata',
'pipeline_for_resolver_test.pbtxt'), pipeline)
self._pipeline_info = pipeline.pipeline_info
self._pipeline_runtime_spec = pipeline.runtime_spec
runtime_parameter_utils.substitute_runtime_parameter(
pipeline, {
constants.PIPELINE_RUN_ID_PARAMETER_NAME: 'my_pipeline_run',
})
# Extracts components
self._my_trainer = pipeline.nodes[0].pipeline_node
self._resolver_node = pipeline.nodes[1].pipeline_node
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 _make_pipeline(self, pipeline_root, pipeline_run_id):
pipeline = test_pipeline_with_resolver.create_pipeline()
runtime_parameter_utils.substitute_runtime_parameter(
pipeline, {
constants.PIPELINE_ROOT_PARAMETER_NAME: pipeline_root,
constants.PIPELINE_RUN_ID_PARAMETER_NAME: pipeline_run_id,
})
return pipeline
def _create_sync_pipeline(pipeline_id: str, run_id: str):
pipeline = _test_pipeline(pipeline_id, pipeline_pb2.Pipeline.SYNC)
pipeline.runtime_spec.pipeline_run_id.runtime_parameter.name = 'pipeline_run_id'
pipeline.runtime_spec.pipeline_run_id.runtime_parameter.type = (
pipeline_pb2.RuntimeParameter.STRING)
runtime_parameter_utils.substitute_runtime_parameter(
pipeline, {
'pipeline_run_id': run_id,
})
return pipeline
def _create_sync_pipeline(pipeline_id: str, run_id: str):
pipeline = _test_pipeline(pipeline_id, pipeline_pb2.Pipeline.SYNC)
pipeline.runtime_spec.pipeline_run_id.runtime_parameter.name = (
constants.PIPELINE_RUN_ID_PARAMETER_NAME)
pipeline.runtime_spec.pipeline_run_id.runtime_parameter.type = (
pipeline_pb2.RuntimeParameter.STRING)
runtime_parameter_utils.substitute_runtime_parameter(
pipeline, {
constants.PIPELINE_RUN_ID_PARAMETER_NAME: run_id,
})
return pipeline
def _test_pipeline(pipeline_id: str, run_id: str):
"""Creates test pipeline with pipeline_id and run_id."""
pipeline = pipeline_pb2.Pipeline()
path = os.path.join(
os.path.dirname(__file__), 'testdata', 'sync_pipeline.pbtxt')
io_utils.parse_pbtxt_file(path, pipeline)
pipeline.pipeline_info.id = pipeline_id
runtime_parameter_utils.substitute_runtime_parameter(pipeline, {
'pipeline_run_id': run_id,
})
return pipeline
def _make_pipeline(self, pipeline_root, pipeline_run_id):
pipeline = pipeline_pb2.Pipeline()
self.load_proto_from_text(
os.path.join(os.path.dirname(__file__), 'testdata',
'pipeline_with_importer.pbtxt'), pipeline)
runtime_parameter_utils.substitute_runtime_parameter(
pipeline, {
'pipeline_root': pipeline_root,
'pipeline_run_id': pipeline_run_id,
})
return pipeline
def setUp(self):
super(LauncherTest, self).setUp()
pipeline_root = os.path.join(
os.environ.get('TEST_UNDECLARED_OUTPUTS_DIR', self.get_temp_dir()),
self.id())
# Set a constant version for artifact version tag.
patcher = mock.patch('tfx.version.__version__')
patcher.start()
tfx_version.__version__ = '0.123.4.dev'
self.addCleanup(patcher.stop)
# Makes sure multiple connections within a test always connect to the same
# MLMD instance.
metadata_path = os.path.join(pipeline_root, 'metadata', 'metadata.db')
connection_config = metadata.sqlite_metadata_connection_config(
metadata_path)
connection_config.sqlite.SetInParent()
self._mlmd_connection = metadata.Metadata(
connection_config=connection_config)
# Sets up pipelines
pipeline = pipeline_pb2.Pipeline()
self.load_proto_from_text(
os.path.join(os.path.dirname(__file__), 'testdata',
'pipeline_for_launcher_test.pbtxt'), pipeline)
# Substitute the runtime parameter to be a concrete run_id
runtime_parameter_utils.substitute_runtime_parameter(
pipeline, {
constants.PIPELINE_RUN_ID_PARAMETER_NAME: 'test_run',
})
self._pipeline_info = pipeline.pipeline_info
self._pipeline_runtime_spec = pipeline.runtime_spec
self._pipeline_runtime_spec.pipeline_root.field_value.string_value = (
pipeline_root)
self._pipeline_runtime_spec.pipeline_run_id.field_value.string_value = (
'test_run_0')
# Extracts components
self._example_gen = pipeline.nodes[0].pipeline_node
self._transform = pipeline.nodes[1].pipeline_node
self._trainer = pipeline.nodes[2].pipeline_node
self._importer = pipeline.nodes[3].pipeline_node
self._resolver = pipeline.nodes[4].pipeline_node
# Fakes an ExecutorSpec for Trainer
self._trainer_executor_spec = _PYTHON_CLASS_EXECUTABLE_SPEC()
# Fakes an executor operator
self._test_executor_operators = {
_PYTHON_CLASS_EXECUTABLE_SPEC: _FakeExecutorOperator
}
# Fakes an custom driver spec
self._custom_driver_spec = _PYTHON_CLASS_EXECUTABLE_SPEC()
self._custom_driver_spec.class_path = 'tfx.orchestration.portable.launcher_test._FakeExampleGenLikeDriver'
def testSubstituteRuntimeParameterFail(self):
pipeline = pipeline_pb2.Pipeline()
self.load_proto_from_text('pipeline_with_runtime_parameter.pbtxt',
pipeline)
with self.assertRaisesRegex(RuntimeError, 'Runtime parameter type'):
runtime_parameter_utils.substitute_runtime_parameter(
pipeline,
{
'context_name_rp': 0, # Wrong type, will lead to failure.
'prop_one_rp': 2,
'prop_two_rp': 'X'
})
def _test_pipeline(ir_path: str, pipeline_id: str, run_id: str,
deployment_config: Optional[message.Message]):
"""Creates test pipeline with pipeline_id and run_id."""
pipeline = pipeline_pb2.Pipeline()
io_utils.parse_pbtxt_file(ir_path, pipeline)
pipeline.pipeline_info.id = pipeline_id
runtime_parameter_utils.substitute_runtime_parameter(pipeline, {
constants.PIPELINE_RUN_ID_PARAMETER_NAME: run_id,
})
if deployment_config:
pipeline.deployment_config.Pack(deployment_config)
return pipeline
def testPartiallySubstituteRuntimeParameter(self):
pipeline = pipeline_pb2.Pipeline()
expected = pipeline_pb2.Pipeline()
self.load_proto_from_text('pipeline_with_runtime_parameter.pbtxt',
pipeline)
self.load_proto_from_text(
'pipeline_with_runtime_parameter_partially_substituted.pbtxt',
expected)
runtime_parameter_utils.substitute_runtime_parameter(
pipeline, {
'context_name_rp': 'my_context',
})
self.assertProtoEquals(pipeline, expected)
def setUp(self):
super(SyncPipelineTaskGeneratorTest, self).setUp()
pipeline_root = os.path.join(
os.environ.get('TEST_UNDECLARED_OUTPUTS_DIR', self.get_temp_dir()),
self.id())
self._pipeline_root = pipeline_root
# Makes sure multiple connections within a test always connect to the same
# MLMD instance.
metadata_path = os.path.join(pipeline_root, 'metadata', 'metadata.db')
self._metadata_path = metadata_path
connection_config = metadata.sqlite_metadata_connection_config(
metadata_path)
connection_config.sqlite.SetInParent()
self._mlmd_connection = metadata.Metadata(
connection_config=connection_config)
# Sets up the pipeline.
pipeline = pipeline_pb2.Pipeline()
self.load_proto_from_text(
os.path.join(
os.path.dirname(__file__), 'testdata', 'sync_pipeline.pbtxt'),
pipeline)
self._pipeline_run_id = str(uuid.uuid4())
runtime_parameter_utils.substitute_runtime_parameter(
pipeline, {
'pipeline_root': pipeline_root,
'pipeline_run_id': self._pipeline_run_id
})
self._pipeline = pipeline
# Extracts components.
self._example_gen = _get_node(pipeline, 'my_example_gen')
self._stats_gen = _get_node(pipeline, 'my_statistics_gen')
self._schema_gen = _get_node(pipeline, 'my_schema_gen')
self._transform = _get_node(pipeline, 'my_transform')
self._example_validator = _get_node(pipeline, 'my_example_validator')
self._trainer = _get_node(pipeline, 'my_trainer')
self._task_queue = tq.TaskQueue()
self._mock_service_job_manager = mock.create_autospec(
service_jobs.ServiceJobManager, instance=True)
self._mock_service_job_manager.is_pure_service_node.side_effect = (
lambda _, node_id: node_id == self._example_gen.node_info.id)
self._mock_service_job_manager.is_mixed_service_node.side_effect = (
lambda _, node_id: node_id == self._transform.node_info.id)
def testFullySubstituteRuntimeParameter(self):
pipeline = pipeline_pb2.Pipeline()
expected = pipeline_pb2.Pipeline()
self.load_proto_from_text(
os.path.join(self._testdata_dir,
'pipeline_with_runtime_parameter.pbtxt'), pipeline)
self.load_proto_from_text(
os.path.join(self._testdata_dir,
'pipeline_with_runtime_parameter_substituted.pbtxt'),
expected)
runtime_parameter_utils.substitute_runtime_parameter(
pipeline, {
'context_name_rp': 'my_context',
'prop_one_rp': 2,
'prop_two_rp': 'X'
})
self.assertProtoEquals(pipeline, expected)
def _resolve_runtime_parameters(tfx_ir: pipeline_pb2.Pipeline,
parameters: Optional[List[str]]) -> None:
"""Resolve runtime parameters in the pipeline proto inplace."""
if parameters is None:
return
parameter_bindings = {
# Substitute the runtime parameter to be a concrete run_id
constants.PIPELINE_RUN_ID_PARAMETER_NAME:
os.environ['WORKFLOW_ID'],
}
# Argo will fill runtime parameter values in the parameters.
for param in parameters:
name, value = _parse_runtime_parameter_str(param)
parameter_bindings[name] = value
runtime_parameter_utils.substitute_runtime_parameter(
tfx_ir, parameter_bindings)
def reloadPipelineWithNewRunId(self):
self._pipeline_run_id_counter += 1
# Substitute the runtime parameter to be a new run_id.
pipeline = pipeline_pb2.Pipeline()
pipeline.CopyFrom(self._pipeline)
runtime_parameter_utils.substitute_runtime_parameter(
pipeline, {
constants.PIPELINE_RUN_ID_PARAMETER_NAME:
f'test_run_{self._pipeline_run_id_counter}',
})
self._pipeline_runtime_spec.pipeline_run_id.field_value.string_value = (
f'test_run_{self._pipeline_run_id_counter}')
# Extracts components
self._example_gen = pipeline.nodes[0].pipeline_node
self._transform = pipeline.nodes[1].pipeline_node
self._trainer = pipeline.nodes[2].pipeline_node
self._importer = pipeline.nodes[3].pipeline_node
self._resolver = pipeline.nodes[4].pipeline_node
def initiate_pipeline_start(
mlmd_handle: metadata.Metadata,
pipeline: pipeline_pb2.Pipeline) -> pstate.PipelineState:
"""Initiates a pipeline start operation.
Upon success, MLMD is updated to signal that the pipeline must be started.
Args:
mlmd_handle: A handle to the MLMD db.
pipeline: IR of the pipeline to start.
Returns:
The `PipelineState` object upon success.
Raises:
status_lib.StatusNotOkError: Failure to initiate pipeline start. With code
`INVALILD_ARGUMENT` if a sync pipeline is not configured for runtime
parameter substitution of `pipeline_run_id`.
"""
pipeline = copy.deepcopy(pipeline)
if pipeline.execution_mode == pipeline_pb2.Pipeline.SYNC:
if not pipeline.runtime_spec.pipeline_run_id.HasField(
'runtime_parameter'):
raise status_lib.StatusNotOkError(
code=status_lib.Code.INVALID_ARGUMENT,
message=
('Sync pipeline IR must be configured for runtime substitution of '
'pipeline_run_id'))
pipeline_run_id = datetime.datetime.now().strftime('%Y%m%d-%H%M%S-%f')
runtime_parameter_utils.substitute_runtime_parameter(
pipeline,
{constants.PIPELINE_RUN_ID_PARAMETER_NAME: pipeline_run_id})
assert (pipeline.runtime_spec.pipeline_run_id.field_value.string_value
== pipeline_run_id)
with pstate.PipelineState.new(mlmd_handle, pipeline) as pipeline_state:
pass
return pipeline_state
def testPartiallySubstituteRuntimeParameter(self):
pipeline = pipeline_pb2.Pipeline()
expected = pipeline_pb2.Pipeline()
self.load_proto_from_text(
os.path.join(self._testdata_dir,
'pipeline_with_runtime_parameter.pbtxt'), pipeline)
self.load_proto_from_text(
os.path.join(
self._testdata_dir,
'pipeline_with_runtime_parameter_partially_substituted.pbtxt'),
expected)
parameters = runtime_parameter_utils.substitute_runtime_parameter(
pipeline, {
'context_name_rp': 'my_context',
})
self.assertProtoEquals(pipeline, expected)
self.assertEqual(len(parameters), 3)
self.assertEqual(parameters['context_name_rp'], 'my_context')
self.assertEqual(parameters['prop_one_rp'], 1)
self.assertIsNone(parameters['prop_two_rp'])
def run(self, pipeline: Union[pipeline_pb2.Pipeline,
pipeline_py.Pipeline]) -> None:
"""Deploys given logical pipeline on Beam.
Args:
pipeline: Logical pipeline in IR format.
"""
# For CLI, while creating or updating pipeline, pipeline_args are extracted
# and hence we avoid deploying the pipeline.
if 'TFX_JSON_EXPORT_PIPELINE_ARGS_PATH' in os.environ:
return
if isinstance(pipeline, pipeline_py.Pipeline):
c = compiler.Compiler()
pipeline = c.compile(pipeline)
run_id = datetime.datetime.now().strftime('%Y%m%d-%H%M%S.%f')
# Substitute the runtime parameter to be a concrete run_id
runtime_parameter_utils.substitute_runtime_parameter(
pipeline, {
constants.PIPELINE_RUN_ID_PARAMETER_NAME: run_id,
})
deployment_config = self._extract_deployment_config(pipeline)
connection_config = self._connection_config_from_deployment_config(
deployment_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: 'beam'}):
with beam.Pipeline() as p:
# Uses for triggering the node DoFns.
root = p | 'CreateRoot' >> beam.Create([None])
# Stores mapping of node to its signal.
signal_map = {}
# pipeline.nodes are in topological order.
for node in pipeline.nodes:
# TODO(b/160882349): Support subpipeline
pipeline_node = node.pipeline_node
node_id = pipeline_node.node_info.id
executor_spec = self._extract_executor_spec(
deployment_config, node_id)
custom_driver_spec = self._extract_custom_driver_spec(
deployment_config, node_id)
# Signals from upstream nodes.
signals_to_wait = []
for upstream_node in pipeline_node.upstream_nodes:
assert upstream_node in signal_map, ('Nodes are not in '
'topological order')
signals_to_wait.append(signal_map[upstream_node])
logging.info('Node %s depends on %s.', node_id,
[s.producer.full_label for s in signals_to_wait])
# Each signal is an empty PCollection. AsIter ensures a node will
# be triggered after upstream nodes are finished.
signal_map[node_id] = (
root
| 'Run[%s]' % node_id >> beam.ParDo(
self._PIPELINE_NODE_DO_FN_CLS(
pipeline_node=pipeline_node,
mlmd_connection_config=connection_config,
pipeline_info=pipeline.pipeline_info,
pipeline_runtime_spec=pipeline.runtime_spec,
executor_spec=executor_spec,
custom_driver_spec=custom_driver_spec,
deployment_config=deployment_config),
*[beam.pvalue.AsIter(s) for s in signals_to_wait]))
logging.info('Node %s is scheduled.', node_id)
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)
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(
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 run_with_ir(
self,
pipeline: pipeline_pb2.Pipeline,
run_options: Optional[pipeline_pb2.RunOptions] = None,
) -> None:
"""Deploys given logical pipeline on Beam.
Args:
pipeline: Logical pipeline in IR format.
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.
"""
run_id = datetime.datetime.now().strftime('%Y%m%d-%H%M%S.%f')
# Substitute the runtime parameter to be a concrete run_id
runtime_parameter_utils.substitute_runtime_parameter(
pipeline, {
constants.PIPELINE_RUN_ID_PARAMETER_NAME: run_id,
})
deployment_config = self._extract_deployment_config(pipeline)
connection_config = self._connection_config_from_deployment_config(
deployment_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: 'beam'}):
with beam.Pipeline() as p:
# Uses for triggering the node DoFns.
root = p | 'CreateRoot' >> beam.Create([None])
# Stores mapping of node_id to its signal.
signal_map = {}
snapshot_node_id = None
# pipeline.nodes are in topological order.
for node in pipeline.nodes:
# TODO(b/160882349): Support subpipeline
pipeline_node = node.pipeline_node
node_id = pipeline_node.node_info.id
if pipeline_node.execution_options.HasField('skip'):
logging.info('Node %s is skipped in this partial run.',
node_id)
continue
run_opts = pipeline_node.execution_options.run
if run_opts.perform_snapshot:
snapshot_node_id = node_id
executor_spec = self._extract_executor_spec(
deployment_config, node_id)
custom_driver_spec = self._extract_custom_driver_spec(
deployment_config, node_id)
# Signals from upstream nodes.
signals_to_wait = []
# This ensures that nodes that rely on artifacts reused from previous
# pipelines runs are scheduled after the snapshot node.
if run_opts.depends_on_snapshot and node_id != snapshot_node_id:
signals_to_wait.append(signal_map[snapshot_node_id])
for upstream_node_id in pipeline_node.upstream_nodes:
if upstream_node_id in signal_map:
signals_to_wait.append(
signal_map[upstream_node_id])
else:
logging.info(
'Node %s is upstream of Node %s, but will be skipped in '
'this partial run. Node %s is responsible for ensuring that '
'Node %s\'s dependencies can be resolved.',
upstream_node_id, node_id, snapshot_node_id,
node_id)
logging.info(
'Node %s depends on %s.', node_id,
[s.producer.full_label for s in signals_to_wait])
# Each signal is an empty PCollection. AsIter ensures a node will
# be triggered after upstream nodes are finished.
signal_map[node_id] = (
root
| 'Run[%s]' % node_id >> beam.ParDo(
self._PIPELINE_NODE_DO_FN_CLS(
pipeline_node=pipeline_node,
mlmd_connection_config=connection_config,
pipeline_info=pipeline.pipeline_info,
pipeline_runtime_spec=pipeline.runtime_spec,
executor_spec=executor_spec,
custom_driver_spec=custom_driver_spec,
deployment_config=deployment_config,
pipeline=pipeline), *
[beam.pvalue.AsIter(s) for s in signals_to_wait]))
logging.info('Node %s is scheduled.', node_id)
def run(self, pipeline: pipeline_pb2.Pipeline) -> None:
"""Deploys given logical pipeline on Beam.
Args:
pipeline: Logical pipeline in IR format.
"""
# For CLI, while creating or updating pipeline, pipeline_args are extracted
# and hence we avoid deploying the pipeline.
if 'TFX_JSON_EXPORT_PIPELINE_ARGS_PATH' in os.environ:
return
run_id = datetime.datetime.now().isoformat()
# Substitute the runtime parameter to be a concrete run_id
runtime_parameter_utils.substitute_runtime_parameter(
pipeline, {
constants.PIPELINE_RUN_ID_PARAMETER_NAME: run_id,
})
# TODO(b/163003901): Support beam DAG runner args through IR.
deployment_config = self._extract_deployment_config(pipeline)
connection_config = deployment_config.metadata_connection_config
mlmd_connection = metadata.Metadata(
connection_config=connection_config)
with telemetry_utils.scoped_labels(
{telemetry_utils.LABEL_TFX_RUNNER: 'beam'}):
with beam.Pipeline() as p:
# Uses for triggering the component DoFns.
root = p | 'CreateRoot' >> beam.Create([None])
# Stores mapping of component to its signal.
signal_map = {}
# pipeline.components are in topological order.
for node in pipeline.nodes:
# TODO(b/160882349): Support subpipeline
pipeline_node = node.pipeline_node
component_id = pipeline_node.node_info.id
executor_spec = self._extract_executor_spec(
deployment_config, component_id)
custom_driver_spec = self._extract_custom_driver_spec(
deployment_config, component_id)
# Signals from upstream components.
signals_to_wait = []
for upstream_node in pipeline_node.upstream_nodes:
assert upstream_node in signal_map, (
'Components is not in '
'topological order')
signals_to_wait.append(signal_map[upstream_node])
logging.info(
'Component %s depends on %s.', component_id,
[s.producer.full_label for s in signals_to_wait])
# Each signal is an empty PCollection. AsIter ensures component will
# be triggered after upstream components are finished.
# LINT.IfChange
signal_map[component_id] = (
root
| 'Run[%s]' % component_id >> beam.ParDo(
self._PIPELINE_NODE_DO_FN_CLS(
pipeline_node=pipeline_node,
mlmd_connection=mlmd_connection,
pipeline_info=pipeline.pipeline_info,
pipeline_runtime_spec=pipeline.runtime_spec,
executor_spec=executor_spec,
custom_driver_spec=custom_driver_spec), *
[beam.pvalue.AsIter(s) for s in signals_to_wait]))
# LINT.ThenChange(../beam/beam_dag_runner.py)
logging.info('Component %s is scheduled.', component_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
