flowyml 1.1.0__py3-none-any.whl

flowyml/core/step.py

@@ -0,0 +1,201 @@
+"""Step Decorator - Define pipeline steps with automatic context injection."""
+
+import hashlib
+import inspect
+import json
+from typing import Any, Union
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+# Import resource types
+try:
+    from flowyml.core.resources import ResourceRequirements
+except ImportError:
+    ResourceRequirements = None  # Type: ignore
+
+
+@dataclass
+class StepConfig:
+    """Configuration for a pipeline step."""
+
+    name: str
+    func: Callable
+    inputs: list[str] = field(default_factory=list)
+    outputs: list[str] = field(default_factory=list)
+    cache: bool | str | Callable = "code_hash"
+    retry: int = 0
+    timeout: int | None = None
+    resources: Union[dict[str, Any], "ResourceRequirements", None] = None
+    tags: dict[str, str] = field(default_factory=dict)
+    condition: Callable | None = None
+    execution_group: str | None = None
+
+    def __hash__(self):
+        """Make StepConfig hashable."""
+        return hash(self.name)
+
+
+class Step:
+    """A pipeline step that can be executed with automatic context injection."""
+
+    def __init__(
+        self,
+        func: Callable,
+        name: str | None = None,
+        inputs: list[str] | None = None,
+        outputs: list[str] | None = None,
+        cache: bool | str | Callable = "code_hash",
+        retry: int = 0,
+        timeout: int | None = None,
+        resources: Union[dict[str, Any], "ResourceRequirements", None] = None,
+        tags: dict[str, str] | None = None,
+        condition: Callable | None = None,
+        execution_group: str | None = None,
+    ):
+        self.func = func
+        self.name = name or func.__name__
+        self.inputs = inputs or []
+        self.outputs = outputs or []
+        self.cache = cache
+        self.retry = retry
+        self.timeout = timeout
+
+        # Store resources (accept both dict for backward compatibility and ResourceRequirements)
+        self.resources = resources
+
+        self.tags = tags or {}
+        self.condition = condition
+        self.execution_group = execution_group
+
+        # Capture source code for UI display
+        try:
+            self.source_code = inspect.getsource(func)
+        except (OSError, TypeError):
+            self.source_code = "# Source code not available"
+
+        self.config = StepConfig(
+            name=self.name,
+            func=func,
+            inputs=self.inputs,
+            outputs=self.outputs,
+            cache=self.cache,
+            retry=self.retry,
+            timeout=self.timeout,
+            resources=self.resources,
+            tags=self.tags,
+            condition=self.condition,
+            execution_group=self.execution_group,
+        )
+
+    def __call__(self, *args, **kwargs):
+        """Execute the step function."""
+        # Check condition if present
+        if self.condition:
+            # We might need to inject context into condition too,
+            # but for now assume it takes no args or same args as step?
+            # This is tricky without context injection logic here.
+            # The executor handles execution, so maybe we just store it here.
+            pass
+
+        return self.func(*args, **kwargs)
+
+    def get_code_hash(self) -> str:
+        """Compute hash of the step's source code."""
+        try:
+            source = inspect.getsource(self.func)
+            return hashlib.md5(source.encode()).hexdigest()
+        except (OSError, TypeError):
+            # Fallback for dynamically defined functions or when source is unavailable
+            return hashlib.md5(self.name.encode()).hexdigest()[:16]
+
+    def get_input_hash(self, inputs: dict[str, Any]) -> str:
+        """Generate hash of inputs for caching."""
+        input_str = json.dumps(inputs, sort_keys=True, default=str)
+        return hashlib.sha256(input_str.encode()).hexdigest()[:16]
+
+    def get_cache_key(self, inputs: dict[str, Any] | None = None) -> str:
+        """Generate cache key based on caching strategy.
+
+        Args:
+            inputs: Input data for the step
+
+        Returns:
+            Cache key string
+        """
+        if self.cache == "code_hash":
+            return f"{self.name}:{self.get_code_hash()}"
+        elif self.cache == "input_hash" and inputs:
+            return f"{self.name}:{self.get_input_hash(inputs)}"
+        elif callable(self.cache) and inputs:
+            return self.cache(inputs, {})
+        else:
+            return f"{self.name}:no-cache"
+
+    def __repr__(self) -> str:
+        return f"Step(name='{self.name}', inputs={self.inputs}, outputs={self.outputs})"
+
+
+def step(
+    _func: Callable | None = None,
+    *,
+    inputs: list[str] | None = None,
+    outputs: list[str] | None = None,
+    cache: bool | str | Callable = "code_hash",
+    retry: int = 0,
+    timeout: int | None = None,
+    resources: Union[dict[str, Any], "ResourceRequirements", None] = None,
+    tags: dict[str, str] | None = None,
+    name: str | None = None,
+    condition: Callable | None = None,
+    execution_group: str | None = None,
+):
+    """Decorator to define a pipeline step with automatic context injection.
+
+    Can be used as @step or @step(inputs=...)
+
+    Args:
+        _func: Function being decorated (when used as @step)
+        inputs: List of input asset names
+        outputs: List of output asset names
+        cache: Caching strategy ("code_hash", "input_hash", callable, or False)
+        retry: Number of retry attempts on failure
+        timeout: Maximum execution time in seconds
+        resources: Resource requirements (ResourceRequirements object or dict for backward compat)
+        tags: Metadata tags for the step
+        name: Optional custom name for the step
+        condition: Optional callable that returns True if step should run
+        execution_group: Optional group name for executing multiple steps together
+
+    Example:
+        >>> @step
+        ... def simple_step():
+        ...     ...
+        >>> @step(inputs=["data/train"], outputs=["model/trained"])
+        ... def train_model(train_data):
+        ...     ...
+        >>> # With resource requirements
+        >>> from flowyml.core.resources import ResourceRequirements, GPUConfig
+        >>> @step(resources=ResourceRequirements(cpu="4", memory="16Gi", gpu=GPUConfig(gpu_type="nvidia-v100", count=2)))
+        ... def gpu_train(data):
+        ...     ...
+    """
+
+    def decorator(func: Callable) -> Step:
+        return Step(
+            func=func,
+            name=name,
+            inputs=inputs,
+            outputs=outputs,
+            cache=cache,
+            retry=retry,
+            timeout=timeout,
+            resources=resources,
+            tags=tags,
+            condition=condition,
+            execution_group=execution_group,
+        )
+
+    if _func is None:
+        return decorator
+    else:
+        return decorator(_func)

