flowyml 1.7.1__py3-none-any.whl → 1.8.0__py3-none-any.whl

flowyml/core/step.py

@@ -31,6 +31,8 @@ class StepConfig:
     tags: dict[str, str] = field(default_factory=dict)
     condition: Callable | None = None
     execution_group: str | None = None
+    source_file: str | None = None
+    source_line: int | None = None
 
     def __hash__(self):
         """Make StepConfig hashable."""
@@ -84,11 +86,15 @@ class Step:
         self.condition = condition
         self.execution_group = execution_group
 
-        # Capture source code for UI display
+        # Capture source code and location for UI display
         try:
             self.source_code = inspect.getsource(func)
+            self.source_file = inspect.getsourcefile(func)
+            _, self.source_line = inspect.getsourcelines(func)
         except (OSError, TypeError):
             self.source_code = "# Source code not available"
+            self.source_file = None
+            self.source_line = None
 
         self.config = StepConfig(
             name=self.name,
@@ -102,6 +108,8 @@ class Step:
             tags=self.tags,
             condition=self.condition,
             execution_group=self.execution_group,
+            source_file=self.source_file,
+            source_line=self.source_line,
         )
 
     def __call__(self, *args, **kwargs):

flowyml/core/step_grouping.py

@@ -172,26 +172,23 @@ class StepGroupAnalyzer:
         Returns:
             True if steps can execute consecutively
         """
-        # Get all dependencies of step2
-        step2_deps = dag.get_all_dependencies(step2.name)
+        # Get ALL transitively producing and consuming nodes between step1 and step2
+        # Steps are consecutive if there are no intermediate steps NOT in this group
+        # that must execute between step1 and step2.
+        all_deps_of_s2 = 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 step1 is not even a dependency of step2, they are independent.
+        # They can be grouped as long as there is no path from step1 to step2
+        # through an external step.
 
-        # 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}
+        # All nodes on any path from step1 to step2:
+        all_successors_of_s1 = dag.get_all_dependents(step1.name)
+        intermediate_nodes = all_successors_of_s1 & all_deps_of_s2
 
-            # If there are NO intermediate group steps, they're consecutive
-            return len(intermediate) == 0
+        # If any node on a path from s1 to s2 is NOT in the group, they are not consecutive
+        external_intermediates = intermediate_nodes - group_step_names
 
-        # step2 doesn't depend on step1, not consecutive
-        return False
+        return len(external_intermediates) == 0
 
     def _get_execution_order(self, steps: list[Step], dag: DAG) -> list[str]:
         """Get topological execution order for steps in a group.
@@ -264,29 +261,46 @@ def get_execution_units(dag: DAG, steps: list[Step]) -> list[Step | StepGroup]:
         for step in group.steps:
             step_to_group[step.name] = group
 
-    # Get topological order of all nodes
-    all_nodes = dag.topological_sort()
+    # To correctly determine execution order of units (which may have changed due to grouping),
+    # we build a new DAG where each node is an execution unit (Step or StepGroup).
+    from flowyml.core.graph import Node as DAGNode
 
-    # Build execution units, avoiding duplicates for grouped steps
-    execution_units: list[Step | StepGroup] = []
-    processed_groups: set[str] = set()
+    units_dag = DAG()
+    unit_map: dict[str, Step | StepGroup] = {}
 
-    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:
+    # Add units as nodes
+    processed_steps = set()
+    for step in steps:
+        if step.name in processed_steps:
             continue
 
-        # Check if this step belongs to a group
+        unit: Step | StepGroup
         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)
+            unit = step_to_group[step.name]
+            unit_name = f"group:{unit.group_name}"
+            # Extract names for inputs/outputs
+            u_inputs_set = set()
+            u_outputs_set = set()
+            for s in unit.steps:
+                u_inputs_set.update(s.inputs)
+                u_outputs_set.update(s.outputs)
+                processed_steps.add(s.name)
+
+            # External inputs are those not produced within the group
+            u_inputs = list(u_inputs_set - u_outputs_set)
+            u_outputs = list(u_outputs_set)
         else:
