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