File size: 7,110 Bytes
74da989 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | from abc import ABC, abstractmethod
from typing import Any, Optional
from trainer.trainer_args import ExecutionMode, TrainerArgs
from trainer.logger import init_logger
from trainer.pipelines import ComposedPipelineBase, build_pipeline
from trainer.pipelines.pipeline_registry import PipelineType
logger = init_logger(__name__)
class WorkflowBase(ABC):
"""
Abstract base class for defining video processing workflows.
A workflow serves as the top-level orchestrator that coordinates multiple pipelines
and components to accomplish a specific video processing task. The workflow pattern
provides several key benefits:
1. **Separation of Concerns**: Workflows separate high-level orchestration logic
from low-level processing implementations in pipelines.
2. **Modularity**: Different workflows can be created for different execution modes
(preprocess, inference, etc.) while sharing common pipeline components.
3. **Configuration Management**: Workflows manage the configuration and initialization
of multiple related pipelines and components in a centralized manner.
4. **Environment Setup**: Workflows handle system-level setup and resource
allocation before pipeline execution begins.
5. **Lifecycle Management**: Workflows control the complete lifecycle from
initialization through execution to cleanup.
The workflow acts as a factory and coordinator, creating the appropriate pipelines
based on configuration, setting up the execution environment, and orchestrating
the overall processing flow.
"""
def __init__(self, trainer_args: TrainerArgs):
"""
Initialize the workflow with configuration arguments.
Args:
trainer_args: Configuration object containing all parameters
needed for workflow and pipeline setup.
"""
self.trainer_args = trainer_args
# TODO: pipeline_config should be: dict[str, PipelineConfig]
# pipeline_type should be included in the PipelineConfig
# pipeline_config[pipeline_name] = (pipeline_type, trainer_args)
self._pipeline_configs: dict[str, tuple[PipelineType,
TrainerArgs]] = {}
self._pipelines: dict[str, ComposedPipelineBase] = {}
self._components: dict[str, Any] = {}
self.register_pipelines()
self.register_components()
self.prepare_system_environment()
self.load_pipelines()
def load_pipelines(self) -> None:
"""
Create and initialize all registered pipelines.
This method instantiates pipeline objects from their configurations
and makes them available as both dictionary entries and instance
attributes for convenient access.
"""
for pipeline_name, pipeline_config in self._pipeline_configs.items():
pipeline_type, trainer_args = pipeline_config
pipeline = build_pipeline(trainer_args, pipeline_type)
self._pipelines[pipeline_name] = pipeline
setattr(self, pipeline_name, pipeline)
def add_pipeline_config(
self, pipeline_name: str,
pipeline_config: tuple[PipelineType, TrainerArgs]) -> None:
"""
Register a pipeline configuration for later instantiation.
Args:
pipeline_name: Unique identifier for the pipeline.
pipeline_config: Tuple containing the pipeline type and
configuration arguments.
"""
self._pipeline_configs[pipeline_name] = pipeline_config
def add_component(self, component_name: str, component: Any) -> None:
"""
Register a component instance with the workflow.
Components are auxiliary objects that may be shared across pipelines
or used for workflow-level functionality (e.g., databases, caches,
external services).
Args:
component_name: Unique identifier for the component.
component: The component instance to register.
"""
self._components[component_name] = component
setattr(self, component_name, component)
def get_component(self, component_name: str) -> Any:
"""
Retrieve a registered component by name.
Args:
component_name: The name of the component to retrieve.
Returns:
The component instance.
"""
return self._components[component_name]
@abstractmethod
def register_components(self) -> None:
"""
Register workflow-specific components.
Subclasses must implement this method to register any components
needed for their specific workflow (e.g., databases, external APIs,
shared resources).
"""
pass
@abstractmethod
def register_pipelines(self) -> None:
"""
Register workflow-specific pipelines.
Subclasses must implement this method to define which pipelines
are needed for their specific workflow and how they should be
configured.
"""
pass
@abstractmethod
def prepare_system_environment(self) -> None:
"""
Prepare the system environment for workflow execution.
Subclasses must implement this method to handle any system-level
setup required before pipeline execution (e.g., GPU initialization,
temporary directories, resource allocation).
"""
pass
@abstractmethod
def run(self):
"""
Execute the main workflow logic.
Subclasses must implement this method to define the specific
execution flow for their workflow, coordinating the registered
pipelines and components to accomplish the desired task.
"""
pass
@classmethod
def get_workflow_cls(
cls, trainer_args: TrainerArgs) -> Optional["WorkflowBase"]:
"""
Factory method to get the appropriate workflow class based on execution mode.
This method acts as a workflow factory, returning the appropriate
workflow class implementation based on the specified execution mode
in the configuration arguments.
Args:
trainer_args: Configuration object containing the execution mode
and other parameters.
Returns:
The appropriate workflow class for the specified execution mode,
or None if no workflow is available for the given mode.
"""
if trainer_args.mode == ExecutionMode.PREPROCESS:
from trainer.workflow.preprocess.preprocess_workflow import (
PreprocessWorkflow)
return PreprocessWorkflow.get_workflow_cls(trainer_args)
else:
raise ValueError(
f"Execution mode: {trainer_args.mode} is not supported in workflow."
)
|