From 2f9952cdbfaf823d7d54eedea012f06f53a44c25 Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 12 Nov 2025 11:51:24 +0800 Subject: [PATCH] diffusion: remove unused workflows folder (#13114) --- .../runtime/workflow/__init__.py | 1 - .../runtime/workflow/preprocess/__init__.py | 1 - .../runtime/workflow/preprocess/components.py | 341 ------------------ .../preprocess/preprocess_workflow.py | 140 ------- .../preprocess/preprocess_workflow_i2v.py | 70 ---- .../preprocess/preprocess_workflow_t2v.py | 70 ---- .../runtime/workflow/workflow_base.py | 188 ---------- 7 files changed, 811 deletions(-) delete mode 100644 python/sglang/multimodal_gen/runtime/workflow/__init__.py delete mode 100644 python/sglang/multimodal_gen/runtime/workflow/preprocess/__init__.py delete mode 100644 python/sglang/multimodal_gen/runtime/workflow/preprocess/components.py delete mode 100644 python/sglang/multimodal_gen/runtime/workflow/preprocess/preprocess_workflow.py delete mode 100644 python/sglang/multimodal_gen/runtime/workflow/preprocess/preprocess_workflow_i2v.py delete mode 100644 python/sglang/multimodal_gen/runtime/workflow/preprocess/preprocess_workflow_t2v.py delete mode 100644 python/sglang/multimodal_gen/runtime/workflow/workflow_base.py diff --git a/python/sglang/multimodal_gen/runtime/workflow/__init__.py b/python/sglang/multimodal_gen/runtime/workflow/__init__.py deleted file mode 100644 index af2eb7d10..000000000 --- a/python/sglang/multimodal_gen/runtime/workflow/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo diff --git a/python/sglang/multimodal_gen/runtime/workflow/preprocess/__init__.py b/python/sglang/multimodal_gen/runtime/workflow/preprocess/__init__.py deleted file mode 100644 index af2eb7d10..000000000 --- a/python/sglang/multimodal_gen/runtime/workflow/preprocess/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo diff --git a/python/sglang/multimodal_gen/runtime/workflow/preprocess/components.py b/python/sglang/multimodal_gen/runtime/workflow/preprocess/components.py deleted file mode 100644 index 1f7890aec..000000000 --- a/python/sglang/multimodal_gen/runtime/workflow/preprocess/components.py +++ /dev/null @@ -1,341 +0,0 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo - -import dataclasses -import gc -import os -import random -from collections.abc import Callable -from typing import Any - -import numpy as np -import pyarrow as pa -import torch -from datasets import Dataset, Video, load_dataset - -from sglang.multimodal_gen.configs.configs import ( - DatasetType, - PreprocessConfig, - VideoLoaderType, -) -from sglang.multimodal_gen.configs.sample.base import DataType -from sglang.multimodal_gen.dataset.dataloader.parquet_io import ( - ParquetDatasetWriter, - records_to_table, -) -from sglang.multimodal_gen.runtime.distributed.parallel_state import ( - get_world_rank, - get_world_size, -) -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import PreprocessBatch -from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger - -logger = init_logger(__name__) - - -class PreprocessingDataValidator: - - def __init__( - self, - max_height: int = 1024, - max_width: int = 1024, - max_h_div_w_ratio: float = 17 / 16, - min_h_div_w_ratio: float = 8 / 16, - num_frames: int = 16, - train_fps: int = 24, - speed_factor: float = 1.0, - video_length_tolerance_range: float = 5.0, - drop_short_ratio: float = 0.0, - hw_aspect_threshold: float = 1.5, - ): - self.max_height = max_height - self.max_width = max_width - self.max_h_div_w_ratio = max_h_div_w_ratio - self.min_h_div_w_ratio = min_h_div_w_ratio - self.num_frames = num_frames - self.train_fps = train_fps - self.speed_factor = speed_factor - self.video_length_tolerance_range = video_length_tolerance_range - self.drop_short_ratio = drop_short_ratio - self.hw_aspect_threshold = hw_aspect_threshold - self.validators: dict[str, Callable[[dict[str, Any]], bool]] = {} - self.filter_counts: dict[str, int] = {} - - self.num_items_before_filtering = 0 - self.num_items_after_filtering = 0 - - self.register_validators() - - def register_validators(self) -> None: - self.add_validator("data_type_validator", self._validate_data_type) - self.add_validator("resolution_validator", self._validate_resolution) - self.add_validator("frame_sampling_validator", self._validate_frame_sampling) - - def add_validator( - self, name: str, validator: Callable[[dict[str, Any]], bool] - ) -> None: - self.validators[name] = validator - self.filter_counts[name] = 0 - - def __call__(self, batch: dict[str, Any]) -> bool: - """ - Validate whether the preprocessing data batch is valid. - """ - self.num_items_before_filtering += 1 - - for name, validator in self.validators.items(): - if not validator(batch): - self.filter_counts[name] += 1 - return False - - self.num_items_after_filtering += 1 - return True - - def _validate_data_type(self, batch: dict[str, Any]) -> bool: - """Validate basic validity of data items""" - return not ( - batch["caption"] is None - or batch["caption"] == "" - or batch["fps"] is None - or batch["fps"] <= 0 - or batch["num_frames"] is None - or batch["num_frames"] <= 0 - ) - - def _validate_resolution(self, batch: dict[str, Any]) -> bool: - """Validate resolution constraints""" - - aspect = self.max_height / self.max_width - if batch["resolution"] is not None: - height = batch["resolution"].get("height", None) - width = batch["resolution"].get("width", None) - - if height is None or width is None: - return False - - return self._filter_resolution( - height, - width, - max_h_div_w_ratio=self.hw_aspect_threshold * aspect, - min_h_div_w_ratio=1 / self.hw_aspect_threshold * aspect, - ) - - def _filter_resolution( - self, h: int, w: int, max_h_div_w_ratio: float, min_h_div_w_ratio: float - ) -> bool: - """Filter based on aspect ratio""" - return (min_h_div_w_ratio <= h / w <= max_h_div_w_ratio) and ( - self.min_h_div_w_ratio <= h / w <= self.max_h_div_w_ratio - ) - - def _validate_frame_sampling(self, batch: dict[str, Any]) -> bool: - """Validate frame sampling constraints""" - - if batch["num_frames"] / batch["fps"] > self.video_length_tolerance_range * ( - self.num_frames / self.train_fps * self.speed_factor - ): - return False - - frame_interval = batch["fps"] / self.train_fps - start_frame_idx = 0 - frame_indices = np.arange( - start_frame_idx, batch["num_frames"], frame_interval - ).astype(int) - return not ( - len(frame_indices) < self.num_frames - and random.random() < self.drop_short_ratio - ) - - def log_validation_stats(self): - info = "" - for name, count in self.filter_counts.items(): - info += f"failed in {name}: {count}, " - info += f"number of items before filtering: {self.num_items_before_filtering}, " - info += f"number of items after filtering: {self.num_items_after_filtering}" - - logger.info(info) - - -class VideoForwardBatchBuilder: - - def __init__(self, seed: int): - self.seed = seed - - def __call__(self, batch: list) -> PreprocessBatch: - forward_batch = PreprocessBatch( - video_loader=[item["video"] for item in batch], - video_file_name=[item["name"] for item in batch], - height=[item["resolution"]["height"] for item in batch], - width=[item["resolution"]["width"] for item in batch], - fps=[item["fps"] for item in batch], - num_frames=[item["num_frames"] for item in batch], - prompt=[item["caption"] for item in batch], - prompt_attention_mask=[], - data_type=DataType.VIDEO, - generator=torch.Generator("cpu").manual_seed(self.seed), - ) - return forward_batch - - -class ParquetDatasetSaver: - """Component for saving and writing Parquet datasets using shared parquet_io.""" - - def __init__( - self, - flush_frequency: int, - samples_per_file: int, - schema: pa.Schema, - record_creator: Callable[..., list[dict[str, Any]]], - ): - self.flush_frequency = flush_frequency - self.samples_per_file = samples_per_file - self.schema = schema - self.create_records_from_batch = record_creator - self.num_processed_samples: int = 0 - self._writer: ParquetDatasetWriter | None = None - - def save_and_write_parquet_batch( - self, - batch: PreprocessBatch, - output_dir: str, - extra_features: dict[str, Any] | None = None, - ) -> None: - """ - Save and write Parquet dataset batch - - Args: - batch: PreprocessBatch containing video and metadata information - output_dir: Output directory - extra_features: Extra features - - Returns: - Number of processed samples - """ - assert isinstance(batch.latents, torch.Tensor) - assert isinstance(batch.prompt_embeds, list) - assert isinstance(batch.prompt_attention_mask, list) - - # Process non-padded embeddings (if needed) - if batch.prompt_attention_mask is not None: - batch.prompt_embeds = self._process_non_padded_embeddings( - batch.prompt_embeds[0], batch.prompt_attention_mask[0] - ) - else: - raise ValueError("prompt_attention_mask is None") - - # Prepare batch data for Parquet dataset - batch_data: list[dict[str, Any]] = [] - - for key in dataclasses.fields(batch): - value = getattr(batch, key.name) - if isinstance(value, list): - for idx in range(len(value)): - if isinstance(value[idx], torch.Tensor): - value[idx] = value[idx].cpu().numpy() - elif isinstance(value, torch.Tensor): - value = value.cpu().numpy() - setattr(batch, key.name, value) - - # Create record for Parquet dataset - records = self.create_records_from_batch(batch) - batch_data.extend(records) - - if batch_data: - self.num_processed_samples += len(batch_data) - table = records_to_table(batch_data, self.schema) - if self._writer is None: - os.makedirs(output_dir, exist_ok=True) - self._writer = ParquetDatasetWriter( - out_dir=output_dir, samples_per_file=self.samples_per_file - ) - self._writer.append_table(table) - logger.debug("Collected batch with %s samples", len(table)) - - # If flush is needed - if self.num_processed_samples >= self.flush_frequency: - self.flush_tables() - - def _process_non_padded_embeddings( - self, prompt_embeds: torch.Tensor, prompt_attention_mask: torch.Tensor - ) -> list[torch.Tensor]: - """Process non-padded embeddings""" - assert isinstance(prompt_embeds, torch.Tensor) - assert isinstance(prompt_attention_mask, torch.Tensor) - assert prompt_embeds.shape[0] == prompt_attention_mask.shape[0] - - # Get sequence lengths from attention masks (number of 1s) - seq_lens = prompt_attention_mask.sum(dim=1) - - non_padded_embeds = [] - - # Process each item in the batch - for i in range(prompt_embeds.size(0)): - seq_len = seq_lens[i].item() - # Slice the embeddings and masks to keep only non-padding parts - non_padded_embeds.append(prompt_embeds[i, :seq_len]) - - return non_padded_embeds - - def flush_tables(self, write_remainder: bool = False): - """Flush buffered records to disk. - - Args: - output_dir: Directory where parquet files are written. Kept for API - symmetry (writer already configured with this path). - write_remainder: If True, also write any leftover rows smaller than - ``samples_per_file`` as a final small file. Useful for the last flush. - """ - if self._writer is None: - return - _ = self._writer.flush(write_remainder=write_remainder) - # Reset processed sample count modulo samples_per_file - remainder = self.num_processed_samples % self.samples_per_file - self.num_processed_samples = 0 if write_remainder else remainder - - def clean_up(self) -> None: - """Clean up all tables""" - self.flush_tables(write_remainder=True) - self._writer = None - self.num_processed_samples = 0 - gc.collect() - - def __del__(self): - self.clean_up() - - -def build_dataset( - preprocess_config: PreprocessConfig, - split: str, - validator: Callable[[dict[str, Any]], bool], -) -> Dataset: - if preprocess_config.dataset_type == DatasetType.HF: - dataset = load_dataset(preprocess_config.dataset_path, split=split) - dataset = dataset.filter(validator) - dataset = dataset.shard(num_shards=get_world_size(), index=get_world_rank()) - elif preprocess_config.dataset_type == DatasetType.MERGED: - metadata_json_path = os.path.join( - preprocess_config.dataset_path, "videos2caption.json" - ) - video_folder = os.path.join(preprocess_config.dataset_path, "videos") - dataset = load_dataset("json", data_files=metadata_json_path, split=split) - column_names = dataset.column_names - # rename columns to match the schema - if "cap" in column_names: - dataset = dataset.rename_column("cap", "caption") - if "path" in column_names: - dataset = dataset.rename_column("path", "name") - - dataset = dataset.filter(validator) - dataset = dataset.shard(num_shards=get_world_size(), index=get_world_rank()) - - # add video column - def add_video_column(item: dict[str, Any]) -> dict[str, Any]: - item["video"] = os.path.join(video_folder, item["name"]) - return item - - dataset = dataset.map(add_video_column) - if preprocess_config.video_loader_type == VideoLoaderType.TORCHCODEC: - dataset = dataset.cast_column("video", Video()) - else: - raise ValueError(f"Invalid dataset type: {preprocess_config.dataset_type}") - - return dataset diff --git a/python/sglang/multimodal_gen/runtime/workflow/preprocess/preprocess_workflow.py b/python/sglang/multimodal_gen/runtime/workflow/preprocess/preprocess_workflow.py deleted file mode 100644 index 38fba4214..000000000 --- a/python/sglang/multimodal_gen/runtime/workflow/preprocess/preprocess_workflow.py +++ /dev/null @@ -1,140 +0,0 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo - -import os -from typing import cast - -from torch.utils.data import DataLoader - -from sglang.multimodal_gen.configs.configs import PreprocessConfig -from sglang.multimodal_gen.dataset.dataloader.record_schema import ( - basic_t2v_record_creator, - i2v_record_creator, -) -from sglang.multimodal_gen.dataset.dataloader.schema import ( - pyarrow_schema_i2v, - pyarrow_schema_t2v, -) -from sglang.multimodal_gen.runtime.distributed.parallel_state import get_world_rank -from sglang.multimodal_gen.runtime.server_args import ServerArgs, WorkloadType -from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger -from sglang.multimodal_gen.runtime.workflow.preprocess.components import ( - ParquetDatasetSaver, - PreprocessingDataValidator, - VideoForwardBatchBuilder, - build_dataset, -) -from sglang.multimodal_gen.runtime.workflow.workflow_base import WorkflowBase - -logger = init_logger(__name__) - - -class PreprocessWorkflow(WorkflowBase): - - def register_pipelines(self) -> None: - self.add_pipeline_config("preprocess_pipeline", self.server_args) - - def register_components(self) -> None: - assert self.server_args.preprocess_config is not None - preprocess_config: PreprocessConfig = self.server_args.preprocess_config - - # raw data validator - raw_data_validator = PreprocessingDataValidator( - max_height=preprocess_config.max_height, - max_width=preprocess_config.max_width, - num_frames=preprocess_config.num_frames, - train_fps=preprocess_config.train_fps, - speed_factor=preprocess_config.speed_factor, - video_length_tolerance_range=preprocess_config.video_length_tolerance_range, - drop_short_ratio=preprocess_config.drop_short_ratio, - ) - self.add_component("raw_data_validator", raw_data_validator) - - # training dataset - training_dataset = build_dataset( - preprocess_config, split="train", validator=raw_data_validator - ) - # we do not use collate_fn here because we use iterable-style Dataset - # and want to keep the original type of the dataset - training_dataloader = DataLoader( - training_dataset, - batch_size=preprocess_config.preprocess_video_batch_size, - num_workers=preprocess_config.dataloader_num_workers, - collate_fn=lambda x: x, - ) - self.add_component("training_dataloader", training_dataloader) - - # try to load validation dataset if it exists - try: - validation_dataset = build_dataset( - preprocess_config, split="validation", validator=raw_data_validator - ) - validation_dataloader = DataLoader( - validation_dataset, - batch_size=preprocess_config.preprocess_video_batch_size, - num_workers=preprocess_config.dataloader_num_workers, - collate_fn=lambda x: x, - ) - except ValueError: - logger.warning( - "Validation dataset not found, skipping validation dataset preprocessing." - ) - validation_dataloader = None - - self.add_component("validation_dataloader", validation_dataloader) - - # forward batch builder - video_forward_batch_builder = VideoForwardBatchBuilder( - seed=self.server_args.preprocess_config.seed - ) - self.add_component("video_forward_batch_builder", video_forward_batch_builder) - - # record creator - if self.server_args.workload_type == WorkloadType.I2V: - record_creator = i2v_record_creator - schema = pyarrow_schema_i2v - else: - record_creator = basic_t2v_record_creator - schema = pyarrow_schema_t2v - processed_dataset_saver = ParquetDatasetSaver( - flush_frequency=self.server_args.preprocess_config.flush_frequency, - samples_per_file=self.server_args.preprocess_config.samples_per_file, - schema=schema, - record_creator=record_creator, - ) - self.add_component("processed_dataset_saver", processed_dataset_saver) - - def prepare_system_environment(self) -> None: - assert self.server_args.preprocess_config is not None - dataset_output_dir = self.server_args.preprocess_config.dataset_output_dir - os.makedirs(dataset_output_dir, exist_ok=True) - - validation_dataset_output_dir = os.path.join( - dataset_output_dir, "validation_dataset", f"worker_{get_world_rank()}" - ) - os.makedirs(validation_dataset_output_dir, exist_ok=True) - self.validation_dataset_output_dir = validation_dataset_output_dir - - training_dataset_output_dir = os.path.join( - dataset_output_dir, "training_dataset", f"worker_{get_world_rank()}" - ) - os.makedirs(training_dataset_output_dir, exist_ok=True) - self.training_dataset_output_dir = training_dataset_output_dir - - @classmethod - def get_workflow_cls(cls, server_args: ServerArgs) -> "PreprocessWorkflow": - if server_args.workload_type == WorkloadType.T2V: - from sglang.multimodal_gen.runtime.workflow.preprocess.preprocess_workflow_t2v import ( - PreprocessWorkflowT2V, - ) - - return cast(PreprocessWorkflow, PreprocessWorkflowT2V) - elif server_args.workload_type == WorkloadType.I2V: - from sglang.multimodal_gen.runtime.workflow.preprocess.preprocess_workflow_i2v import ( - PreprocessWorkflowI2V, - ) - - return cast(PreprocessWorkflow, PreprocessWorkflowI2V) - else: - raise ValueError( - f"Workload type: {server_args.workload_type} is not supported in preprocessing workflow." - ) diff --git a/python/sglang/multimodal_gen/runtime/workflow/preprocess/preprocess_workflow_i2v.py b/python/sglang/multimodal_gen/runtime/workflow/preprocess/preprocess_workflow_i2v.py deleted file mode 100644 index 876d7bd41..000000000 --- a/python/sglang/multimodal_gen/runtime/workflow/preprocess/preprocess_workflow_i2v.py +++ /dev/null @@ -1,70 +0,0 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo - -from typing import TYPE_CHECKING - -from tqdm import tqdm - -from sglang.multimodal_gen.dataset.preprocessing_datasets import PreprocessBatch -from sglang.multimodal_gen.runtime.workflow.preprocess.components import ( - ParquetDatasetSaver, -) -from sglang.multimodal_gen.runtime.workflow.preprocess.preprocess_workflow import ( - PreprocessWorkflow, -) - -if TYPE_CHECKING: - from torch.utils.data import DataLoader - - from sglang.multimodal_gen.runtime.pipelines.composed_pipeline_base import ( - ComposedPipelineBase, - ) - from sglang.multimodal_gen.runtime.workflow.preprocess.components import ( - VideoForwardBatchBuilder, - ) - - -class PreprocessWorkflowI2V(PreprocessWorkflow): - training_dataloader: "DataLoader" - validation_dataloader: "DataLoader" - preprocess_pipeline: "ComposedPipelineBase" - processed_dataset_saver: "ParquetDatasetSaver" - video_forward_batch_builder: "VideoForwardBatchBuilder" - - def run(self) -> None: - # Training dataset preprocessing - for batch in tqdm( - self.training_dataloader, - desc="Preprocessing training dataset", - unit="batch", - ): - forward_batch: PreprocessBatch = self.video_forward_batch_builder(batch) - - forward_batch = self.preprocess_pipeline.forward( - forward_batch, self.server_args - ) - - self.processed_dataset_saver.save_and_write_parquet_batch( - forward_batch, self.training_dataset_output_dir - ) - - self.processed_dataset_saver.flush_tables() - self.processed_dataset_saver.clean_up() - - # Validation dataset preprocessing - if self.validation_dataloader is not None: - for batch in tqdm( - self.validation_dataloader, - desc="Preprocessing validation dataset", - unit="batch", - ): - forward_batch = self.video_forward_batch_builder(batch) - - forward_batch = self.preprocess_pipeline.forward( - forward_batch, self.server_args - ) - - self.processed_dataset_saver.save_and_write_parquet_batch( - forward_batch, self.validation_dataset_output_dir - ) - self.processed_dataset_saver.flush_tables() - self.processed_dataset_saver.clean_up() diff --git a/python/sglang/multimodal_gen/runtime/workflow/preprocess/preprocess_workflow_t2v.py b/python/sglang/multimodal_gen/runtime/workflow/preprocess/preprocess_workflow_t2v.py deleted file mode 100644 index b8f1df011..000000000 --- a/python/sglang/multimodal_gen/runtime/workflow/preprocess/preprocess_workflow_t2v.py +++ /dev/null @@ -1,70 +0,0 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo - -from typing import TYPE_CHECKING, Optional - -from tqdm import tqdm - -from sglang.multimodal_gen.runtime.pipelines.pipeline_batch_info import PreprocessBatch -from sglang.multimodal_gen.runtime.workflow.preprocess.components import ( - ParquetDatasetSaver, -) -from sglang.multimodal_gen.runtime.workflow.preprocess.preprocess_workflow import ( - PreprocessWorkflow, -) - -if TYPE_CHECKING: - from torch.utils.data import DataLoader - - from sglang.multimodal_gen.runtime.pipelines.composed_pipeline_base import ( - ComposedPipelineBase, - ) - from sglang.multimodal_gen.runtime.workflow.preprocess.components import ( - VideoForwardBatchBuilder, - ) - - -class PreprocessWorkflowT2V(PreprocessWorkflow): - training_dataloader: "DataLoader" - validation_dataloader: Optional["DataLoader"] - preprocess_pipeline: "ComposedPipelineBase" - processed_dataset_saver: "ParquetDatasetSaver" - video_forward_batch_builder: "VideoForwardBatchBuilder" - - def run(self) -> None: - # Training dataset preprocessing - for batch in tqdm( - self.training_dataloader, - desc="Preprocessing training dataset", - unit="batch", - ): - forward_batch: PreprocessBatch = self.video_forward_batch_builder(batch) - - forward_batch = self.preprocess_pipeline.forward( - forward_batch, self.server_args - ) - - self.processed_dataset_saver.save_and_write_parquet_batch( - forward_batch, self.training_dataset_output_dir - ) - - self.processed_dataset_saver.flush_tables() - self.processed_dataset_saver.clean_up() - - # Validation dataset preprocessing - if self.validation_dataloader is not None: - for batch in tqdm( - self.validation_dataloader, - desc="Preprocessing validation dataset", - unit="batch", - ): - forward_batch = self.video_forward_batch_builder(batch) - - forward_batch = self.preprocess_pipeline.forward( - forward_batch, self.server_args - ) - - self.processed_dataset_saver.save_and_write_parquet_batch( - forward_batch, self.validation_dataset_output_dir - ) - self.processed_dataset_saver.flush_tables() - self.processed_dataset_saver.clean_up() diff --git a/python/sglang/multimodal_gen/runtime/workflow/workflow_base.py b/python/sglang/multimodal_gen/runtime/workflow/workflow_base.py deleted file mode 100644 index 88522f69e..000000000 --- a/python/sglang/multimodal_gen/runtime/workflow/workflow_base.py +++ /dev/null @@ -1,188 +0,0 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo - -from abc import ABC, abstractmethod -from typing import Any, Optional - -from sglang.multimodal_gen.runtime.pipelines import ComposedPipelineBase, build_pipeline -from sglang.multimodal_gen.runtime.pipelines.pipeline_registry import PipelineType -from sglang.multimodal_gen.runtime.server_args import ExecutionMode, ServerArgs -from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger - -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, server_args: ServerArgs): - """ - Initialize the workflow with configuration arguments. - - Args: - server_args: Configuration object containing all parameters - needed for workflow and pipeline setup. - """ - self.server_args = server_args - - # TODO: pipeline_config should be: dict[str, PipelineConfig] - # pipeline_type should be included in the PipelineConfig - # pipeline_config[pipeline_name] = (pipeline_type, server_args) - self._pipeline_configs: dict[str, tuple[PipelineType, ServerArgs]] = {} - 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, server_args = pipeline_config - pipeline = build_pipeline(server_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, ServerArgs] - ) -> 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, server_args: ServerArgs) -> 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: - server_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 server_args.mode == ExecutionMode.PREPROCESS: - from sglang.multimodal_gen.runtime.workflow.preprocess.preprocess_workflow import ( - PreprocessWorkflow, - ) - - return PreprocessWorkflow.get_workflow_cls(server_args) - else: - raise ValueError( - f"Execution mode: {server_args.mode} is not supported in workflow." - )