-            # Ungrouped step, add as-is
-            execution_units.append(step)
+            unit = step
+            unit_name = step.name
+            u_inputs = step.inputs
+            u_outputs = step.outputs
+            processed_steps.add(step.name)
+
+        unit_map[unit_name] = unit
+        units_dag.add_node(DAGNode(name=unit_name, step=unit, inputs=u_inputs, outputs=u_outputs))
+
+    # Build edges and sort
+    units_dag.build_edges()
+    sorted_unit_nodes = units_dag.topological_sort()
 
-    return execution_units
+    return [unit_map[node.name] for node in sorted_unit_nodes]

flowyml/core/types.py

@@ -0,0 +1,407 @@
+"""FlowyML Core Types - Artifact Types for Automatic Routing.
+
+This module defines artifact types that enable automatic routing to appropriate
+infrastructure based on type annotations. Just annotate your step outputs with
+these types, and FlowyML will route them to the correct stores and registries.
+
+Usage:
+    from flowyml.core import step, Model, Dataset, Metrics
+
+    @step
+    def train_model(data: Dataset) -> Model:
+        model = train(data)
+        return Model(model, name="my_classifier", version="1.0.0")
+
+    @step
+    def evaluate(model: Model) -> Metrics:
+        return Metrics({"accuracy": 0.95, "f1": 0.92})
+
+The stack configuration determines where each type is routed:
+    - Model → artifact_store + optional model_registry
+    - Dataset → artifact_store
+    - Metrics → experiment_tracker
+"""
+
+from abc import ABC
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Union
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+# =============================================================================
+# BASE ARTIFACT TYPE
+# =============================================================================
+
+
+@dataclass
+class Artifact(ABC):  # noqa: B024
+    """Base artifact type for automatic routing.
+
+    All routable artifacts inherit from this class. The runtime inspects
+    step return types and routes outputs based on their artifact type.
+
+    Attributes:
+        data: The actual artifact data (model, dataset, etc.)
+        name: Optional name for the artifact
+        metadata: Additional metadata to store with the artifact
+        uri: URI where the artifact is stored (set after saving)
+    """
+
+    data: Any = None
+    name: str | None = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+    uri: str | None = None
+
+    def __post_init__(self):
+        """Validate artifact after initialization."""
+        if self.metadata is None:
+            self.metadata = {}
+
+    def with_metadata(self, **kwargs) -> "Artifact":
+        """Add metadata to the artifact.
+
+        Args:
+            **kwargs: Key-value pairs to add to metadata.
+
+        Returns:
+            Self for method chaining.
+        """
+        self.metadata.update(kwargs)
+        return self
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize artifact metadata to dictionary.
+
+        Returns:
+            Dictionary representation (excluding the data itself).
+        """
+        return {
+            "type": self.__class__.__name__,
+            "name": self.name,
+            "metadata": self.metadata,
+            "uri": self.uri,
+        }
+
+
+# =============================================================================
+# MODEL ARTIFACT
+# =============================================================================
+
+
+@dataclass
+class Model(Artifact):
+    """ML model artifact - routes to artifact store + optional model registry.
+
+    Use this type for step outputs that are machine learning models.
+    When configured, models are automatically:
+    1. Saved to the artifact store (GCS, S3, local)
+    2. Registered in the model registry (Vertex AI, SageMaker, MLflow)
+
+    Attributes:
+        data: The model object (sklearn, pytorch, tensorflow, etc.)
+        name: Model name for registry
+        version: Optional version string
+        framework: ML framework (auto-detected if not provided)
+        serving_config: Optional serving configuration
+
+    Example:
+        @step
+        def train() -> Model:
+            clf = RandomForestClassifier().fit(X, y)
+            return Model(
+                data=clf,
+                name="fraud_detector",
+                version="1.0.0",
+                framework="sklearn"
+            )
+    """
+
+    version: str | None = None
+    framework: str | None = None
+    serving_config: dict[str, Any] | None = None
+    input_schema: dict[str, Any] | None = None
+    output_schema: dict[str, Any] | None = None
+
+    def __post_init__(self):
+        """Auto-detect framework if not provided."""
+        super().__post_init__()
+        if self.framework is None and self.data is not None:
+            self.framework = self._detect_framework()
+
+    def _detect_framework(self) -> str:
+        """Detect the ML framework from the model object."""
+        model = self.data
+        model_type = type(model).__module__
+
+        if "sklearn" in model_type:
+            return "sklearn"
+        elif "torch" in model_type:
+            return "pytorch"
+        elif "tensorflow" in model_type or "keras" in model_type:
+            return "tensorflow"
+        elif "xgboost" in model_type:
+            return "xgboost"
+        elif "lightgbm" in model_type:
+            return "lightgbm"
+        elif "catboost" in model_type:
+            return "catboost"
+        else:
+            return "unknown"
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize model metadata."""
+        base = super().to_dict()
+        base.update(
+            {
+                "version": self.version,
+                "framework": self.framework,
+                "serving_config": self.serving_config,
+                "input_schema": self.input_schema,
+                "output_schema": self.output_schema,
+            },
+        )
+        return base
+
+
+# =============================================================================
+# DATASET ARTIFACT
+# =============================================================================
+
+
+@dataclass
+class Dataset(Artifact):
+    """Dataset artifact - routes to artifact store.
+
+    Use this type for step outputs that are datasets (training data,
+    feature tables, processed data, etc.).
+
+    Attributes:
+        data: The dataset (DataFrame, numpy array, file path, etc.)
+        name: Dataset name
+        format: Data format (parquet, csv, json, etc.)
+        schema: Optional schema definition
+        statistics: Optional statistics about the dataset
+
+    Example:
+        @step
+        def preprocess(raw_data: pd.DataFrame) -> Dataset:
+            processed = clean_and_transform(raw_data)
+            return Dataset(
+                data=processed,
+                name="training_features",
+                format="parquet",
+                statistics={"rows": len(processed)}
+            )
+    """
+
+    format: str | None = None  # noqa: A003
+    schema: dict[str, Any] | None = None
+    statistics: dict[str, Any] | None = None
+    num_rows: int | None = None
+    num_columns: int | None = None
+
+    def __post_init__(self):
+        """Auto-detect format and compute statistics if possible."""
+        super().__post_init__()
+        if self.data is not None:
+            self._detect_properties()
+
+    def _detect_properties(self):
+        """Detect dataset properties from the data."""
+        data = self.data
+
+        # Detect format
+        if self.format is None:
+            if hasattr(data, "to_parquet"):
+                self.format = "parquet"
+            elif hasattr(data, "to_csv"):
+                self.format = "csv"
+            elif isinstance(data, (str, Path)):
+                path = Path(data)
+                self.format = path.suffix.lstrip(".")
+
+        # Detect dimensions
+        if hasattr(data, "shape"):
+            shape = data.shape
+            self.num_rows = shape[0] if len(shape) > 0 else None
+            self.num_columns = shape[1] if len(shape) > 1 else None
+        elif hasattr(data, "__len__"):
+            self.num_rows = len(data)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize dataset metadata."""
+        base = super().to_dict()
+        base.update(
+            {
+                "format": self.format,
+                "schema": self.schema,
+                "statistics": self.statistics,
+                "num_rows": self.num_rows,
+                "num_columns": self.num_columns,
+            },
+        )
+        return base
+
+
+# =============================================================================
+# METRICS ARTIFACT
+# =============================================================================
+
+
+class Metrics(dict):
+    """Metrics dictionary - routes to experiment tracker.
+
+    Use this type for step outputs that are evaluation metrics.
+    Metrics are automatically logged to the configured experiment
+    tracker (MLflow, Vertex AI Experiments, etc.).
+
+    This is a dict subclass for easy use - just return a Metrics dict
+    from your step and it will be automatically logged.
+
+    Example:
+        @step
+        def evaluate(model: Model, test_data: Dataset) -> Metrics:
+            predictions = model.predict(test_data)
+            return Metrics({
+                "accuracy": accuracy_score(y_true, predictions),
+                "f1": f1_score(y_true, predictions),
+                "precision": precision_score(y_true, predictions),
+            })
+    """
+
+    def __init__(self, data: dict[str, Union[int, float]] | None = None, **kwargs):
+        """Initialize metrics.
+
+        Args:
+            data: Dictionary of metric names to values.
+            **kwargs: Additional metrics as keyword arguments.
+        """
+        if data is None:
+            data = {}
+        super().__init__(data)
+        self.update(kwargs)
+        self._step: int | None = None
+        self._run_id: str | None = None
+        self._metadata: dict[str, Any] = {}
+
+    def at_step(self, step: int) -> "Metrics":
+        """Set the step number for these metrics.
+
+        Args:
+            step: Step/epoch number.
+
+        Returns:
+            Self for method chaining.
+        """
+        self._step = step
+        return self
+
+    def with_metadata(self, **kwargs) -> "Metrics":
+        """Add metadata to the metrics.
+
+        Args:
+            **kwargs: Key-value pairs to add to metadata.
+
+        Returns:
+            Self for method chaining.
+        """
+        self._metadata.update(kwargs)
+        return self
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize metrics to dictionary."""
+        return {
+            "type": "Metrics",
+            "values": dict(self),
+            "step": self._step,
+            "run_id": self._run_id,
+            "metadata": self._metadata,
+        }
+
+
+# =============================================================================
+# PARAMETERS ARTIFACT
+# =============================================================================
+
+
+class Parameters(dict):
+    """Parameters dictionary - logs to experiment tracker as params.
+
+    Use this type for step inputs/outputs that are hyperparameters
+    or configuration values. Parameters are logged for reproducibility.
+
+    Example:
+        @step
+        def train(params: Parameters) -> Model:
+            model = RandomForestClassifier(**params)
+            return Model(model.fit(X, y))
+    """
+
+    def __init__(self, data: dict[str, Any] | None = None, **kwargs):
+        """Initialize parameters.
+
+        Args:
+            data: Dictionary of parameter names to values.
+            **kwargs: Additional parameters as keyword arguments.
+        """
+        if data is None:
+            data = {}
+        super().__init__(data)
+        self.update(kwargs)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize parameters to dictionary."""
+        return {
+            "type": "Parameters",
+            "values": dict(self),
+        }
+
+
+# =============================================================================
+# HELPER FUNCTIONS
+# =============================================================================
+
+
+def is_artifact_type(obj: Any) -> bool:
+    """Check if an object is a FlowyML artifact type.
+
+    Args:
+        obj: Object or type to check.
+
+    Returns:
+        True if it's an artifact type or instance.
+    """
+    if isinstance(obj, type):
+        return issubclass(obj, (Artifact, Metrics, Parameters))
+    return isinstance(obj, (Artifact, Metrics, Parameters))
+
+
+def get_artifact_type_name(obj: Any) -> str | None:
+    """Get the artifact type name for routing.
+
+    Args:
+        obj: Artifact instance or type.
+
+    Returns:
+        Type name string or None.
+    """
+    if isinstance(obj, type):
+        if issubclass(obj, Artifact) or obj in (Metrics, Parameters):
+            return obj.__name__
+    elif isinstance(obj, (Artifact, Metrics, Parameters)):
+        return type(obj).__name__
+    return None
+
+
+# =============================================================================
+# TYPE ALIASES FOR CONVENIENCE
+# =============================================================================
+
+# Common type aliases
+ModelArtifact = Model
+DatasetArtifact = Dataset
+MetricsDict = Metrics
+ParamsDict = Parameters