from copy import copy
from dataclasses import dataclass
from typing import List, Callable, Optional, Dict, Any, Sequence, Union
from ..data.local_dataset import LocalDataSet
from ..models.local_model_zoo import ModelBase, LocalModelZoo
from ..testsuite.local_testsuite import LocalTestSuite
from ..utils.results import ValidationResult
@dataclass
class Parameter(object):
name: str
description: str
value: Union[list, dict]
@dataclass
class Node(object):
name: str
func: Callable
parent: Optional[Sequence[str]] = None
func_inputs: Optional[Dict[str, Any]] = None
func_outputs: Optional[List[str]] = None
status: str = "pending" # Node Status (pending, running, failed, success)
save_data: bool = False
save_model: bool = False
save_testsuite: bool = False
title: str = None
[docs]
class Pipeline(object):
"""
Pipeline controller.
Pipeline is a DAG of base steps, each step will be executed.
Parameters
----------
name : str, default="modeva-pipeline"
Name of the pipeline.
save : bool, default=False
Whether to save the pipeline results.
"""
def __init__(self, name: str = 'modeva-pipeline', save: bool = False) -> None:
self.name = name
self.save = save
self._params = {}
self._nodes = {}
self._adj = {} # graph adjacent list representation
self._stack = []
[docs]
def add_step(
self,
name: str,
func: Callable,
func_inputs: Optional[Dict[str, Any]] = None,
parent: Optional[Sequence[str]] = None,
save_data: bool = False,
save_model: bool = False,
save_testsuite: bool = False,
title: str = None
):
"""Add a step to pipeline.
Parameters
----------
name : str or Path
The name of the pipeline.
func : Callable
The callable function for this step.
func_inputs : Optional[Dict[str, Any]], default = None
Inputs for the callable function.
parent : Optional[Sequence[str]] = None
Parent step of this step
save_data : bool, default = False
Whether to register the data object to database.
save_model : bool, default = False
Whether to register the model object to database.
save_testsuite : bool, default = False
Whether to register the validation result object to database.
title : str, default = None
Title of this step.
"""
n = Node(name=name, parent=parent, func=func, func_inputs=func_inputs,
save_data=save_data, save_model=save_model,
save_testsuite=save_testsuite, title=title)
self._nodes[name] = n
self._adj[name] = []
if parent:
if isinstance(parent, str):
self._verify_parent_node(parent)
self._adj[parent].append(n)
elif isinstance(parent, list):
for p in parent:
self._verify_parent_node(p)
self._adj[p].append(n)
def _verify_parent_node(self, parent):
if parent not in self._adj:
raise RuntimeError(f"Could not find parent step '{parent}'. Please add parent step first.")
def _topologicalSortUtil(self, v, visited):
# Mark the current node as visited
visited.add(v)
# Recur for all adjacent vertices
for u in reversed(self._adj[v]):
if not u.name in visited:
self._topologicalSortUtil(u.name, visited)
# Push current vertex to stack which stores the result
self._stack.append(v)
# Function to perform Topological Sort
def _topologicalSort(self):
visited = set()
# Call the recursive helper function to store
# Topological Sort starting from all vertices one by
# one
for name, nodes in self._adj.items():
if not name in visited:
self._topologicalSortUtil(name, visited)
[docs]
def run(self):
"""Run the pipeline.
"""
self.ds = None
self.model_zoo = None
self.test_result_run_id = None
if self._verify_dag():
raise RuntimeError("Found cycle in the piepline. Please make sure your pipeline does not contain cycle.")
# TODO: parallel computing
self._topologicalSort()
while self._stack:
node_name = self._stack.pop()
node = self._nodes[node_name]
print(f"Executing step: {node_name}")
# gather inputs from parent nodes and kwargs from func_inputs argument
args = []
if node.parent:
if isinstance(node.parent, str):
args = args + self._nodes[node.parent].func_outputs
elif isinstance(node.parent, list):
for p in node.parent:
args = args + self._nodes[p].func_outputs
# execute callable func
node.func_outputs = node.func(*args, **node.func_inputs)
# flatten func outputs into list
if isinstance(node.func_outputs, tuple):
node.func_outputs = list(node.func_outputs)
else:
node.func_outputs = [node.func_outputs]
for item in node.func_outputs:
if isinstance(item, LocalDataSet):
self.ds = item
if node.save_data:
item.register(name=self.name, override=True)
elif isinstance(item, ModelBase):
## TODO: here self.ds must be non-None, which means users need to return ds in one of the steps.
self.model_zoo = LocalModelZoo(dataset=self.ds,
experiment_name=self.name + "-" + "Models")
self.model_zoo.add_model(item)
if node.save_model:
self.model_zoo.register(name=item.name)
elif isinstance(item, ValidationResult):
ts = LocalTestSuite(name=self.name + "-" + "TestSuite")
names = ts.list_registered_tests().Name.tolist()
if node.save_testsuite:
if item.key not in names:
ts.register(name=item.key, test_result=item)
else:
cnt = 0
for name in names:
if name.startswith(item.key):
cnt += 1
ky = item.key + f"_v{cnt}"
ts.register(name=ky, test_result=item)
def _verify_dag(self) -> bool:
"""
return: True iff the pipeline dag is fully accessible and contains no cycles
"""
visited = set()
prev_visited = None
while prev_visited != visited:
prev_visited = copy(visited)
for k, node in list(self._nodes.items()):
if k in visited:
continue
if any(p == node.name for p in node.parent or []):
# node cannot have itself as parent
return False
if not all(p in visited for p in node.parent or []):
continue
visited.add(k)
# return False if we did not cover all the nodes
return not bool(set(self._nodes.keys()) - visited)