Implemented LinearFlow for both local and prefect runner (and everywhere else)

Author: sveinugu

src/unifair/compute/flow.py

@@ -5,7 +5,10 @@
 from typing import Any, Callable, Dict, Optional, Tuple, Type, Union
 
 from unifair.compute.job import CallableDecoratingJobTemplateMixin, Job, JobConfig, JobTemplate
-from unifair.engine.protocols import IsDagFlowRunnerEngine, IsFuncFlowRunnerEngine, IsTaskTemplate
+from unifair.engine.protocols import (IsDagFlowRunnerEngine,
+                                      IsFuncFlowRunnerEngine,
+                                      IsLinearFlowRunnerEngine,
+                                      IsTaskTemplate)
 from unifair.util.callable_decorator_cls import callable_decorator_cls
 
 
@@ -50,6 +53,73 @@ def _call_func(self, *args: Any, **kwargs: Any) -> Any:
         pass
 
 
+class LinearFlowConfig(FlowConfig):
+    def __init__(
+        self,
+        linear_flow_func: Callable,
+        *task_templates: IsTaskTemplate,
+        name: Optional[str] = None,
+        **kwargs: Any,
+    ):
+        self._linear_flow_func = linear_flow_func
+        self._linear_flow_func_signature = inspect.signature(self._linear_flow_func)
+        self._task_templates: Tuple[IsTaskTemplate] = task_templates
+        name = name if name is not None else self._linear_flow_func.__name__
+        super().__init__(name=name, **kwargs)
+
+    def _get_init_arg_values(self) -> Union[Tuple[()], Tuple[Any, ...]]:
+        return self._linear_flow_func, *self._task_templates
+
+    def _get_init_kwarg_public_property_keys(self) -> Tuple[str, ...]:
+        return ()
+
+    @property
+    def task_templates(self) -> Tuple[IsTaskTemplate]:
+        return self._task_templates
+
+    def has_coroutine_func(self) -> bool:
+        return asyncio.iscoroutinefunction(self._linear_flow_func)
+
+    @property
+    def param_signatures(self) -> MappingProxyType:
+        return self._linear_flow_func_signature.parameters
+
+    @property
+    def return_type(self) -> Type[Any]:
+        return self._linear_flow_func_signature.return_annotation
+
+
+@callable_decorator_cls
+class LinearFlowTemplate(CallableDecoratingJobTemplateMixin['LinearFlowTemplate'],
+                         FlowTemplate,
+                         LinearFlowConfig):
+    @classmethod
+    def _get_job_subcls_for_apply(cls) -> Type[Job]:
+        return LinearFlow
+
+    def _apply_engine_decorator(self, flow: 'LinearFlow') -> 'LinearFlow':
+        if self.engine is not None and isinstance(self.engine, IsLinearFlowRunnerEngine):
+            return self.engine.linear_flow_decorator(flow)  # noqa  # Pycharm static type checker bug
+        else:
+            raise RuntimeError(f'Engine "{self.engine}" does not support DAG flows')
+
+
+class LinearFlow(Flow, LinearFlowConfig):
+    @classmethod
+    def _get_job_config_subcls_for_init(cls) -> Type[JobConfig]:
+        return LinearFlowConfig
+
+    @classmethod
+    def _get_job_template_subcls_for_revise(cls) -> Type[JobTemplate]:
+        return LinearFlowTemplate  # noqa  # Pycharm static type checker bug
+
+    def _call_func(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError
+
+    def get_call_args(self, *args: object, **kwargs: object) -> Dict[str, object]:
+        return inspect.signature(self._linear_flow_func).bind(*args, **kwargs).arguments
+
+
 class DagFlowConfig(FlowConfig):
     def __init__(
         self,

src/unifair/engine/job_runner.py

@@ -2,11 +2,11 @@
 import inspect
 import sys
 from types import AsyncGeneratorType, GeneratorType
-from typing import Any, Awaitable, Callable, cast
+from typing import Any, Awaitable, Callable, cast, Iterable
 
 from unifair.engine.base import Engine
 from unifair.engine.constants import RunState
-from unifair.engine.protocols import IsDagFlow, IsFuncFlow, IsJob, IsTask
+from unifair.engine.protocols import IsDagFlow, IsFuncFlow, IsJob, IsLinearFlow, IsTask
 
 
 class JobRunnerEngine(Engine, ABC):
@@ -81,6 +81,48 @@ def _run_task(self, state: Any, task: IsTask, call_func: Callable, *args, **kwar
         ...
 
 
+class LinearFlowRunnerEngine(JobRunnerEngine):
+    def linear_flow_decorator(self, linear_flow: IsLinearFlow) -> IsLinearFlow:
+        # prev_call_func = flow._call_func  # Only raises error anyway
+
+        self._register_job_state(linear_flow, RunState.INITIALIZED)
+        state = self._init_linear_flow(linear_flow)
+
+        def _call_func(*args: object, **kwargs: object) -> Any:
+            self._register_job_state(linear_flow, RunState.RUNNING)
+            flow_result = self._run_linear_flow(state, linear_flow, *args, **kwargs)
+            return self._decorate_result_with_job_finalization_detector(linear_flow, flow_result)
+
+        setattr(linear_flow, '_call_func', _call_func)
+        return linear_flow
+
+    @staticmethod
+    def default_linear_flow_run_decorator(linear_flow: IsLinearFlow) -> Any:
+        def _inner_run_linear_flow(*args: object, **kwargs: object):
+
+            result = None
+            for i, job in enumerate(linear_flow.task_templates):
+                with linear_flow.flow_context:
+                    # TODO: Better handling or kwargs
+                    if i == 0:
+                        result = job(*args, **kwargs)
+                    else:
+                        result = job(*args)
+
+                args = result if isinstance(result, Iterable) else [result]
+            return result
+
+        return _inner_run_linear_flow
+
+    @abstractmethod
+    def _init_linear_flow(self, linear_flow: IsLinearFlow) -> Any:
+        ...
+
+    @abstractmethod
+    def _run_linear_flow(self, state: Any, linear_flow: IsLinearFlow, *args, **kwargs) -> Any:
+        ...
+
+
 class DagFlowRunnerEngine(JobRunnerEngine):
     def dag_flow_decorator(self, dag_flow: IsDagFlow) -> IsDagFlow:
         # prev_call_func = flow._call_func  # Only raises error anyway

src/unifair/engine/local.py

@@ -1,11 +1,21 @@
 from typing import Any, Callable, Type
 
 from unifair.config.engine import LocalRunnerConfig
-from unifair.engine.job_runner import DagFlowRunnerEngine, FuncFlowRunnerEngine, TaskRunnerEngine
-from unifair.engine.protocols import IsDagFlow, IsFuncFlow, IsLocalRunnerConfig, IsTask
-
-
-class LocalRunner(TaskRunnerEngine, DagFlowRunnerEngine, FuncFlowRunnerEngine):
+from unifair.engine.job_runner import (DagFlowRunnerEngine,
+                                       FuncFlowRunnerEngine,
+                                       LinearFlowRunnerEngine,
+                                       TaskRunnerEngine)
+from unifair.engine.protocols import (IsDagFlow,
+                                      IsFuncFlow,
+                                      IsLinearFlow,
+                                      IsLocalRunnerConfig,
+                                      IsTask)
+
+
+class LocalRunner(TaskRunnerEngine,
+                  LinearFlowRunnerEngine,
+                  DagFlowRunnerEngine,
+                  FuncFlowRunnerEngine):
     def _init_engine(self) -> None:
         ...
 
@@ -22,6 +32,12 @@ def _init_task(self, task: IsTask, call_func: Callable) -> Any:
     def _run_task(self, state: Any, task: IsTask, call_func: Callable, *args, **kwargs) -> Any:
         return call_func(*args, **kwargs)
 
+    def _init_linear_flow(self, flow: IsLinearFlow) -> Any:
+        ...
+
+    def _run_linear_flow(self, state: Any, flow: IsLinearFlow, *args, **kwargs) -> Any:
+        return self.default_linear_flow_run_decorator(flow)(*args, **kwargs)
+
     def _init_dag_flow(self, flow: IsDagFlow) -> Any:
         ...
 

src/unifair/engine/prefect.py

@@ -9,12 +9,22 @@
 from prefect.tasks import task_input_hash
 
 from unifair.config.engine import PrefectEngineConfig
-from unifair.engine.job_runner import DagFlowRunnerEngine, FuncFlowRunnerEngine, TaskRunnerEngine
-from unifair.engine.protocols import IsDagFlow, IsFuncFlow, IsPrefectEngineConfig, IsTask
+from unifair.engine.job_runner import (DagFlowRunnerEngine,
+                                       FuncFlowRunnerEngine,
+                                       LinearFlowRunnerEngine,
+                                       TaskRunnerEngine)
+from unifair.engine.protocols import (IsDagFlow,
+                                      IsFuncFlow,
+                                      IsLinearFlow,
+                                      IsPrefectEngineConfig,
+                                      IsTask)
 from unifair.util.helpers import resolve
 
 
-class PrefectEngine(TaskRunnerEngine, DagFlowRunnerEngine, FuncFlowRunnerEngine):
+class PrefectEngine(TaskRunnerEngine,
+                    LinearFlowRunnerEngine,
+                    DagFlowRunnerEngine,
+                    FuncFlowRunnerEngine):
     def _init_engine(self) -> None:
         ...
 
@@ -71,6 +81,33 @@ def task_flow(*inner_args, **inner_kwargs):
 
             return task_flow(*args, **kwargs)
 
+    # LinearFlowRunnerEngine
+    def _init_linear_flow(self, linear_flow: IsLinearFlow) -> Any:
+        assert isinstance(self._config, PrefectEngineConfig)
+        flow_kwargs = dict(name=linear_flow.name,)
+        call_func = self.default_linear_flow_run_decorator(linear_flow)
+
+        if linear_flow.has_coroutine_func():
+
+            @prefect_flow(**flow_kwargs)
+            async def run_linear_flow(*inner_args, **inner_kwargs):
+                with linear_flow.flow_context:
+                    return await resolve(call_func(*inner_args, **inner_kwargs))
+        else:
+
+            @prefect_flow(**flow_kwargs)
+            def run_linear_flow(*inner_args, **inner_kwargs):
+                with linear_flow.flow_context:
+                    return call_func(*inner_args, **inner_kwargs)
+
+        return run_linear_flow
+
+    def _run_linear_flow(self, state: Any, linear_flow: IsLinearFlow, *args, **kwargs) -> Any:
+
+        _prefect_flow = state
+
+        return _prefect_flow(*args, **kwargs)
+
     # DagFlowRunnerEngine
 
     def _init_dag_flow(self, dag_flow: IsDagFlow) -> Any:

src/unifair/engine/protocols.py

@@ -136,6 +136,20 @@ def apply(self) -> IsFlow:
         ...
 
 
+class IsLinearFlow(IsFlow, Protocol):
+    task_templates: Tuple[IsTaskTemplate]
+
+    def has_coroutine_func(self) -> bool:
+        ...
+
+    def get_call_args(self, *args, **kwargs) -> Dict[str, object]:
+        ...
+
+
+class IsLinearFlowTemplate(IsLinearFlow, IsFlowTemplate, Protocol):
+    ...
+
+
 class IsDagFlow(IsFlow, Protocol):
     task_templates: Tuple[IsTaskTemplate]
 
@@ -178,6 +192,12 @@ def task_decorator(self, task: IsTask) -> IsTask:
         ...
 
 
+@runtime_checkable
+class IsLinearFlowRunnerEngine(IsEngine, Protocol):
+    def linear_flow_decorator(self, linear_flow: IsLinearFlow) -> IsLinearFlow:
+        ...
+
+
 @runtime_checkable
 class IsDagFlowRunnerEngine(IsEngine, Protocol):
     def dag_flow_decorator(self, dag_flow: IsDagFlow) -> IsDagFlow:

tests/compute/cases/decorators.py

@@ -1,6 +1,6 @@
 import pytest_cases as pc
 
-from unifair.compute.flow import DagFlowTemplate, FuncFlowTemplate
+from unifair.compute.flow import DagFlowTemplate, FuncFlowTemplate, LinearFlowTemplate
 from unifair.compute.task import TaskTemplate
 
 
@@ -31,6 +31,34 @@ def plus_other(number: int, other: int) -> int:
     return plus_other  # noqa  # Pycharm static type checker bug
 
 
+@pc.case(
+    id='linear_flow-plus_five(number)',
+    tags=['sync', 'function', 'linear_flow', 'plain'],
+)
+@pc.parametrize_with_cases('plus_one_template', cases='.', has_tag='task')
+def case_linear_flow_number_plus_five_template(plus_one_template) -> LinearFlowTemplate:
+    @LinearFlowTemplate(*((plus_one_template,) * 5))
+    def plus_five(number: int) -> int:  # noqa
+        ...
+
+    return plus_five  # noqa  # Pycharm static type checker bug
+
+
+@pc.case(
+    id='linear_flow-plus_five(x)',
+    tags=['sync', 'function', 'linear_flow', 'with_kw_params'],
+)
+@pc.parametrize_with_cases('plus_one_template', cases='.', has_tag='task')
+def case_linear_flow_x_plus_five_template(plus_one_template) -> FuncFlowTemplate:
+    iterative_x_plus_one_template = plus_one_template.refine(param_key_map=dict(number='x'),)
+
+    @LinearFlowTemplate(*((iterative_x_plus_one_template,) * 5))
+    def plus_five(x: int) -> int:  # noqa
+        ...
+
+    return plus_five  # noqa  # Pycharm static type checker bug
+
+
 @pc.case(
     id='dag_flow-plus_five(number)',
     tags=['sync', 'function', 'dag_flow', 'plain'],

tests/compute/cases/flows.py

@@ -5,7 +5,7 @@
 import pytest
 import pytest_cases as pc
 
-from unifair.compute.flow import DagFlowTemplate, FlowTemplate, FuncFlowTemplate
+from unifair.compute.flow import DagFlowTemplate, FlowTemplate, FuncFlowTemplate, LinearFlowTemplate
 from unifair.compute.task import TaskTemplate
 
 from .tasks import TaskCase
@@ -28,6 +28,26 @@ class FlowCase(Generic[ArgT, ReturnT]):
 # TODO: Add assert_signature_and_return_type_func
 
 
+@pc.case(
+    id='sync-linearflow-single-task',
+    tags=['sync', 'linearflow', 'singlethread', 'success'],
+)
+@pc.parametrize_with_cases('task_case', cases='.tasks')
+def case_sync_linearflow_single_task(task_case: TaskCase) -> FlowCase[[], None]:
+    task_template = TaskTemplate(task_case.task_func)
+    linear_flow = LinearFlowTemplate(task_template)(task_case.task_func)
+
+    return FlowCase(
+        name=task_case.name,
+        flow_func=task_case.task_func,
+        flow_template=linear_flow,
+        args=task_case.args,
+        kwargs=task_case.kwargs,
+        assert_results_func=task_case.assert_results_func,
+        # assert_signature_and_return_type_func=task_case.assert_signature_and_return_type_func,
+    )
+
+
 @pc.case(
     id='sync-dagflow-single-task',
     tags=['sync', 'dagflow', 'singlethread', 'success'],

tests/compute/helpers/mocks.py

@@ -2,10 +2,11 @@
 from typing import Any, Callable, Dict, Mapping, Optional, Tuple, Type, Union
 
 from unifair.compute.job import CallableDecoratingJobTemplateMixin, Job, JobConfig, JobTemplate
-from unifair.engine.job_runner import DagFlowRunnerEngine
+from unifair.engine.job_runner import DagFlowRunnerEngine, LinearFlowRunnerEngine
 from unifair.engine.protocols import (IsDagFlow,
                                       IsEngineConfig,
                                       IsFuncFlow,
+                                      IsLinearFlow,
                                       IsRunStateRegistry,
                                       IsTask)
 from unifair.util.callable_decorator_cls import callable_decorator_cls
@@ -213,6 +214,10 @@ def _call_func(*args: Any, **kwargs: Any) -> Any:
         setattr(task, '_call_func', _call_func)
         return task
 
+    def linear_flow_decorator(self, flow: IsLinearFlow) -> IsLinearFlow:  # noqa
+        setattr(flow, '_call_func', LinearFlowRunnerEngine.default_linear_flow_run_decorator(flow))
+        return flow
+
     def dag_flow_decorator(self, flow: IsDagFlow) -> IsDagFlow:  # noqa
         setattr(flow, '_call_func', DagFlowRunnerEngine.default_dag_flow_run_decorator(flow))
         return flow