deriva-ml 1.17.10__py3-none-any.whl → 1.17.12__py3-none-any.whl

deriva_ml/execution/runner.py

@@ -0,0 +1,595 @@
+"""
+Deriva ML Model Runner
+======================
+
+Generic model runner for executing ML workflows within DerivaML execution contexts.
+
+This module provides the infrastructure to run ML models with full provenance tracking,
+configuration management via Hydra, and support for parameter sweeps.
+
+Key Features
+------------
+- **Automatic execution context**: Creates execution records in the catalog
+- **Multirun/sweep support**: Parent-child execution nesting for parameter sweeps
+- **Hydra configuration**: Composable configs with command-line overrides
+- **Subclass support**: Works with DerivaML subclasses (EyeAI, GUDMAP, etc.)
+- **Provenance tracking**: Links inputs, outputs, and configuration
+
+Model Protocol
+--------------
+Models must follow this signature pattern to work with run_model:
+
+    def my_model(
+        param1: int = 10,
+        param2: float = 0.01,
+        # ... other model parameters ...
+        ml_instance: DerivaML = None,  # Injected at runtime
+        execution: Execution = None,   # Injected at runtime
+    ) -> None:
+        '''Train/run the model within the execution context.'''
+        # Access input datasets
+        for dataset in execution.datasets:
+            bag = execution.download_dataset_bag(dataset)
+            # ... process data ...
+
+        # Register output files
+        model_path = execution.asset_file_path("Model", "model.pt")
+        torch.save(model.state_dict(), model_path)
+
+        metrics_path = execution.asset_file_path("Execution_Metadata", "metrics.json")
+        with open(metrics_path, "w") as f:
+            json.dump({"accuracy": 0.95}, f)
+
+The `ml_instance` and `execution` parameters are injected by run_model at runtime.
+All other parameters are configured via Hydra.
+
+Quick Start
+-----------
+1. Create your model function following the protocol above.
+
+2. Create a hydra-zen configuration for your model:
+
+    from hydra_zen import builds, store
+
+    # Wrap model with builds() and zen_partial=True for deferred execution
+    MyModelConfig = builds(my_model, param1=10, param2=0.01, zen_partial=True)
+    store(MyModelConfig, group="model_config", name="default_model")
+
+3. Set up the main runner script:
+
+    from deriva_ml import DerivaML
+    from deriva_ml.execution import run_model, create_model_config
+    from hydra_zen import store, zen
+
+    # Create the main config (uses DerivaML by default)
+    deriva_model = create_model_config(DerivaML)
+    store(deriva_model, name="deriva_model")
+
+    # Load your config modules
+    store.add_to_hydra_store()
+
+    # Launch
+    if __name__ == "__main__":
+        zen(run_model).hydra_main(config_name="deriva_model", version_base="1.3")
+
+4. Run from command line:
+
+    python my_runner.py                           # Run with defaults
+    python my_runner.py model_config.param1=20    # Override parameter
+    python my_runner.py dry_run=true              # Test without catalog writes
+    python my_runner.py --multirun model_config.param1=10,20,30  # Parameter sweep
+
+Domain Subclasses
+-----------------
+For domain-specific classes like EyeAI:
+
+    from eye_ai import EyeAI
+
+    # Create config with EyeAI instead of DerivaML
+    deriva_model = create_model_config(EyeAI, description="EyeAI analysis")
+
+    # Your model receives an EyeAI instance:
+    def my_eyeai_model(
+        ...,
+        ml_instance: EyeAI = None,  # Now an EyeAI instance
+        execution: Execution = None,
+    ):
+        # Access EyeAI-specific methods
+        ml_instance.some_eyeai_method()
+
+Multirun Parameter Sweeps
+-------------------------
+When using Hydra's multirun mode (--multirun or -m), run_model automatically:
+
+1. Creates a parent execution to group all sweep jobs
+2. Links each child execution to the parent with sequence ordering
+3. Records sweep configuration in the parent's description
+
+Example sweep:
+
+    python my_runner.py --multirun model_config.learning_rate=0.001,0.01,0.1
+
+This creates:
+- Parent execution: "Multirun sweep: ..." (contains sweep metadata)
+- Child 0: learning_rate=0.001 (sequence=0)
+- Child 1: learning_rate=0.01 (sequence=1)
+- Child 2: learning_rate=0.1 (sequence=2)
+
+Query nested executions via the catalog or MCP tools:
+- list_nested_executions(parent_rid)
+- list_parent_executions(child_rid)
+
+Configuration Groups
+--------------------
+The default hydra_defaults in create_model_config() expect these config groups:
+
+- deriva_ml: Connection settings (hostname, catalog_id, credentials)
+- datasets: Dataset specifications (RIDs, versions)
+- assets: Asset RIDs (model weights, etc.)
+- workflow: Workflow definition (name, type, description)
+- model_config: Model parameters (your model's config)
+
+Each group should have at least a "default_*" entry. Override at runtime:
+
+    python my_runner.py deriva_ml=production datasets=full_training
+
+See Also
+--------
+- DerivaMLModel protocol: defines the expected model signature
+- ExecutionConfiguration: bundles inputs for an execution
+- Execution: context manager for execution lifecycle
+"""
+
+from __future__ import annotations
+
+import atexit
+import logging
+from pathlib import Path
+from typing import Any, TypeVar, TYPE_CHECKING
+
+from hydra.core.hydra_config import HydraConfig
+from hydra_zen import builds
+
+if TYPE_CHECKING:
+    from deriva_ml import DerivaML
+    from deriva_ml.core.config import DerivaMLConfig
+    from deriva_ml.dataset import DatasetSpec
+    from deriva_ml.execution import ExecutionConfiguration, Workflow
+    from deriva_ml.core.definitions import RID
+
+
+# Type variable for DerivaML and its subclasses
+T = TypeVar("T", bound="DerivaML")
+
+
+# =============================================================================
+# Multirun State Management
+# =============================================================================
+
+class MultirunState:
+    """Manages state for multirun (sweep) executions.
+
+    In multirun mode, we create a parent execution that groups all the
+    individual sweep jobs. This class holds the shared state needed to
+    coordinate between jobs.
+
+    Attributes:
+        parent_execution_rid: RID of the parent execution (created on first job)
+        parent_execution: The parent Execution object
+        ml_instance: Shared DerivaML instance
+        job_sequence: Counter for ordering child executions
+        sweep_dir: Path to the sweep output directory
+    """
+    parent_execution_rid: str | None = None
+    parent_execution: Any = None
+    ml_instance: Any = None  # DerivaML or subclass
+    job_sequence: int = 0
+    sweep_dir: Path | None = None
+
+
+# Global instance - persists across jobs in a multirun
+_multirun_state = MultirunState()
+
+
+def _is_multirun() -> bool:
+    """Check if we're running in Hydra multirun mode."""
+    try:
+        hydra_cfg = HydraConfig.get()
+        # RunMode.MULTIRUN has value 2
+        return hydra_cfg.mode.value == 2
+    except Exception:
+        return False
+
+
+def _get_job_num() -> int:
+    """Get the current job number in a multirun."""
+    try:
+        hydra_cfg = HydraConfig.get()
+        return hydra_cfg.job.num
+    except Exception:
+        return 0
+
+
+def _complete_parent_execution() -> None:
+    """Complete the parent execution at the end of a multirun sweep.
+
+    This is registered as an atexit handler to ensure the parent execution
+    is properly completed and its outputs uploaded when the process exits.
+    """
+    global _multirun_state
+
+    if _multirun_state.parent_execution is None:
+        return
+
+    try:
+        parent = _multirun_state.parent_execution
+
+        # Stop the parent execution timing
+        parent.execution_stop()
+
+        # Upload any outputs and clean up
+        parent.upload_execution_outputs()
+
+        logging.info(
+            f"Completed parent execution: {_multirun_state.parent_execution_rid} "
+            f"({_multirun_state.job_sequence} child jobs)"
+        )
+    except Exception as e:
+        logging.warning(f"Failed to complete parent execution: {e}")
+    finally:
+        # Clear the state
+        reset_multirun_state()
+
+
+# Track if atexit handler is registered
+_atexit_registered = False
+
+
+def _create_parent_execution(
+    ml_instance: "DerivaML",
+    workflow: "Workflow",
+    description: str,
+    dry_run: bool = False,
+) -> None:
+    """Create the parent execution for a multirun sweep.
+
+    This is called on the first job of a multirun to create the parent
+    execution that will group all child executions together.
+
+    Args:
+        ml_instance: The DerivaML (or subclass) instance.
+        workflow: The workflow to associate with the parent execution.
+        description: Description for the parent execution. When using multirun_config,
+            this is the rich markdown description from the config.
+        dry_run: If True, don't write to the catalog.
+    """
+    global _multirun_state, _atexit_registered
+
+    # Import here to avoid circular imports
+    from deriva_ml.execution import ExecutionConfiguration
+
+    # Use the description directly - it comes from multirun_config or the CLI
+    parent_description = description
+
+    # Create parent execution configuration (no datasets - those are for children)
+    parent_config = ExecutionConfiguration(
+        description=parent_description,
+    )
+
+    # Create the parent execution
+    parent_execution = ml_instance.create_execution(
+        parent_config,
+        workflow=workflow,
+        dry_run=dry_run,
+    )
+
+    # Start the parent execution
+    parent_execution.execution_start()
+
+    # Store in global state
+    _multirun_state.parent_execution = parent_execution
+    _multirun_state.parent_execution_rid = parent_execution.execution_rid
+    _multirun_state.ml_instance = ml_instance
+
+    # Register atexit handler to complete parent execution when process exits
+    if not _atexit_registered:
+        atexit.register(_complete_parent_execution)
+        _atexit_registered = True
+
+    logging.info(f"Created parent execution: {parent_execution.execution_rid}")
+
+
+def run_model(
+    deriva_ml: "DerivaMLConfig",
+    datasets: list["DatasetSpec"],
+    assets: list["RID"],
+    description: str,
+    workflow: "Workflow",
+    model_config: Any,
+    dry_run: bool = False,
+    ml_class: type["DerivaML"] | None = None,
+) -> None:
+    """
+    Execute a machine learning model within a DerivaML execution context.
+
+    This function serves as the main entry point called by hydra-zen after
+    configuration resolution. It orchestrates the complete execution lifecycle:
+    connecting to Deriva, creating an execution record, running the model,
+    and uploading results.
+
+    In multirun mode, this function also:
+    - Creates a parent execution on the first job to group all sweep jobs
+    - Links each child execution to the parent with sequence ordering
+
+    Parameters
+    ----------
+    deriva_ml : DerivaMLConfig
+        Configuration for the DerivaML connection. Contains server URL,
+        catalog ID, credentials, and other connection parameters.
+
+    datasets : list[DatasetSpec]
+        Specifications for datasets to use in this execution. Each DatasetSpec
+        identifies a dataset in the Deriva catalog to be made available to
+        the model.
+
+    assets : list[RID]
+        Resource IDs (RIDs) of assets to include in the execution. Typically
+        used for model weight files, pretrained checkpoints, or other
+        artifacts needed by the model.
+
+    description : str
+        Human-readable description of this execution run. Stored in the
+        Deriva catalog for provenance tracking. In multirun mode, this is
+        also used for the parent execution if running via multirun_config.
+
+    workflow : Workflow
+        The workflow definition to associate with this execution. Defines
+        the computational pipeline and its metadata.
+
+    model_config : Any
+        A hydra-zen callable that wraps the actual model code. When called
+        with `ml_instance` and `execution` arguments, it runs the model
+        training or inference logic.
+
+    dry_run : bool, optional
+        If True, create the execution record but skip actual model execution.
+        Useful for testing configuration without running expensive computations.
+        Default is False.
+
+    ml_class : type[DerivaML], optional
+        The DerivaML class (or subclass) to instantiate. If None, uses the
+        base DerivaML class. Use this to instantiate domain-specific classes
+        like EyeAI or GUDMAP.
+
+    Returns
+    -------
+    None
+        Results are uploaded to the Deriva catalog as execution outputs.
+
+    Examples
+    --------
+    This function is typically not called directly, but through hydra:
+
+        # From command line:
+        python deriva_run.py +experiment=cifar10_cnn dry_run=True
+
+        # Multirun (creates parent + child executions):
+        python deriva_run.py --multirun +experiment=cifar10_quick,cifar10_extended
+
+        # With a custom DerivaML subclass (in your script):
+        from functools import partial
+        run_model_eyeai = partial(run_model, ml_class=EyeAI)
+    """
+    global _multirun_state
+
+    # Import here to avoid circular imports
+    from deriva_ml import DerivaML
+    from deriva_ml.execution import ExecutionConfiguration
+
+    # ---------------------------------------------------------------------------
+    # Clear hydra's logging configuration
+    # ---------------------------------------------------------------------------
+    # Hydra sets up its own logging handlers which can interfere with DerivaML's
+    # logging. Remove them to ensure consistent log output.
+    root = logging.getLogger()
+    for handler in root.handlers[:]:
+        root.removeHandler(handler)
+
+    # ---------------------------------------------------------------------------
+    # Connect to the Deriva catalog
+    # ---------------------------------------------------------------------------
+    # Use the provided ml_class or default to DerivaML
+    if ml_class is None:
+        ml_class = DerivaML
+
+    ml_instance = ml_class.instantiate(deriva_ml)
+
+    # ---------------------------------------------------------------------------
+    # Handle multirun mode - create parent execution on first job
+    # ---------------------------------------------------------------------------
+    is_multirun = _is_multirun()
+    if is_multirun and _multirun_state.parent_execution is None:
+        _create_parent_execution(ml_instance, workflow, description, dry_run)
+
+    # ---------------------------------------------------------------------------
+    # Capture Hydra runtime choices for provenance
+    # ---------------------------------------------------------------------------
+    # The choices dict maps config group names to the selected config names
+    # e.g., {"model_config": "cifar10_quick", "datasets": "cifar10_training"}
+    # Filter out None values (some Hydra internal groups have None choices)
+    config_choices: dict[str, str] = {}
+    try:
+        hydra_cfg = HydraConfig.get()
+        config_choices = {k: v for k, v in hydra_cfg.runtime.choices.items() if v is not None}
+    except Exception:
+        pass  # HydraConfig not available outside Hydra context
+
+    # ---------------------------------------------------------------------------
+    # Create the execution context
+    # ---------------------------------------------------------------------------
+    # The ExecutionConfiguration bundles together all the inputs for this run:
+    # which datasets to use, which assets (model weights, etc.), and metadata.
+
+    # In multirun mode, enhance the description with job info
+    job_description = description
+    if is_multirun:
+        job_num = _get_job_num()
+        job_description = f"[Job {job_num}] {description}"
+
+    execution_config = ExecutionConfiguration(
+        datasets=datasets,
+        assets=assets,
+        description=job_description,
+        config_choices=config_choices,
+    )
+
+    # Create the execution record in the catalog. This generates a unique
+    # execution ID and sets up the working directories for this run.
+    execution = ml_instance.create_execution(
+        execution_config,
+        workflow=workflow,
+        dry_run=dry_run
+    )
+
+    # ---------------------------------------------------------------------------
+    # Link to parent execution in multirun mode
+    # ---------------------------------------------------------------------------
+    if is_multirun and _multirun_state.parent_execution is not None:
+        if not dry_run:
+            try:
+                # Get the current job sequence from the global state
+                job_sequence = _multirun_state.job_sequence
+                _multirun_state.parent_execution.add_nested_execution(
+                    execution,
+                    sequence=job_sequence
+                )
+                logging.info(
+                    f"Linked execution {execution.execution_rid} to parent "
+                    f"{_multirun_state.parent_execution_rid} (sequence={job_sequence})"
+                )
+                # Increment the sequence for the next job
+                _multirun_state.job_sequence += 1
+            except Exception as e:
+                logging.warning(f"Failed to link execution to parent: {e}")
+
+    # ---------------------------------------------------------------------------
+    # Run the model within the execution context
+    # ---------------------------------------------------------------------------
+    # The context manager handles setup (downloading datasets, creating output
+    # directories) and teardown (recording completion status, timing).
+    with execution.execute() as exec_context:
+        if dry_run:
+            # In dry run mode, skip model execution but still test the setup
+            logging.info("Dry run mode: skipping model execution")
+        else:
+            # Invoke the model configuration callable. The model_config is a
+            # hydra-zen wrapped function that has been partially configured with
+            # all model-specific parameters (e.g., learning rate, batch size).
+            # We provide the runtime context here.
+            model_config(ml_instance=ml_instance, execution=exec_context)
+
+    # ---------------------------------------------------------------------------
+    # Upload results to the catalog
+    # ---------------------------------------------------------------------------
+    # After the model completes, upload any output files (metrics, predictions,
+    # model checkpoints) to the Deriva catalog for permanent storage.
+    if not dry_run:
+        uploaded_assets = execution.upload_execution_outputs()
+
+        # Print summary of uploaded assets
+        total_files = sum(len(files) for files in uploaded_assets.values())
+        if total_files > 0:
+            print(f"\nUploaded {total_files} asset(s) to catalog:")
+            for asset_type, files in uploaded_assets.items():
+                for f in files:
+                    print(f"  - {asset_type}: {f}")
+
+
+def create_model_config(
+    ml_class: type["DerivaML"] | None = None,
+    description: str = "Model execution",
+    hydra_defaults: list | None = None,
+) -> Any:
+    """Create a hydra-zen configuration for run_model.
+
+    This helper creates a properly configured hydra-zen builds() for run_model
+    with the specified DerivaML class bound via partial application.
+
+    Parameters
+    ----------
+    ml_class : type[DerivaML], optional
+        The DerivaML class (or subclass) to use. If None, uses the base DerivaML.
+
+    description : str, optional
+        Default description for executions. Can be overridden at runtime.
+
+    hydra_defaults : list, optional
+        Custom hydra defaults. If None, uses standard defaults for deriva_ml,
+        datasets, assets, workflow, and model_config groups.
+
+    Returns
+    -------
+    Any
+        A hydra-zen builds() configuration ready to be registered with store().
+
+    Examples
+    --------
+    Basic usage with DerivaML:
+
+        >>> from deriva_ml.execution.runner import create_model_config
+        >>> model_config = create_model_config()
+        >>> store(model_config, name="deriva_model")
+
+    With a custom subclass:
+
+        >>> from eye_ai import EyeAI
+        >>> model_config = create_model_config(EyeAI, description="EyeAI analysis")
+        >>> store(model_config, name="eyeai_model")
+
+    With custom hydra defaults:
+
+        >>> model_config = create_model_config(
+        ...     hydra_defaults=[
+        ...         "_self_",
+        ...         {"deriva_ml": "production"},
+        ...         {"datasets": "full_dataset"},
+        ...     ]
+        ... )
+    """
+    from functools import partial
+
+    if hydra_defaults is None:
+        hydra_defaults = [
+            "_self_",
+            {"deriva_ml": "default_deriva"},
+            {"datasets": "default_dataset"},
+            {"assets": "default_asset"},
+            {"workflow": "default_workflow"},
+            {"model_config": "default_model"},
+        ]
+
+    # Create a partial function with ml_class bound
+    if ml_class is not None:
+        run_func = partial(run_model, ml_class=ml_class)
+    else:
+        run_func = run_model
+
+    return builds(
+        run_func,
+        description=description,
+        populate_full_signature=True,
+        hydra_defaults=hydra_defaults,
+    )
+
+
+def reset_multirun_state() -> None:
+    """Reset the global multirun state.
+
+    This is primarily useful for testing to ensure clean state between tests.
+    """
+    global _multirun_state
+    _multirun_state.parent_execution_rid = None
+    _multirun_state.parent_execution = None
+    _multirun_state.ml_instance = None
+    _multirun_state.job_sequence = 0
+    _multirun_state.sweep_dir = None