def log_component_run( component_run: ComponentRun, set_dependencies_from_inputs=True, staleness_threshold: int = (60 * 60 * 24 * 30), ): """Takes client-facing ComponentRun object and logs it to the DB.""" store = Store(_db_uri) # Make dictionary object component_run_dict = component_run.to_dictionary() component_run_sql = store.initialize_empty_component_run( component_run.component_name) # Add relevant attributes if component_run_dict["start_timestamp"]: component_run_sql.set_start_timestamp( component_run_dict["start_timestamp"]) if component_run_dict["end_timestamp"]: component_run_sql.set_end_timestamp( component_run_dict["end_timestamp"]) if component_run_dict["notes"]: component_run_sql.add_notes(component_run_dict["notes"]) component_run_sql.set_git_hash(component_run_dict["git_hash"]) component_run_sql.set_git_tags(component_run_dict["git_tags"]) component_run_sql.set_code_snapshot(component_run_dict["code_snapshot"]) # Add I/O component_run_sql.add_inputs([ store.get_io_pointer(inp.name, inp.value, pointer_type=inp.pointer_type) for inp in component_run_dict["inputs"] ]) component_run_sql.add_outputs([ store.get_io_pointer(out.name, out.value, pointer_type=out.pointer_type) for out in component_run_dict["outputs"] ]) # Create component if it does not exist create_component(component_run.component_name, "", "") # Add dependencies if there is flag to automatically set if set_dependencies_from_inputs: store.set_dependencies_from_inputs(component_run_sql) # Add dependencies explicitly stored in the component run for dependency in component_run_dict["dependencies"]: cr = store.get_history(dependency, 1)[0] component_run_sql.set_upstream(cr) store.commit_component_run(component_run_sql, staleness_threshold=staleness_threshold)
class TestComponentRun(unittest.TestCase): def setUp(self): self.mock_component_run = ComponentRun("mock_component_run") self.mock_component_run_dict = { "component_name": "mock_component_run", "notes": "", "inputs": [], "outputs": [], "git_hash": None, "git_tags": None, "code_snapshot": None, "start_timestamp": None, "end_timestamp": None, "dependencies": [], "id": None, "stale": [], "test_result": None, "mlflow_run_id": None, "mlflow_run_params": None, "mlflow_run_metrics": None, } self.mock_inputs = [ IOPointer("mock_input_1"), IOPointer("mock_input_2"), ] self.mock_outputs = [ IOPointer("mock_output_1"), IOPointer("mock_output_2"), ] def testSerialize(self): """ Test the serialization functionality. """ self.assertEqual( self.mock_component_run.to_dictionary(), self.mock_component_run_dict, ) def testSetStartEndError(self): """ Test that setting start and end ts as non datetime types throws an error. """ with self.assertRaises(TypeError): self.mock_component_run.set_start_timestamp("incorrect_type") with self.assertRaises(TypeError): self.mock_component_run.set_end_timestamp("incorrect_type") def testAddInputOutput(self): cr = copy.deepcopy(self.mock_component_run) for inp in self.mock_inputs: cr.add_input(inp) for out in self.mock_outputs: cr.add_output(out) self.assertEqual(cr.inputs, list(set(self.mock_inputs))) self.assertEqual(cr.outputs, list(set(self.mock_outputs))) def testAddInputsOutputs(self): cr = copy.deepcopy(self.mock_component_run) cr.add_inputs(self.mock_inputs) cr.add_outputs(self.mock_outputs) self.assertEqual(cr.inputs, list(set(self.mock_inputs))) self.assertEqual(cr.outputs, list(set(self.mock_outputs))) def testAddDuplicateInputs(self): cr = copy.deepcopy(self.mock_component_run) cr.add_inputs(self.mock_inputs) cr.add_inputs(self.mock_inputs) self.assertEqual(cr.inputs, list(set(self.mock_inputs))) def testAddNotes(self): cr = copy.deepcopy(self.mock_component_run) expected_output = "this is a test note" cr.notes = "this is a test note" self.assertEqual(cr.notes, expected_output) def testAddNotesError(self): """ Test that adding non-str input to the notes attribute gives a TypeError """ with self.assertRaises(TypeError): self.mock_component_run.notes = ["incorrect_type"]
class TestComponentRun(unittest.TestCase): def setUp(self): self.mock_component_run = ComponentRun("mock_component_run") self.mock_component_run_dict = { "component_name": "mock_component_run", "inputs": [], "outputs": [], "git_hash": None, "code_snapshot": None, "start_timestamp": None, "end_timestamp": None, "dependencies": [], "id": None, "stale": [], } self.mock_inputs = [IOPointer("mock_input_1"), IOPointer("mock_input_2")] self.mock_outputs = [IOPointer("mock_output_1"), IOPointer("mock_output_2")] def testSerialize(self): """ Test the serialization functionality. """ self.assertEqual( self.mock_component_run.to_dictionary(), self.mock_component_run_dict ) def testSetStartEndError(self): """ Test that setting start and end ts as non datetime types throws an error. """ with self.assertRaises(TypeError): self.mock_component_run.set_start_timestamp("incorrect_type") with self.assertRaises(TypeError): self.mock_component_run.set_end_timestamp("incorrect_type") def testAddInputOutput(self): cr = copy.deepcopy(self.mock_component_run) for inp in self.mock_inputs: cr.add_input(inp) for out in self.mock_outputs: cr.add_output(out) self.assertEqual(cr.inputs, list(set(self.mock_inputs))) self.assertEqual(cr.outputs, list(set(self.mock_outputs))) def testAddInputsOutputs(self): cr = copy.deepcopy(self.mock_component_run) cr.add_inputs(self.mock_inputs) cr.add_outputs(self.mock_outputs) self.assertEqual(cr.inputs, list(set(self.mock_inputs))) self.assertEqual(cr.outputs, list(set(self.mock_outputs))) def testAddDuplicateInputs(self): cr = copy.deepcopy(self.mock_component_run) cr.add_inputs(self.mock_inputs) cr.add_inputs(self.mock_inputs) self.assertEqual(cr.inputs, list(set(self.mock_inputs)))