flowyml/core/step_grouping.py

@@ -0,0 +1,292 @@
+"""Step Grouping - Analyze and group pipeline steps for efficient execution.
+
+This module provides functionality to group multiple pipeline steps that should execute
+together in the same environment (e.g., Docker container, remote worker). It analyzes
+the DAG to ensure only consecutive steps are grouped and aggregates their resource
+requirements intelligently.
+"""
+
+from collections import defaultdict
+from dataclasses import dataclass
+
+from flowyml.core.graph import DAG
+from flowyml.core.step import Step
+from flowyml.core.resources import ResourceRequirements
+
+
+@dataclass
+class StepGroup:
+    """Represents a group of steps that execute together.
+
+    Args:
+        group_name: Name identifier for this group
+        steps: List of Step objects in this group
+        aggregated_resources: Combined resource requirements for the group
+        execution_order: Ordered list of step names (topological order within group)
+    """
+
+    group_name: str
+    steps: list[Step]
+    aggregated_resources: ResourceRequirements | None
+    execution_order: list[str]
+
+    def __repr__(self) -> str:
+        step_names = [s.name for s in self.steps]
+        return f"StepGroup(name='{self.group_name}', steps={step_names})"
+
+
+class StepGroupAnalyzer:
+    """Analyzes pipeline DAG to create valid step groups.
+
+    This analyzer ensures that:
+    1. Only steps with the same execution_group name are grouped
+    2. Grouped steps can execute consecutively (no gaps in DAG)
+    3. Resources are aggregated intelligently (max CPU, memory, etc.)
+    4. Execution order within groups is preserved
+    """
+
+    def analyze_groups(self, dag: DAG, steps: list[Step]) -> list[StepGroup]:
+        """Analyze DAG and create valid step groups.
+
+        Args:
+            dag: Pipeline DAG
+            steps: List of all pipeline steps
+
+        Returns:
+            List of StepGroup objects (excludes ungrouped steps)
+        """
+        # Collect steps by execution_group
+        groups_dict: dict[str, list[Step]] = defaultdict(list)
+
+        for step in steps:
+            if step.execution_group:
+                groups_dict[step.execution_group].append(step)
+
+        # Process each group
+        step_groups = []
+        for group_name, group_steps in groups_dict.items():
+            # Split into consecutive subgroups if needed
+            subgroups = self._split_into_consecutive_groups(group_steps, dag)
+
+            # Create StepGroup for each subgroup
+            for i, subgroup in enumerate(subgroups):
+                # If original group was split, append index to name
+                final_name = group_name if len(subgroups) == 1 else f"{group_name}_{i}"
+
+                # Get execution order for this subgroup
+                exec_order = self._get_execution_order(subgroup, dag)
+
+                # Aggregate resources
+                aggregated = self._aggregate_resources(subgroup)
+
+                step_groups.append(
+                    StepGroup(
+                        group_name=final_name,
+                        steps=subgroup,
+                        aggregated_resources=aggregated,
+                        execution_order=exec_order,
+                    ),
+                )
+
+        return step_groups
+
+    def _split_into_consecutive_groups(
+        self,
+        steps: list[Step],
+        dag: DAG,
+    ) -> list[list[Step]]:
+        """Split steps into subgroups that can execute consecutively.
+
+        This handles cases where steps with the same execution_group are not
+        actually consecutive in the DAG (e.g., parallel branches).
+
+        Args:
+            steps: Steps with the same execution_group
+            dag: Pipeline DAG
+
+        Returns:
+            List of step sublists, where each sublist can execute consecutively
+        """
+        if len(steps) <= 1:
+            return [steps] if steps else []
+
+        # Build a mapping of step names to steps
+        step_map = {s.name: s for s in steps}
+        step_names = set(step_map.keys())
+
+        # Get topological order for all nodes
+        try:
+            all_nodes = dag.topological_sort()
+        except ValueError:
+            # DAG has cycles, return each step as separate group
+            return [[s] for s in steps]
+
+        # Filter to only our steps, preserving topological order
+        ordered_steps = [step_map[node.name] for node in all_nodes if node.name in step_names]
+
+        # Now split into consecutive sequences
+        # Two steps are consecutive if there are no other group steps between them
+        subgroups: list[list[Step]] = []
+        current_group: list[Step] = []
+
+        for step in ordered_steps:
+            if not current_group:
+                # Start new group
+                current_group.append(step)
+            else:
+                # Check if this step can join current group
+                last_step = current_group[-1]
+
+                if self._are_consecutive(last_step, step, dag, step_names):
+                    current_group.append(step)
+                else:
+                    # Gap detected, finalize current group and start new one
+                    subgroups.append(current_group)
+                    current_group = [step]
+
+        # Add final group
+        if current_group:
+            subgroups.append(current_group)
+
+        return subgroups
+
+    def _are_consecutive(
+        self,
+        step1: Step,
+        step2: Step,
+        dag: DAG,
+        group_step_names: set[str],
+    ) -> bool:
+        """Check if two steps can execute consecutively in a group.
+
+        Steps are consecutive if:
+        - step2 depends on step1 (directly or transitively) OR they're independent
+        - All intermediate dependencies are NOT in this group
+
+        Args:
+            step1: First step
+            step2: Second step
+            dag: Pipeline DAG
+            group_step_names: Set of all step names in this group
+
+        Returns:
+            True if steps can execute consecutively
+        """
+        # Get all dependencies of step2
+        step2_deps = dag.get_all_dependencies(step2.name)
+
+        # If step2 doesn't depend on anything in the group, they can be consecutive
+        # (parallel steps in same group are OK if no dependencies)
+        group_deps = step2_deps & group_step_names
+        if not group_deps:
+            # No dependencies from this group, consecutive is OK
+            return True
+
+        # If step2 depends on step1, check for intermediate group steps
+        if step1.name in step2_deps:
+            # Get all group steps that step2 depends on (excluding step1)
+            intermediate = group_deps - {step1.name}
+
+            # If there are NO intermediate group steps, they're consecutive
+            return len(intermediate) == 0
+
+        # step2 doesn't depend on step1, not consecutive
+        return False
+
+    def _get_execution_order(self, steps: list[Step], dag: DAG) -> list[str]:
+        """Get topological execution order for steps in a group.
+
+        Args:
+            steps: Steps in the group
+            dag: Pipeline DAG
+
+        Returns:
+            Ordered list of step names
+        """
+        step_names = {s.name for s in steps}
+
+        # Get full topological order
+        all_nodes = dag.topological_sort()
+
+        # Filter to only our steps
+        return [node.name for node in all_nodes if node.name in step_names]
+
+    def _aggregate_resources(self, steps: list[Step]) -> ResourceRequirements | None:
+        """Aggregate resource requirements from multiple steps.
+
+        Strategy:
+        - CPU: Take maximum
+        - Memory: Take maximum
+        - GPU: Merge configs (max count, best type)
+        - Storage: Take maximum
+        - Node affinity: Merge required/preferred labels
+
+        Args:
+            steps: Steps to aggregate resources from
+
+        Returns:
+            Aggregated ResourceRequirements or None if no steps have resources
+        """
+        resource_reqs = [s.resources for s in steps if s.resources and isinstance(s.resources, ResourceRequirements)]
+
+        if not resource_reqs:
+            return None
+
+        # Start with first resource requirement
+        aggregated = resource_reqs[0]
+
+        # Merge with remaining
+        for req in resource_reqs[1:]:
+            aggregated = aggregated.merge_with(req)
+
+        return aggregated
+
+
+def get_execution_units(dag: DAG, steps: list[Step]) -> list[Step | StepGroup]:
+    """Get ordered execution units (individual steps or groups).
+
+    This is a convenience function that analyzes groups and returns a mixed list
+    of ungrouped steps and StepGroups in topological order.
+
+    Args:
+        dag: Pipeline DAG
+        steps: All pipeline steps
+
+    Returns:
+        List of execution units (Step or StepGroup) in execution order
+    """
+    analyzer = StepGroupAnalyzer()
+    step_groups = analyzer.analyze_groups(dag, steps)
+
+    # Build a mapping of step names to their groups
+    step_to_group: dict[str, StepGroup] = {}
+    for group in step_groups:
+        for step in group.steps:
+            step_to_group[step.name] = group
+
+    # Get topological order of all nodes
+    all_nodes = dag.topological_sort()
+
+    # Build execution units, avoiding duplicates for grouped steps
+    execution_units: list[Step | StepGroup] = []
+    processed_groups: set[str] = set()
+
+    for node in all_nodes:
+        # Find the step object
+        step = next((s for s in steps if s.name == node.name), None)
+        if not step:
+            continue
+
+        # Check if this step belongs to a group
+        if step.name in step_to_group:
+            group = step_to_group[step.name]
+
+            # Only add the group once (when we encounter its first step)
+            if group.group_name not in processed_groups:
+                execution_units.append(group)
+                processed_groups.add(group.group_name)
+        else:
+            # Ungrouped step, add as-is
+            execution_units.append(step)
+
+    return execution_units