axl-workflows 0.1.0__py3-none-any.whl

axl/__init__.py

@@ -0,0 +1,47 @@
+"""
+AXL Workflows - Lightweight framework for building data and ML workflows.
+
+Write once → run anywhere (Dagster locally or Kubeflow in production).
+"""
+
+__version__ = "0.1.0"
+__author__ = "Pedro Spinosa"
+
+# Core imports
+from .core import Workflow, step, workflow
+
+# IO imports
+from .io import (
+    CloudpickleIOHandler,
+    InMemoryStorage,
+    IOHandler,
+    LocalFileStorage,
+    PickleIOHandler,
+    StorageBackend,
+    create_storage_from_path,
+    registry,
+    storage_registry,
+)
+
+# Logging imports
+from .logging import WorkflowLogger
+
+# Runtime imports
+from .runtime import LocalRuntime
+
+__all__ = [
+    "workflow",
+    "step",
+    "Workflow",
+    "IOHandler",
+    "PickleIOHandler",
+    "CloudpickleIOHandler",
+    "registry",
+    "StorageBackend",
+    "LocalFileStorage",
+    "InMemoryStorage",
+    "storage_registry",
+    "create_storage_from_path",
+    "LocalRuntime",
+    "WorkflowLogger",
+]

axl/cli.py

