def test_get_dag(self):
operations = [OperationFactory() for _ in range(4)]
operations[0].upstream_operations.set(operations[2:])
operations[1].upstream_operations.set(operations[2:])
operation_by_ids = {op.id: op for op in operations}
def get_downstream(op):
return op.downstream_operations.values_list('id', flat=True)
assert dags.get_dag(nodes=operations, downstream_fn=get_downstream) == (
{
operations[0].id: set(),
operations[1].id: set(),
operations[2].id: {operations[0].id, operations[1].id},
operations[3].id: {operations[0].id, operations[1].id},
},
operation_by_ids
)
# Add operations outside the dag
operation1 = OperationFactory()
operation1.downstream_operations.set([operations[1], operations[2], operations[3]])
operation2 = OperationFactory()
operation2.upstream_operations.set([operations[0], operations[2]])
assert dags.get_dag(nodes=operations, downstream_fn=get_downstream) == (
{
operations[0].id: {operation2.id, },
operations[1].id: set(),
operations[2].id: {operations[0].id, operations[1].id, operation2.id},
operations[3].id: {operations[0].id, operations[1].id},
},
operation_by_ids
)
def dag(self):
"""Construct the DAG of this pipeline based on the its operations and their downstream."""
operations = self.operations.all().prefetch_related('downstream_operations')
def get_downstream(op):
return op.downstream_operations.values_list('id', flat=True)
return dags.get_dag(operations, get_downstream)
def dag(self):
"""Construct the DAG of this pipeline based on the its operations and their downstream."""
from pipelines import dags
operations = self.operations.all().prefetch_related('downstream_operations')
def get_downstream(op):
return op.downstream_operations.values_list('id', flat=True)
return dags.get_dag(operations, get_downstream)
def dag(self) -> Tuple[Dict, Dict]:
"""Construct the DAG of this pipeline run
based on the its operation runs and their downstream.
"""
from pipelines import dags
operation_runs = self.operation_runs.all().prefetch_related('downstream_runs')
def get_downstream(op_run):
return op_run.downstream_runs.values_list('id', flat=True)
return dags.get_dag(operation_runs, get_downstream)
