dirac-cwl 1.0.2__py3-none-any.whl

dirac_cwl/execution_hooks/core.py

@@ -0,0 +1,342 @@
+"""Core metadata framework for DIRAC CWL integration.
+
+This module provides the foundational classes and interfaces for the extensible
+metadata plugin system in DIRAC/DIRACX.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import (
+    Any,
+    ClassVar,
+    Dict,
+    List,
+    Mapping,
+    Optional,
+    Self,
+    Sequence,
+    TypeVar,
+    Union,
+)
+
+from DIRAC.DataManagementSystem.Client.DataManager import (  # type: ignore[import-untyped]
+    DataManager,
+)
+from DIRACCommon.Core.Utilities.ReturnValues import (  # type: ignore[import-untyped]
+    returnSingleResult,
+)
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
+
+from dirac_cwl.commands import PostProcessCommand, PreProcessCommand
+from dirac_cwl.data_management_mocks.data_manager import MockDataManager
+
+logger = logging.getLogger(__name__)
+
+# TypeVar for generic class methods
+T = TypeVar("T", bound="SchedulingHint")
+
+
+class ExecutionHooksBasePlugin(BaseModel):
+    """Base class for all runtime plugin models with execution hooks.
+
+    This class uses composition instead of inheritance for data catalog operations,
+    providing better separation of concerns and flexibility.
+    """
+
+    model_config = ConfigDict(
+        extra="ignore",
+        validate_assignment=True,
+        arbitrary_types_allowed=True,
+        json_schema_extra={
+            "title": "DIRAC Metadata Model",
+            "description": "Base metadata model for DIRAC jobs",
+        },
+    )
+
+    # Class-level metadata for plugin discovery
+    vo: ClassVar[Optional[str]] = None
+    version: ClassVar[str] = "1.0.0"
+    description: ClassVar[str] = "Base metadata model"
+
+    output_paths: Dict[str, Any] = {}
+    output_sandbox: list[str] = []
+    output_se: list[str] = []
+
+    _datamanager: DataManager = PrivateAttr(default_factory=DataManager)
+
+    def __init__(self, **kwargs):
+        """Initialize the execution hooks base plugin.
+
+        :param kwargs: Additional keyword arguments passed to the parent class.
+        """
+        super().__init__(**kwargs)
+        if os.getenv("DIRAC_PROTO_LOCAL") == "1":
+            self._datamanager = MockDataManager()
+
+    _preprocess_commands: List[type[PreProcessCommand]] = PrivateAttr(default=[])
+    _postprocess_commands: List[type[PostProcessCommand]] = PrivateAttr(default=[])
+
+    @property
+    def preprocess_commands(self) -> List[type[PreProcessCommand]]:
+        """Get the list of pre-processing commands."""
+        return self._preprocess_commands
+
+    @preprocess_commands.setter
+    def preprocess_commands(self, value: List[type[PreProcessCommand]]) -> None:
+        """Set the list of pre-processing commands."""
+        self._preprocess_commands = value
+
+    @property
+    def postprocess_commands(self) -> List[type[PostProcessCommand]]:
+        """Get the list of post-processing commands."""
+        return self._postprocess_commands
+
+    @postprocess_commands.setter
+    def postprocess_commands(self, value: List[type[PostProcessCommand]]) -> None:
+        """Set the list of post-processing commands."""
+        self._postprocess_commands = value
+
+    @classmethod
+    def name(cls) -> str:
+        """Auto-derive hook plugin identifier from class name."""
+        return cls.__name__
+
+    def store_output(
+        self,
+        outputs: dict[str, str | Path | Sequence[str | Path]],
+        **kwargs: Any,
+    ) -> None:
+        """Store an output file or set of files via the appropriate storage interface.
+
+        :param dict[str, str | Path | Sequence[str | Path]] outputs:
+            Dictionary containing the path or list of paths to the source file(s) to be stored
+            for each cwl output.
+        :param Any **kwargs:
+            Additional keyword arguments for extensibility.
+        """
+        for output_name, src_path in outputs.items():
+            logger.info("Storing output %s, with source %s", output_name, src_path)
+
+            if not src_path:
+                raise RuntimeError(f"src_path parameter required for filesystem storage of {output_name}")
+
+            lfn = self.output_paths.get(output_name, None)
+
+            if lfn:
+                if isinstance(src_path, str) or isinstance(src_path, Path):
+                    src_path = [src_path]
+                for src in src_path:
+                    file_lfn = Path(lfn) / Path(src).name
+                    res = None
+                    for se in self.output_se:
+                        res = returnSingleResult(self._datamanager.putAndRegister(str(file_lfn), src, se))
+                        if res["OK"]:
+                            logger.info("Successfully saved file %s with LFN %s", src, file_lfn)
+                            break
+                    if res and not res["OK"]:
+                        raise RuntimeError(f"Could not save file {src} with LFN {str(lfn)} : {res['Message']}")
+
+    def get_input_query(self, input_name: str, **kwargs: Any) -> Union[Path, List[Path], None]:
+        """Generate LFN-based input query path.
+
+        Accepts and ignores extra kwargs for interface compatibility.
+        """
+        # Build LFN: /query_root/vo/campaign/site/data_type/input_name
+        pass
+
+    @classmethod
+    def get_schema_info(cls) -> Dict[str, Any]:
+        """Get schema information for this metadata model."""
+        return {
+            "hook_plugin": cls.name(),
+            "vo": cls.vo,
+            "version": cls.version,
+            "description": cls.description,
+            "schema": cls.model_json_schema(),
+        }
+
+
+class Hint(ABC):
+    """Base class for all DIRAC hints and requirements models."""
+
+    @classmethod
+    @abstractmethod
+    def from_cwl(cls, cwl_object: Any) -> "Hint":
+        """Extract hint information from a CWL object."""
+        pass
+
+
+class SchedulingHint(BaseModel, Hint):
+    """Descriptor for job execution configuration."""
+
+    model_config = ConfigDict(extra="forbid", validate_assignment=True)
+
+    platform: Optional[str] = Field(default=None, description="Target platform (e.g., 'DIRAC', 'DIRACX')")
+
+    priority: int = Field(default=10, description="Job priority (higher values = higher priority)")
+
+    sites: Optional[List[str]] = Field(default=None, description="Candidate execution sites")
+
+    @classmethod
+    def from_cwl(cls: type[T], cwl_object: Any) -> T:
+        """Extract task descriptor from CWL hints."""
+        descriptor = cls()
+
+        hints = getattr(cwl_object, "hints", []) or []
+        for hint in hints:
+            if hint.get("class") == "dirac:Scheduling":
+                hint_data = {k: v for k, v in hint.items() if k != "class"}
+                descriptor = descriptor.model_copy(update=hint_data)
+
+        return descriptor
+
+
+class ExecutionHooksHint(BaseModel, Hint):
+    """Descriptor for data management configuration in CWL hints.
+
+    This class represents the serializable data management configuration that
+    can be embedded in CWL hints and later instantiated into concrete
+    metadata models.
+
+    Enhanced with submission functionality for DIRAC CWL integration.
+    """
+
+    model_config = ConfigDict(
+        extra="allow",  # Allow vo-specific fields
+        validate_assignment=True,
+        json_schema_extra={
+            "title": "DIRAC Data Manager",
+            "description": "Data management configuration for DIRAC jobs",
+        },
+    )
+
+    hook_plugin: str = Field(
+        default="QueryBasedPlugin",
+        description="Registry key for the metadata implementation class",
+    )
+
+    # Enhanced fields for submission functionality
+    configuration: Dict[str, Any] = Field(
+        default_factory=dict, description="Additional parameters for metadata plugins"
+    )
+
+    output_paths: Dict[str, Any] = Field(default_factory=dict, description="LFNs for outputs on the Data Catalog")
+
+    output_sandbox: list[str] = Field(
+        default_factory=list,
+        description="List of the outputs stored in the output sandbox",
+    )
+
+    output_se: list[str] = Field(
+        default_factory=lambda: ["SE-USER"],
+        description="List of Storage Elements that can be used to store the outputs",
+    )
+
+    def model_copy(
+        self,
+        update: Optional[Mapping[str, Any]] = None,
+        *,
+        deep: bool = False,
+    ) -> Self:
+        """Enhanced model copy with intelligent merging of dict fields (including configuration)."""
+        if update is None:
+            update = {}
+        else:
+            update = dict(update)
+
+        merged_update = {}
+        for key, value in update.items():
+            if hasattr(self, key) and isinstance(getattr(self, key), dict) and isinstance(value, dict):
+                existing_value = getattr(self, key).copy()
+                existing_value.update(value)
+                merged_update[key] = existing_value
+            else:
+                merged_update[key] = value
+
+        return super().model_copy(update=merged_update, deep=deep)
+
+    def to_runtime(self, submitted: Optional[Any] = None) -> "ExecutionHooksBasePlugin":
+        """Build and instantiate the runtime metadata implementation.
+
+        The returned object is an instance of :class:`ExecutionHooksBasePlugin` created
+        by the metadata registry. The instantiation parameters are constructed
+        by merging, in order:
+
+        1. Input defaults declared on the CWL task (if ``submitted`` is provided).
+        2. The first set of CWL parameter overrides (``submitted.parameters``),
+           if present.
+        3. The descriptor's ``configuration``.
+
+        During merging, keys are normalized from dash-case to snake_case to
+        align with typical Python argument names used by runtime implementations.
+
+        :param submitted: Optional submission context used to resolve CWL input defaults
+            and parameter overrides.
+        :type submitted: JobSubmissionModel | None
+        :return: Runtime plugin implementation instantiated from the registry.
+        :rtype: ExecutionHooksBasePlugin
+        """
+        # Import here to avoid circular imports
+        from .registry import get_registry
+
+        # Quick helper to convert dash-case to snake_case without importing utils
+        def _dash_to_snake(s: str) -> str:
+            return s.replace("-", "_")
+
+        if submitted is None:
+            descriptor = ExecutionHooksHint(
+                hook_plugin=self.hook_plugin,
+                output_paths=self.output_paths,
+                output_sandbox=self.output_sandbox,
+                output_se=self.output_se,
+                **self.configuration,
+            )
+            return get_registry().instantiate_plugin(descriptor)
+
+        # Build inputs from task defaults and parameter overrides
+        inputs: dict[str, Any] = {}
+        for inp in submitted.task.inputs:
+            input_name = inp.id.split("#")[-1].split("/")[-1]
+            input_value = getattr(inp, "default", None)
+            params_list = getattr(submitted, "parameters", None)
+            if params_list and params_list[0]:
+                input_value = params_list[0].cwl.get(input_name, input_value)
+            inputs[input_name] = input_value
+
+        # Merge with explicit configuration
+        if self.configuration:
+            inputs.update(self.configuration)
+
+        params = {_dash_to_snake(key): value for key, value in inputs.items()}
+
+        descriptor = ExecutionHooksHint(
+            hook_plugin=self.hook_plugin,
+            output_paths=self.output_paths,
+            output_sandbox=self.output_sandbox,
+            output_se=self.output_se,
+            **params,
+        )
+        return get_registry().instantiate_plugin(descriptor)
+
+    @classmethod
+    def from_cwl(cls, cwl_object: Any) -> Self:
+        """Extract metadata descriptor from CWL object using Hint interface."""
+        descriptor = cls()
+        hints = getattr(cwl_object, "hints", []) or []
+        for hint in hints:
+            if hint.get("class") == "dirac:ExecutionHooks":
+                hint_data = {k: v for k, v in hint.items() if k != "class"}
+                descriptor = descriptor.model_copy(update=hint_data)
+        return descriptor
+
+
+class TransformationExecutionHooksHint(ExecutionHooksHint):
+    """Extended data manager for transformations."""
+
+    group_size: Optional[Dict[str, int]] = Field(
+        default=None, description="Input grouping configuration for transformation jobs"
+    )

dirac_cwl/execution_hooks/plugins/__init__.py

@@ -0,0 +1,16 @@
+"""Plugin package initialization.
+
+This module ensures that core plugins are automatically registered
+when the metadata system is imported.
+"""
+
+from .core import (
+    QueryBasedPlugin,
+)
+
+# Plugins will be auto-registered through the metaclass or discovery system
+
+__all__ = [
+    # Core plugins
+    "QueryBasedPlugin",
+]

dirac_cwl/execution_hooks/plugins/core.py

@@ -0,0 +1,58 @@
+"""Core DIRAC metadata models.
+
+This module contains the standard metadata models provided by DIRAC core.
+These serve as examples and provide basic functionality for common use cases.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Any, ClassVar, List, Optional, Union
+
+from pydantic import Field
+
+from ..core import (
+    ExecutionHooksBasePlugin,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class QueryBasedPlugin(ExecutionHooksBasePlugin):
+    """Metadata plugin using LFN-based data catalog for structured data discovery.
+
+    This plugin demonstrates filesystem-based data organization using
+    Logical File Names (LFNs) with campaign, site, and data type parameters.
+    """
+
+    description: ClassVar[str] = "LFN-based metadata for structured data discovery"
+
+    # LFN parameters
+    query_root: str = Field(default="/grid/data", description="Base path for LFN structure")
+    site: Optional[str] = Field(default=None, description="Site identifier for LFN path")
+    campaign: Optional[str] = Field(default=None, description="Campaign name for LFN path")
+    data_type: Optional[str] = Field(default=None, description="Data type classification")
+
+    def get_input_query(self, input_name: str, **kwargs: Any) -> Union[Path, List[Path], None]:
+        """Generate LFN-based input query path.
+
+        Accepts and ignores extra kwargs for interface compatibility.
+        """
+        # Build LFN: /query_root/vo/campaign/site/data_type/input_name
+        path_parts = []
+
+        if self.vo:
+            path_parts.append(self.vo)
+
+        if self.campaign:
+            path_parts.append(self.campaign)
+        if self.site:
+            path_parts.append(self.site)
+        if self.data_type:
+            path_parts.append(self.data_type)
+
+        if len(path_parts) > 0:  # More than just VO
+            return Path(self.query_root) / Path(*path_parts) / Path(input_name)
+
+        return Path(self.query_root) / Path(input_name)

dirac_cwl/execution_hooks/registry.py

@@ -0,0 +1,209 @@
+"""Enhanced plugin registry for metadata models.
+
+This module provides a sophisticated plugin discovery and registration system
+for DIRAC metadata models, supporting virtual organization-specific extensions and
+automatic discovery.
+"""
+
+from __future__ import annotations
+
+import logging
+from importlib.metadata import entry_points
+from typing import Any, Dict, List, Optional, Type
+
+from .core import ExecutionHooksBasePlugin, ExecutionHooksHint
+
+logger = logging.getLogger(__name__)
+
+
+class ExecutionHooksPluginRegistry:
+    """
+    Registry for execution hooks plugins.
+
+    This class manages the registration and retrieval of execution hooks plugins
+    for different steps in CWL workflows. Plugins are registered using
+    entry points and can be retrieved by name.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the execution hooks registry."""
+        self._plugins: Dict[str, Type[ExecutionHooksBasePlugin]] = {}
+        self._vo_plugins: Dict[str, Dict[str, Type[ExecutionHooksBasePlugin]]] = {}
+        self._plugin_info: Dict[str, Dict[str, Any]] = {}
+
+    def register_plugin(self, plugin_class: Type[ExecutionHooksBasePlugin], override: bool = False) -> None:
+        """Register a metadata plugin.
+
+        :param plugin_class: The metadata model class to register.
+        :type plugin_class: Type[ExecutionHooksBasePlugin]
+        :param override: Whether to override existing registrations, by default False.
+        :type override: bool
+        :raises ValueError: If plugin is already registered and override=False.
+        """
+        if not issubclass(plugin_class, ExecutionHooksBasePlugin):
+            raise ValueError(f"Plugin {plugin_class} must inherit from ExecutionHooksBasePlugin")
+
+        plugin_key = plugin_class.name()
+        vo = plugin_class.vo
+
+        # Check for conflicts
+        if plugin_key in self._plugins and not override:
+            existing = self._plugins[plugin_key]
+            raise ValueError(
+                f"Plugin '{plugin_key}' already registered by {existing.__module__}.{existing.__name__}. "
+                f"Use override=True to replace."
+            )
+
+        # Register globally
+        self._plugins[plugin_key] = plugin_class
+        self._plugin_info[plugin_key] = plugin_class.get_schema_info()
+
+        # Register by VO if specified
+        if vo:
+            if vo not in self._vo_plugins:
+                self._vo_plugins[vo] = {}
+            self._vo_plugins[vo][plugin_key] = plugin_class
+
+        vo_suffix = f" (VO: {vo})" if vo else ""
+        logger.info(
+            "Registered metadata plugin '%s' from %s.%s%s",
+            plugin_key,
+            plugin_class.__module__,
+            plugin_class.__name__,
+            vo_suffix,
+        )
+
+    def get_plugin(self, plugin_key: str, vo: Optional[str] = None) -> Optional[Type[ExecutionHooksBasePlugin]]:
+        """Get a registered plugin.
+
+        :param plugin_key: The plugin identifier.
+        :type plugin_key: str
+        :param vo: Virtual Organization namespace to search first, by default None.
+        :type vo: Optional[str]
+        :return: The plugin class or None if not found.
+        :rtype: Optional[Type[ExecutionHooksBasePlugin]]
+        """
+        # Try VO-specific first if specified
+        if vo and vo in self._vo_plugins:
+            if plugin_key in self._vo_plugins[vo]:
+                return self._vo_plugins[vo][plugin_key]
+
+        # Fall back to global registry
+        return self._plugins.get(plugin_key)
+
+    def instantiate_plugin(self, descriptor: ExecutionHooksHint, **kwargs: Any) -> ExecutionHooksBasePlugin:
+        """Instantiate a metadata plugin from a descriptor.
+
+        :param descriptor: The data manager containing configuration.
+        :type descriptor: ExecutionHooksHint
+        :param kwargs: Additional parameters to pass to the plugin constructor.
+        :type kwargs: Any
+        :return: Instantiated metadata model.
+        :rtype: ExecutionHooksBasePlugin
+        :raises KeyError: If the requested plugin is not registered.
+        :raises ValueError: If plugin instantiation fails.
+        """
+        plugin_class = self.get_plugin(descriptor.hook_plugin)
+
+        if plugin_class is None:
+            available = self.list_plugins()
+            raise KeyError(f"Unknown execution hooks plugin: '{descriptor.hook_plugin}'" f"Available: {available}")
+
+        # Extract plugin parameters from descriptor
+        plugin_params = descriptor.model_dump(
+            exclude={
+                "hook_plugin",
+            }
+        )
+        plugin_params.update(kwargs)
+
+        try:
+            return plugin_class(**plugin_params)
+        except Exception as e:
+            raise ValueError(f"Failed to instantiate plugin '{descriptor.hook_plugin}': {e}") from e
+
+    def list_plugins(self, vo: Optional[str] = None) -> List[str]:
+        """List available plugins.
+
+        :param vo: Filter by Virtual Organization, by default None.
+        :type vo: Optional[str]
+        :return: List of available plugin keys.
+        :rtype: List[str]
+        """
+        if vo and vo in self._vo_plugins:
+            return list(self._vo_plugins[vo].keys())
+        return list(self._plugins.keys())
+
+    def list_virtual_organizations(self) -> List[str]:
+        """List available Virtual Organizations."""
+        return list(self._vo_plugins.keys())
+
+    def get_plugin_info(self, plugin_key: str) -> Optional[Dict[str, Any]]:
+        """Get detailed information about a plugin."""
+        return self._plugin_info.get(plugin_key)
+
+    def discover_plugins(self) -> int:
+        """Discover and register plugins from the entry points defined in the pyproject.toml.
+
+        :return: Number of plugins discovered and registered.
+        :rtype: int
+        """
+        entrypoints = entry_points(group="dirac_cwl.execution_hooks")
+        discovered = 0
+        for hook_name in entrypoints.names:
+            try:
+                hook = entrypoints[hook_name].load()
+                if issubclass(hook, ExecutionHooksBasePlugin):
+                    self.register_plugin(hook)
+                    discovered += 1
+                else:
+                    logger.warning(
+                        "Tried to discover execution hook with name '%s' that does not inherit %s",
+                        hook_name,
+                        ExecutionHooksBasePlugin.__name__,
+                    )
+            except Exception as e:
+                logger.error("Failed to import plugin %s: %s", hook_name, e)
+
+        return discovered
+
+    def validate_descriptor(self, descriptor: ExecutionHooksHint) -> List[str]:
+        """Validate a data manager against registered plugins.
+
+        :param descriptor: The data manager to validate.
+        :type descriptor: ExecutionHooksHint
+        :return: List of validation errors (empty if valid).
+        :rtype: List[str]
+        """
+        errors = []
+
+        plugin_class = self.get_plugin(descriptor.hook_plugin)
+
+        if plugin_class is None:
+            available = self.list_plugins()
+            errors.append(f"Unknown metadata plugin: '{descriptor.hook_plugin}'. " f"Available: {available}")
+            return errors
+
+        # Validate descriptor against plugin schema
+        try:
+            plugin_params = descriptor.model_dump(exclude={"hook_plugin"})
+            plugin_class.model_validate(plugin_params)
+        except Exception as e:
+            errors.append(f"Plugin validation failed: {e}")
+
+        return errors
+
+
+# Global registry instance
+_registry = ExecutionHooksPluginRegistry()
+
+
+# Public API
+def get_registry() -> ExecutionHooksPluginRegistry:
+    """Get the global execution hooks plugin registry."""
+    return _registry
+
+
+def discover_plugins() -> int:
+    """Discover and register plugins from packages."""
+    return _registry.discover_plugins()