@@ -0,0 +1,322 @@
+"""
+CLI entry point for AXL Workflows.
+
+This module provides the command-line interface for the axl tool.
+"""
+
+import importlib
+import sys
+from pathlib import Path
+from typing import Any
+
+import typer
+import yaml
+from rich.console import Console
+from rich.panel import Panel
+from rich.table import Table
+
+from . import __version__
+from .ir import build_ir, get_workflow_info
+from .runtime import LocalRuntime
+
+app = typer.Typer(
+    name="axl",
+    help="AXL Workflows - Lightweight framework for building data and ML workflows",
+    add_completion=False,
+)
+console = Console()
+
+
+@app.callback()
+def main(
+    version: bool = typer.Option(
+        None,
+        "--version",
+        "-v",
+        help="Show version and exit",
+        callback=lambda v: typer.echo(f"axl {__version__}") if v else None,
+        is_eager=True,
+    ),
+) -> None:
+    """
+    AXL Workflows - Write once → run anywhere (Dagster locally or Kubeflow in production).
+    """
+    pass
+
+
+@app.command()
+def version() -> None:
+    """Show version information."""
+    console.print(Panel(f"AXL Workflows v{__version__}", title="Version"))
+
+
+def load_parameters(params_file: str | None) -> dict[str, Any]:
+    """
+    Load parameters from YAML file.
+
+    Args:
+        params_file: Path to YAML parameters file
+
+    Returns:
+        Dictionary of parameters
+
+    Raises:
+        typer.Exit: If file not found or invalid YAML
+    """
+    if not params_file:
+        return {}
+
+    try:
+        params_path = Path(params_file)
+        if not params_path.exists():
+            raise FileNotFoundError(f"Parameters file not found: {params_path}")
+
+        with open(params_path) as f:
+            params = yaml.safe_load(f)
+
+        if not isinstance(params, dict):
+            raise ValueError("Parameters file must contain a YAML dictionary")
+
+        console.print(
+            f"[green]Loaded {len(params)} parameters from {params_file}[/green]"
+        )
+        return params
+
+    except Exception as e:
+        console.print(f"[red]Error loading parameters: {e}[/red]")
+        raise typer.Exit(1) from None
+
+
+def load_workflow_class(module_path: str) -> type:
+    """
+    Load workflow class from module path.
+
+    Args:
+        module_path: Module path in format 'module:Class' or 'path/to/file.py:Class'
+
+    Returns:
+        Workflow class
+
+    Raises:
+        typer.Exit: If module or class not found
+    """
+    try:
+        if ":" not in module_path:
+            raise ValueError(
+                "Module path must be in format 'module:Class' or 'path/to/file.py:Class'"
+            )
+
+        module_spec, class_name = module_path.split(":", 1)
+
+        # Handle file paths
+        if module_spec.endswith(".py"):
+            file_path = Path(module_spec)
+            if not file_path.exists():
+                raise FileNotFoundError(f"File not found: {file_path}")
+
+            # Add parent directory to sys.path
+            sys.path.insert(0, str(file_path.parent))
+            module_spec = file_path.stem
+
+        # Import module
+        module = importlib.import_module(module_spec)
+
+        # Get class
+        if not hasattr(module, class_name):
+            raise AttributeError(
+                f"Class '{class_name}' not found in module '{module_spec}'"
+            )
+
+        workflow_class = getattr(module, class_name)
+
+        # Validate it's a workflow class
+        if not hasattr(workflow_class, "_is_workflow"):
+            raise ValueError(
+                f"Class '{class_name}' is not a workflow class (missing @workflow decorator)"
+            )
+
+        return workflow_class  # type: ignore[no-any-return]
+
+    except Exception as e:
+        console.print(f"[red]Error loading workflow: {e}[/red]")
+        raise typer.Exit(1) from None
+
+
+@app.command()
+def validate(
+    module: str = typer.Argument(
+        ...,
+        help="Module path to workflow class (e.g., 'myflow:MyWorkflow')",
+    ),
+) -> None:
+    """Validate a workflow definition."""
+    try:
+        console.print(f"[yellow]Validating workflow: {module}[/yellow]")
+
+        # Load workflow class
+        workflow_class = load_workflow_class(module)
+
+        # Get basic info
+        info = get_workflow_info(workflow_class)
+
+        # Display workflow info
+        table = Table(title="Workflow Information")
+        table.add_column("Property", style="cyan")
+        table.add_column("Value", style="green")
+
+        table.add_row("Name", info["name"])
+        table.add_row("Image", info["image"])
+        table.add_row("IO Handler", info["io_handler"])
+        table.add_row("Steps", str(info["step_count"]))
+
+        console.print(table)
+
+        # Validate workflow
+        console.print("\n[yellow]Building IR and validating...[/yellow]")
+        ir_workflow = build_ir(workflow_class)
+
+        # Display validation results
+        console.print("[green]✅ Workflow is valid![/green]")
+
+        # Show step details
+        step_table = Table(title="Steps")
+        step_table.add_column("Step", style="cyan")
+        step_table.add_column("Inputs", style="blue")
+        step_table.add_column("IO Handler", style="green")
+
+        for node in ir_workflow.nodes:
+            inputs = ", ".join(node.inputs) if node.inputs else "None"
+            io_handler = node.get_io_handler() or "default"
+            step_table.add_row(node.name, inputs, io_handler)
+
+        console.print(step_table)
+
+    except Exception as e:
+        console.print(f"[red]❌ Validation failed: {e}[/red]")
+        raise typer.Exit(1) from None
+
+
+@app.command()
+def compile(
+    module: str = typer.Argument(
+        ...,
+        help="Module path to workflow class (e.g., 'myflow:MyWorkflow')",
+    ),
+    target: str = typer.Option(
+        "argo",
+        "--target",
+        "-t",
+        help="Target backend (argo, dagster)",
+    ),
+    out: str = typer.Option(
+        None,
+        "--out",
+        "-o",
+        help="Output file path",
+    ),
+) -> None:
+    """Compile a workflow to target backend."""
+    # TODO: Implement workflow compilation
+    console.print(f"[yellow]Compiling {module} to {target}[/yellow]")
+    if out:
+        console.print(f"[yellow]Output: {out}[/yellow]")
+    console.print("[red]Not implemented yet[/red]")
+
+
+@app.command()
+def run(
+    backend: str = typer.Argument(
+        "local",
+        help="Backend to run on (local, argo, dagster)",
+    ),
+    module: str = typer.Option(
+        None,
+        "--module",
+        "-m",
+        help="Module path to workflow class (e.g., 'examples/iris_example.py:IrisKNN')",
+    ),
+    storage_backend: str = typer.Option(
+        "memory",
+        "--storage",
+        help="Storage backend for artifacts (memory|local)",
+    ),
+    workspace: str = typer.Option(
+        "axl_workspace",
+        "--workspace",
+        help="Workspace directory for local storage backend",
+    ),
+    params: str = typer.Option(
+        None,
+        "--params",
+        "-p",
+        help="Parameters file (YAML) to pass to workflow constructor",
+    ),
+) -> None:
+    """Run a workflow on the specified backend.
+
+    Currently implemented: local backend using LocalRuntime.
+    """
+    try:
+        console.print(f"[yellow]Running workflow on {backend}[/yellow]")
+
+        if backend != "local":
+            console.print("[red]Only 'local' backend is implemented for now[/red]")
+            raise typer.Exit(1)
+
+        if not module:
+            console.print("[red]--module is required for local execution[/red]")
+            raise typer.Exit(1)
+
+        # Load parameters if provided
+        workflow_params = load_parameters(params)
+
+        # Load workflow class
+        workflow_class = load_workflow_class(module)
+
+        # Build IR
+        ir_workflow = build_ir(workflow_class)
+
+        # Instantiate workflow with parameters
+        wf_instance = workflow_class(**workflow_params)
+
+        # Execute locally
+        console.print(
+            f"[yellow]Executing locally (storage={storage_backend}, workspace='{workspace}')[/yellow]"
+        )
+        runtime = LocalRuntime(
+            workspace_path=workspace, storage_backend=storage_backend
+        )
+        runtime.execute_workflow(ir_workflow, wf_instance)
+
+        console.print(Panel("Local run completed", title="Run"))
+
+    except typer.Exit:
+        raise
+    except Exception as e:
+        console.print(f"[red]Run failed: {e}[/red]")
+        raise typer.Exit(1) from None
+
+
+@app.command()
+def render(
+    module: str = typer.Argument(
+        ...,
+        help="Module path to workflow class",
+    ),
+    out: str = typer.Option(
+        None,
+        "--out",
+        "-o",
+        help="Output file path (e.g., dag.png)",
+    ),
+) -> None:
+    """Render workflow DAG as a graph."""
+    # TODO: Implement DAG rendering
+    console.print(f"[yellow]Rendering DAG for {module}[/yellow]")
+    if out:
+        console.print(f"[yellow]Output: {out}[/yellow]")
+    console.print("[red]Not implemented yet[/red]")
+
+
+if __name__ == "__main__":
+    app()

