def _finish_node_execution(self, use_task_queue, exec_node_task): """Simulates successful execution of a node.""" test_utils.fake_execute_node(self._mlmd_connection, exec_node_task) if use_task_queue: dequeued_task = self._task_queue.dequeue() self._task_queue.task_done(dequeued_task) self.assertEqual(exec_node_task.task_id, dequeued_task.task_id)
def test_cached_execution(self): """Tests that cached execution is used if one is available.""" # Fake ExampleGen run. example_gen_exec = otu.fake_example_gen_run(self._mlmd_connection, self._example_gen, 1, 1) # Invoking generator should produce an ExecNodeTask for StatsGen. [stats_gen_task] = self._generate_and_test( False, num_initial_executions=1, num_tasks_generated=1, num_new_executions=1, num_active_executions=1) self.assertEqual('my_statistics_gen', stats_gen_task.node_uid.node_id) # Finish StatsGen execution. otu.fake_execute_node(self._mlmd_connection, stats_gen_task) # Prepare another pipeline with a new pipeline_run_id. pipeline_run_id = str(uuid.uuid4()) new_pipeline = self._make_pipeline(self._pipeline_root, pipeline_run_id) with self._mlmd_connection as m: contexts = m.store.get_contexts_by_execution(example_gen_exec.id) # We use node context as cache context for ease of testing. cache_context = [c for c in contexts if c.name == 'my_example_gen'][0] # Fake example_gen cached execution. otu.fake_cached_execution(self._mlmd_connection, cache_context, otu.get_node(new_pipeline, 'my_example_gen')) stats_gen = otu.get_node(new_pipeline, 'my_statistics_gen') # Invoking generator for the new pipeline should result in: # 1. StatsGen execution succeeds with state "CACHED" but no ExecNodeTask # generated. # 2. An ExecNodeTask is generated for SchemaGen (component downstream of # StatsGen) with an active execution in MLMD. [schema_gen_task] = self._generate_and_test( False, pipeline=new_pipeline, num_initial_executions=3, num_tasks_generated=1, num_new_executions=2, num_active_executions=1) self.assertEqual('my_schema_gen', schema_gen_task.node_uid.node_id) # Check that StatsGen execution is successful in state "CACHED". with self._mlmd_connection as m: executions = task_gen_utils.get_executions(m, stats_gen) self.assertLen(executions, 1) execution = executions[0] self.assertTrue(execution_lib.is_execution_successful(execution)) self.assertEqual(metadata_store_pb2.Execution.CACHED, execution.last_known_state)
def _finish_node_execution(self, use_task_queue, exec_node_task, artifact_custom_properties=None): """Simulates successful execution of a node.""" self._start_processing(use_task_queue, exec_node_task) test_utils.fake_execute_node( self._mlmd_connection, exec_node_task, artifact_custom_properties=artifact_custom_properties) self._finish_processing(use_task_queue, exec_node_task)
def test_pipeline_succeeds_when_terminal_nodes_succeed( self, use_task_queue, fail_fast): """Tests that pipeline is finalized only after terminal nodes are successful. Args: use_task_queue: If task queue is enabled, new tasks are only generated if a task with the same task_id does not already exist in the queue. `use_task_queue=False` is useful to test the case of task generation when task queue is empty (for eg: due to orchestrator restart). fail_fast: If `True`, pipeline is aborted immediately if any node fails. """ # Check the expected terminal nodes. layers = sptg._topsorted_layers(self._pipeline) self.assertEqual( { self._example_validator.node_info.id, self._chore_b.node_info.id, # evaluator execution will be skipped as it is run conditionally and # the condition always evaluates to False in the current test. self._evaluator.node_info.id, }, sptg._terminal_node_ids(layers)) # Start executing the pipeline: test_utils.fake_example_gen_run(self._mlmd_connection, self._example_gen, 1, 1) self._run_next(use_task_queue, expect_nodes=[self._stats_gen]) self._run_next(use_task_queue, expect_nodes=[self._schema_gen]) # Both example-validator and transform are ready to execute. [example_validator_task, transform_task] = self._generate(use_task_queue, True, fail_fast=fail_fast) self.assertEqual(self._example_validator.node_info.id, example_validator_task.node_uid.node_id) self.assertEqual(self._transform.node_info.id, transform_task.node_uid.node_id) # Start processing (but do not finish) example-validator. self._start_processing(use_task_queue, example_validator_task) # But finish transform which is in the same layer. self._finish_node_execution(use_task_queue, transform_task) # Readability note: below, example-validator task should continue to be # generated when not using task queue because the execution is active. # Trainer and downstream nodes can execute as transform is finished. self._run_next(use_task_queue, expect_nodes=[self._trainer] if use_task_queue else [self._example_validator, self._trainer], finish_nodes=[self._trainer], fail_fast=fail_fast) self._run_next(use_task_queue, expect_nodes=[self._chore_a] if use_task_queue else [self._example_validator, self._chore_a], finish_nodes=[self._chore_a], fail_fast=fail_fast) self._run_next(use_task_queue, expect_nodes=[self._chore_b] if use_task_queue else [self._example_validator, self._chore_b], finish_nodes=[self._chore_b], fail_fast=fail_fast) self._run_next( use_task_queue, expect_nodes=[] if use_task_queue else [self._example_validator], finish_nodes=[], fail_fast=fail_fast) # FinalizePipelineTask is generated only after example-validator finishes. test_utils.fake_execute_node(self._mlmd_connection, example_validator_task) self._finish_processing(use_task_queue, example_validator_task) [finalize_task] = self._generate(use_task_queue, True, fail_fast=fail_fast) self.assertTrue(task_lib.is_finalize_pipeline_task(finalize_task)) self.assertEqual(status_lib.Code.OK, finalize_task.status.code)