axl/compiler/__init__.py

@@ -0,0 +1,11 @@
+"""
+Compilers for AXL Workflows.
+
+This module contains backend-specific compilers that convert IR to target formats.
+"""
+
+# TODO: Implement compiler modules
+# from .argo import ArgoCompiler
+# from .dagster import DagsterCompiler
+
+__all__: list[str] = []

axl/core/__init__.py

@@ -0,0 +1,10 @@
+"""
+Core DSL components for AXL Workflows.
+
+This module contains the main decorators and base classes for defining workflows.
+"""
+
+from .decorators import step, workflow
+from .workflow import Workflow
+
+__all__ = ["workflow", "step", "Workflow"]

axl/core/decorators.py

@@ -0,0 +1,158 @@
+"""
+Decorators for AXL Workflows.
+
+This module contains the @workflow and @step decorators for defining workflows.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from functools import wraps
+from typing import Any, Protocol, TypedDict, TypeVar, cast, runtime_checkable
+
+from ..ir import validate_step_args
+from .workflow import Workflow
+
+T = TypeVar("T", bound=type[Workflow])
+
+
+class StepConfig(TypedDict, total=False):
+    """Typed configuration attached to step functions at decoration time."""
+
+    io_handler: str
+    input_mode: str | dict[str, str]
+    resources: dict[str, Any]
+    retries: int
+    env: dict[str, str]
+
+
+@runtime_checkable
+class StepFunction(Protocol):
+    """Protocol for functions decorated with @step."""
+
+    _is_step: bool
+    _step_config: StepConfig
+    _original_func: Callable[..., Any]
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
+
+
+def workflow(
+    *,
+    name: str | None = None,
+    image: str | None = None,
+    io_handler: str | None = None,
+) -> Callable[[T], T]:
+    """
+    Decorator for defining workflows with configuration overrides.
+
+    Args:
+        name: Workflow name (optional, defaults to class name)
+        image: Docker image override
+        io_handler: IO handler override
+
+    Returns:
+        Decorated workflow class
+
+    Example:
+        @workflow(name="my-workflow", image="custom:latest")
+        class MyWorkflow(Workflow):
+            def graph(self):
+                return self.step1()
+    """
+
+    def decorator(cls: T) -> T:
+        # Validate that the class inherits from Workflow
+        if not issubclass(cls, Workflow):
+            raise TypeError(f"Class {cls.__name__} must inherit from Workflow")
+
+        # Store original __init__ method
+        original_init = cls.__init__
+
+        def new_init(self: Workflow, *args: Any, **kwargs: Any) -> None:
+            # Apply decorator overrides to kwargs
+            if image is not None:
+                kwargs["image"] = image
+            if io_handler is not None:
+                kwargs["io_handler"] = io_handler
+
+            # Call original __init__ with updated kwargs
+            original_init(self, *args, **kwargs)
+
+            # Set workflow name (prefer decorator override)
+            self.name = name or cls.__name__
+
+        # Replace the __init__ method
+        cls.__init__ = new_init  # type: ignore[method-assign]
+
+        # Add workflow metadata on the class (dynamic attrs by design)
+        cls._is_workflow = True  # type: ignore[attr-defined]
+        cls._workflow_name = name or cls.__name__  # type: ignore[attr-defined]
+
+        return cls
+
+    return decorator
+
+
+def step(
+    *,
+    io_handler: str | None = None,
+    input_mode: str | dict[str, str] | None = None,
+    resources: dict[str, Any] | None = None,
+    retries: int | None = None,
+    env: dict[str, str] | None = None,
+) -> Callable[[Callable[..., Any]], StepFunction]:
+    """
+    Decorator for defining workflow steps with configuration.
+
+    Args:
+        io_handler: IO handler override for this step
+        input_mode: Input mode override ("object", "path", or {arg: mode})
+        resources: Resource requirements (CPU, memory, etc.)
+        retries: Number of retries for this step
+        env: Environment variables for this step
+
+    Returns:
+        Decorated step function
+
+    Example:
+        @step(io_handler="parquet", retries=3)
+        def preprocess(self, data):
+            return processed_data
+    """
+
+    def decorator(func: Callable[..., Any]) -> StepFunction:
+        # Build step configuration (only include provided keys)
+        step_cfg: StepConfig = {}
+        if io_handler is not None:
+            step_cfg["io_handler"] = io_handler
+        if input_mode is not None:
+            step_cfg["input_mode"] = input_mode
+        if resources is not None:
+            step_cfg["resources"] = resources
+        if retries is not None:
+            step_cfg["retries"] = retries
+        if env is not None:
+            step_cfg["env"] = env
+
+        # Attach attributes to the original function (ruff prefers direct assigns)
+        cast(Any, func)._is_step = True
+        cast(Any, func)._step_config = step_cfg
+
+        # Create a wrapper that validates arguments
+        @wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            # Validate that all arguments are OutputRef instances
+            validate_step_args(func.__name__, args, kwargs)
+            # Call the original function
+            return func(*args, **kwargs)
+
+        # Copy attributes to wrapper so runtime can read them
+        cast(Any, wrapper)._is_step = True
+        cast(Any, wrapper)._step_config = step_cfg
+        cast(Any, wrapper)._original_func = func
+
+        # Explicit cast so mypy understands wrapper carries the attributes
+        return cast(StepFunction, wrapper)
+
+    return decorator

axl/core/workflow.py

@@ -0,0 +1,120 @@
+"""
+Workflow base class for AXL Workflows.
+
+This module contains the core Workflow base class with configuration
+and validation capabilities.
+"""
+
+from typing import Any
+
+from ..ir import OutputRef
+from ..logging import WorkflowLogger
+
+
+class Workflow:
+    """Base class for workflows with defaults and configuration."""
+
+    def __init__(
+        self,
+        image: str = "ghcr.io/axl-workflows/runner:latest",
+        io_handler: str = "pickle",
+        **params: Any,
+    ) -> None:
+        """
+        Initialize workflow with configuration and parameters.
+
+        Args:
+            image: Docker image for workflow execution
+            io_handler: IO handler for step outputs (pickle, parquet, numpy, torch)
+            **params: Additional parameters passed to the workflow
+        """
+        self.image = image
+        self.io_handler = io_handler
+        self.name: str = ""  # Will be set by @workflow decorator
+
+        # Store parameters
+        self.workflow_params = params
+
+        # Initialize logger for this workflow instance
+        self.log = WorkflowLogger(self.name)
+
+        # Validate configuration
+        self._validate_config()
+
+    def _validate_config(self) -> None:
+        """Validate workflow configuration."""
+        if self.io_handler not in ["pickle", "parquet", "numpy", "torch"]:
+            raise ValueError(
+                f"io_handler must be one of ['pickle', 'parquet', 'numpy', 'torch'], "
+                f"got {self.io_handler}"
+            )
+
+    def configure(self, **kwargs: Any) -> None:
+        """
+        Configure workflow settings after initialization.
+
+        Args:
+            **kwargs: Configuration options to update
+        """
+        for key, value in kwargs.items():
+            if hasattr(self, key):
+                setattr(self, key, value)
+            else:
+                raise ValueError(f"Unknown configuration option: {key}")
+
+        # Re-validate after configuration changes
+        self._validate_config()
+
+    def get_workflow_config(self) -> dict[str, Any]:
+        """
+        Get workflow configuration for IR building.
+
+        Returns:
+            Dictionary with workflow configuration
+        """
+        return {
+            "image": self.image,
+            "io_handler": self.io_handler,
+        }
+
+    def graph(self) -> Any:
+        """
+        Define the workflow graph. Must be implemented by subclasses.
+
+        Returns:
+            Workflow output (typically the final step result as OutputRef)
+        """
+        raise NotImplementedError("graph method must be implemented")
+
+    def __getattribute__(self, name: str) -> Any:
+        """
+        Intercept step method calls to return OutputRef instances.
+
+        Step method calls always return OutputRef instances for IR building.
+
+        Args:
+            name: Method name to look up
+
+        Returns:
+            OutputRef if it's a step method, otherwise normal attribute
+        """
+        # Get the attribute normally first
+        attr = object.__getattribute__(self, name)
+
+        # Check if this is a step method
+        if hasattr(attr, "_is_step") and attr._is_step:
+            # Create OutputRef with step configuration
+            output_ref = OutputRef(name)
+            if hasattr(attr, "_step_config"):
+                output_ref.metadata["step_config"] = attr._step_config
+            return output_ref
+
+        return attr
+
+    def __repr__(self) -> str:
+        """String representation of the workflow."""
+        return (
+            f"{self.__class__.__name__}("
+            f"image='{self.image}', "
+            f"io_handler='{self.io_handler}')"
+        )

axl/io/__init__.py

@@ -0,0 +1,27 @@
+"""IO handlers for AXL Workflows."""
+
+from .handlers import CloudpickleIOHandler, PickleIOHandler
+from .protocol import IOHandler
+from .registry import IOHandlerRegistry, registry
+from .storage import (
+    InMemoryStorage,
+    LocalFileStorage,
+    StorageBackend,
+    StorageRegistry,
+    create_storage_from_path,
+    storage_registry,
+)
+
+__all__ = [
+    "IOHandler",
+    "PickleIOHandler",
+    "CloudpickleIOHandler",
+    "IOHandlerRegistry",
+    "registry",
+    "StorageBackend",
+    "LocalFileStorage",
+    "InMemoryStorage",
+    "StorageRegistry",
+    "storage_registry",
+    "create_storage_from_path",
+]

axl/io/handlers/__init__.py

@@ -0,0 +1,9 @@
+"""IO handler implementations for AXL Workflows."""
+
+from .cloudpickle import CloudpickleIOHandler
+from .pickle import PickleIOHandler
+
+__all__ = [
+    "PickleIOHandler",
+    "CloudpickleIOHandler",
+]