dirac-cwl 1.0.2__py3-none-any.whl

dirac_cwl/__init__.py

@@ -0,0 +1,28 @@
+"""DIRAC CWL Proto - Common Workflow Language integration for DIRAC."""
+
+import logging
+from importlib.metadata import PackageNotFoundError, version
+
+import typer
+
+from dirac_cwl.job import app as job_app
+from dirac_cwl.production import app as production_app
+from dirac_cwl.transformation import app as transformation_app
+
+logging.basicConfig(level=logging.DEBUG, format="%(asctime)s | %(name)s | %(levelname)s | %(message)s")
+
+try:
+    __version__ = version("dirac-cwl")
+except PackageNotFoundError:
+    # package is not installed
+    pass
+
+app = typer.Typer()
+
+# Add sub-apps
+app.add_typer(production_app, name="production")
+app.add_typer(transformation_app, name="transformation")
+app.add_typer(job_app, name="job")
+
+if __name__ == "__main__":
+    app()

dirac_cwl/commands/__init__.py

@@ -0,0 +1,5 @@
+"""Command classes for workflow pre/post-processing operations."""
+
+from .core import PostProcessCommand, PreProcessCommand
+
+__all__ = ["PreProcessCommand", "PostProcessCommand"]

dirac_cwl/commands/core.py

@@ -0,0 +1,37 @@
+"""Core base classes for workflow processing commands."""
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+
+class CommandBase(ABC):
+    """Base abstract class for pre/post-processing commands.
+
+    New commands **MUST NOT** inherit this class. Instead they should inherit the interface classes
+    :class:`dirac_cwl.commands.base.PreProcessCommand` and
+    :class:`dirac_cwl.commands.base.PostProcessCommand`
+    """
+
+    @abstractmethod
+    def execute(self, job_path: Path, **kwargs) -> None:
+        """Execute the command in the given job path.
+
+        :param job_path: Path to the job working directory.
+        :param kwargs: Additional keyword arguments.
+        :raises NotImplementedError: This method must be implemented by subclasses.
+        """
+        raise NotImplementedError("This method should be implemented by child class")
+
+
+class PreProcessCommand(CommandBase):
+    """Interface class for pre-processing commands.
+
+    Every pre-processing command must inherit this class. Used for type validation.
+    """
+
+
+class PostProcessCommand(CommandBase):
+    """Interface class for post-processing commands.
+
+    Every post-processing command must inherit this class. Used for type validation.
+    """

dirac_cwl/commands/download_config.py

@@ -0,0 +1,22 @@
+"""Example pre-processing command that downloads configuration."""
+
+import os
+
+from dirac_cwl.commands import PreProcessCommand
+
+
+class DownloadConfig(PreProcessCommand):
+    """Example command that creates a file with named 'content.cfg'."""
+
+    def execute(self, job_path, **kwargs):
+        """Execute the configuration download.
+
+        :param job_path: Path to the job working directory.
+        :param kwargs: Additional keyword arguments.
+        """
+        content = """\
+This is an example
+"""
+        file_path = os.path.join(job_path, "content.cfg")
+        with open(file_path, "w") as f:
+            f.write(content)

dirac_cwl/commands/group_outputs.py

@@ -0,0 +1,32 @@
+"""Example post-processing command that groups output files."""
+
+import glob
+import os
+
+from dirac_cwl.commands import PostProcessCommand
+
+
+class GroupOutputs(PostProcessCommand):
+    """Example command that merges all of the outputs in a singular file."""
+
+    def execute(self, job_path, **kwargs):
+        """Execute the output file grouping.
+
+        :param job_path: Path to the job working directory.
+        :param kwargs: Additional keyword arguments.
+        """
+        grouped_outputs = "group.out"
+        output_path = os.path.join(job_path, grouped_outputs)
+        output_files = ["*.out", "*.txt"]
+
+        with open(output_path, "w", encoding="utf-8") as f_in:
+            for file_type in output_files:
+                extension = f"{job_path}/{file_type}"
+                for file in glob.glob(extension):
+                    if file == output_path:
+                        continue
+
+                    with open(file, "r", encoding="utf-8") as f_out:
+                        f_in.write(f"############ {file}\n")
+                        f_in.writelines(f_out.readlines())
+                        f_in.write("\n")

dirac_cwl/core/__init__.py

@@ -0,0 +1 @@
+"""Dirac-cwl core package."""

dirac_cwl/core/exceptions.py

@@ -0,0 +1,5 @@
+"""Dirac-cwl exceptions module."""
+
+
+class WorkflowProcessingException(Exception):
+    """Workflow exception raised during pre and post processing steps."""

dirac_cwl/core/utility.py

@@ -0,0 +1,41 @@
+"""Utility functions for file catalog operations."""
+
+from pathlib import Path
+from typing import Any
+
+from cwl_utils.parser.cwl_v1_2 import (
+    File,
+)
+
+
+def get_lfns(input_data: dict[str, Any]) -> dict[str, list[Path]]:
+    """Get the list of LFNs in the inputs from the parameters.
+
+    :param input_data: The parameters of the job.
+    :return: The list of LFN paths.
+    """
+    # Get the files from the input data
+    files: dict[str, list[Path]] = {}
+    for input_name, input_value in input_data.items():
+        val = []
+        if isinstance(input_value, list):
+            for item in input_value:
+                if isinstance(item, File):
+                    if not item.location and not item.path:
+                        raise NotImplementedError("File location is not defined.")
+
+                    if not item.location:
+                        continue
+                    # Skip files from the File Catalog
+                    if item.location.startswith("lfn:"):
+                        val.append(Path(item.location))
+            files[input_name] = val
+        elif isinstance(input_value, File):
+            if not input_value.location and not input_value.path:
+                raise NotImplementedError("File location is not defined.")
+            if not input_value.location:
+                continue
+            if input_value.location.startswith("lfn:"):
+                val.append(Path(input_value.location))
+            files[input_name] = val
+    return files

dirac_cwl/data_management_mocks/data_manager.py

@@ -0,0 +1,99 @@
+"""Mock DIRAC data manager for local file storage operations."""
+
+from pathlib import Path
+
+from DIRAC.DataManagementSystem.Client.DataManager import DataManager  # type: ignore[import-untyped]
+from DIRAC.Resources.Storage.FileStorage import FileStorage  # type: ignore[import-untyped]
+from DIRACCommon.Core.Utilities.ReturnValues import S_ERROR, S_OK, returnSingleResult  # type: ignore[import-untyped]
+
+from dirac_cwl.data_management_mocks.file_catalog import LocalFileCatalog
+
+
+class MockDataManager(DataManager):
+    """Mock DIRAC DataManager for local file storage."""
+
+    def __init__(self):
+        """Initialize the mock data manager with local file catalog."""
+        self.base_storage_path = "filecatalog"
+        self.storage_element = FileStorage("local", {"Path": self.base_storage_path})
+        self.fileCatalog = LocalFileCatalog()
+
+    def getFile(self, lfn, destinationDir=".", sourceSE=None, diskOnly=False):
+        """Get local copy of LFN(s) from Storage Elements.
+
+        :param mixed lfn: a single LFN or list of LFNs.
+        :param str destinationDir: directory to which the file(s) will be
+        downloaded. (Default: current working directory).
+        :param str sourceSE: source SE from which to download (Default: all replicas will be attempted).
+        :param bool diskOnly: chooses the disk ONLY replica(s). (Default: False)
+        :return: S_OK({"Successful": {}, "Failed": {}})/S_ERROR(errMessage).
+        """
+        if isinstance(lfn, list):
+            lfns = lfn
+        elif isinstance(lfn, str):
+            lfns = [lfn]
+        else:
+            return S_ERROR(f"wrong type for lfn: {lfn}, expected str or list[str]")
+
+        if not sourceSE:
+            sourceSE = self.storage_element
+
+        success = {}
+        fail = {}
+        for lfn in lfns:
+            res = sourceSE.getFile(
+                str(
+                    Path(self.base_storage_path) / str(lfn).removeprefix("lfn:").removeprefix("LFN:").removeprefix("/")
+                ),
+                destinationDir,
+            )
+            if not res["OK"]:
+                fail[lfn] = res["Message"]
+            elif res["Value"]["Failed"]:
+                fail[lfn] = res["Value"]["Failed"]
+            else:
+                success[lfn] = str(Path(destinationDir).resolve() / Path(lfn).name)
+        return S_OK({"Successful": success, "Failed": fail})
+
+    def putAndRegister(
+        self,
+        lfn,
+        fileName,
+        diracSE,
+        guid=None,
+        path=None,
+        checksum=None,
+        overwrite=None,
+    ):
+        """Put a local file to a Storage Element and register in the File Catalogues.
+
+        :param lfn: The file LFN.
+        :param fileName: The full path to the local file.
+        :param diracSE: The Storage Element to which to put the file.
+        :param guid: The guid with which the file is to be registered (if not provided will be generated).
+        :param path: The path on the storage where the file will be put (if not provided the LFN will be used).
+        :param checksum: File checksum (optional).
+        :param overwrite: Removes file from the file catalogue and SE before attempting upload.
+        """
+        self.fileCatalog.addFile(lfn)
+        return self.put(lfn, fileName, diracSE, path)
+
+    def put(self, lfn, fileName, diracSE, path=None):
+        """Put a local file to a Storage Element.
+
+        :param str lfn: LFN.
+        :param str fileName: The full path to the local file.
+        :param str diracSE: The Storage Element to which to put the file.
+        :param str path: The path on the storage where the file will be put (if not provided the LFN will be used).
+        :return: S_OK({"Successful": {...}, "Failed": {...}})/S_ERROR(errMessage).
+        """
+        se = self.storage_element
+        if not se:
+            return S_ERROR("No Storage Element defined")
+        if not path:
+            path = str(lfn).removeprefix("lfn:").removeprefix("LFN:").removeprefix("/")
+        dest = str(Path(self.base_storage_path) / Path(path))
+        res = returnSingleResult(se.putFile({dest: fileName}))
+        if not res["OK"]:
+            return S_OK({"Successful": {}, "Failed": {lfn: res["Message"]}})
+        return S_OK({"Successful": {lfn: res["Value"]}, "Failed": {}})

dirac_cwl/data_management_mocks/file_catalog.py

@@ -0,0 +1,132 @@
+"""Mock file catalog implementations for local testing."""
+
+import json
+import time
+from pathlib import Path
+
+from DIRAC import S_ERROR, S_OK  # type: ignore[import-untyped]
+from DIRAC.Resources.Catalog.FileCatalog import FileCatalog  # type: ignore[import-untyped]
+
+
+class InMemoryFileCatalog(FileCatalog):
+    """Minimal in-memory FileCatalog compatible with DIRAC DataManager."""
+
+    def __init__(self, catalogs=None, vo=None):
+        """Initialize the in-memory file catalog.
+
+        :param catalogs: Catalog configuration (unused).
+        :param vo: Virtual organization (unused).
+        """
+        self._eligibleCatalogs = {}
+        self._files = {}  # store metadata and logical file names
+        super(FileCatalog, self).__init__()
+
+    def _getEligibleCatalogs(self):
+        """Get eligible catalogs for this file catalog.
+
+        :return: S_OK with catalog configuration.
+        """
+        self._eligibleCatalogs = {"MyMockCatalog": {"Type": "MockFileCatalog", "Backend": "Memory"}}
+        return S_OK(self._eligibleCatalogs)
+
+    def findFile(self, lfn):
+        """Find a file in the catalog by LFN.
+
+        :param lfn: Logical file name.
+        :return: S_OK with file metadata or S_ERROR if not found.
+        """
+        if lfn in self._files:
+            return S_OK([self._files[lfn]])
+        return S_ERROR(f"File {lfn} not found")
+
+    def addFile(self, lfn, metadata=None):
+        """Add a file to the catalog.
+
+        :param lfn: Logical file name.
+        :param metadata: Optional file metadata.
+        :return: S_OK with LFN or S_ERROR if file already exists.
+        """
+        if lfn in self._files:
+            return S_ERROR(f"File {lfn} already exists")
+        self._files[lfn] = {"LFN": lfn, "Metadata": metadata or {}}
+        return S_OK(lfn)
+
+
+class LocalFileCatalog(FileCatalog):
+    """File catalog implementation using local filesystem storage."""
+
+    def __init__(self, catalogs=None, vo=None):
+        """Initialize the local file catalog.
+
+        :param catalogs: Catalog configuration (unused).
+        :param vo: Virtual organization (unused).
+        """
+        self._eligibleCatalogs = {"MyMockCatalog": {"Type": "MockFileCatalog", "Backend": "LocalFileSystem"}}
+        self._metadataPath = "filecatalog/metadata.json"
+        super(FileCatalog, self).__init__()
+
+    def _getEligibleCatalogs(self):
+        """Get eligible catalogs for this file catalog.
+
+        :return: S_OK with catalog configuration.
+        """
+        return S_OK(self._eligibleCatalogs)
+
+    def getFileMetadata(self, lfn):
+        """Get metadata for a file.
+
+        :param lfn: Logical file name.
+        :return: S_OK with metadata dict or failed dict.
+        """
+        metaAll = self._getAllMetadata()
+        if lfn not in metaAll:
+            return S_OK({"Successful": {}, "Failed": {lfn: f"File {lfn} not found"}})
+        return S_OK({"Successful": {lfn: metaAll[lfn]}, "Failed": {}})
+
+    def addFile(self, lfn):
+        """Add a file to the catalog.
+
+        :param lfn: Logical file name.
+        :return: S_OK with success/failed dict or S_ERROR if file exists.
+        """
+        if lfn in self._getAllMetadata():
+            return S_ERROR(f"File {lfn} already exists")
+        self.setMetadata(lfn, {"CreationDate": time.time()})
+        return S_OK({"Successful": {lfn: True}, "Failed": {}})
+
+    def setMetadata(self, lfn, metadataDict):
+        """Set metadata for a file.
+
+        :param lfn: Logical file name.
+        :param metadataDict: Metadata dictionary to set.
+        :return: S_OK with success/failed dict or S_ERROR on failure.
+        """
+        meta = self._getAllMetadata()
+        meta[lfn] = metadataDict
+
+        try:
+            self._setAllMetadata(meta)
+        except Exception as e:
+            return S_ERROR(f"Could set metadata: {e}")
+        return S_OK({"Successful": {lfn: True}, "Failed": {}})
+
+    def _getAllMetadata(self):
+        """Get all metadata from the local file.
+
+        :return: Dictionary of all file metadata.
+        """
+        try:
+            with open(self._metadataPath, "r") as file:
+                meta = json.load(file)
+        except Exception:
+            meta = {}
+        return meta
+
+    def _setAllMetadata(self, metadata):
+        """Save all metadata to the local file.
+
+        :param metadata: Dictionary of file metadata to save.
+        """
+        Path(self._metadataPath).parent.mkdir(parents=True, exist_ok=True)
+        with open(self._metadataPath, "w+") as file:
+            json.dump(metadata, file)

dirac_cwl/data_management_mocks/sandbox.py

@@ -0,0 +1,89 @@
+"""Mock DIRAC sandbox store client for local file operations."""
+
+import hashlib
+import logging
+import os
+import tarfile
+import tempfile
+from pathlib import Path
+from typing import Literal, Sequence
+
+import zstandard
+from diracx.core.models import SandboxInfo
+
+logger = logging.getLogger(__name__)
+
+SANDBOX_CHECKSUM_ALGORITHM = "sha256"
+SANDBOX_COMPRESSION: Literal["zst"] = "zst"
+
+# Get the project root directory (where pyproject.toml is located)
+PROJECT_ROOT = Path(__file__).parent.parent.parent.parent
+SANDBOX_STORE_DIR = PROJECT_ROOT / "sandboxstore"
+
+
+def create_sandbox(paths: Sequence[str | Path]):
+    """Upload a sandbox archive to the sandboxstore.
+
+    :param paths: File paths to be uploaded in the sandbox.
+    """
+    with tempfile.TemporaryFile(mode="w+b") as tar_fh:
+        # Create zstd compressed tar with level 18 and long matching enabled
+        compression_params = zstandard.ZstdCompressionParameters.from_level(18, enable_ldm=1)
+        cctx = zstandard.ZstdCompressor(compression_params=compression_params)
+        with cctx.stream_writer(tar_fh, closefd=False) as compressor:
+            with tarfile.open(fileobj=compressor, mode="w|") as tf:
+                for path in paths:
+                    if isinstance(path, str):
+                        path = Path(path)
+                    logger.debug("Adding %s to sandbox as %s", path.resolve(), path.name)
+                    tf.add(path.resolve(), path.name, recursive=True)
+        tar_fh.seek(0)
+
+        # Generate sandbox checksum
+        hasher = getattr(hashlib, SANDBOX_CHECKSUM_ALGORITHM)()
+        while data := tar_fh.read(512 * 1024):
+            hasher.update(data)
+        checksum = hasher.hexdigest()
+        tar_fh.seek(0)
+        logger.debug("Sandbox checksum is %s", checksum)
+
+        # Store sandbox info
+        sandbox_info = SandboxInfo(
+            checksum_algorithm=SANDBOX_CHECKSUM_ALGORITHM,
+            checksum=checksum,
+            size=os.stat(tar_fh.fileno()).st_size,
+            format=f"tar.{SANDBOX_COMPRESSION}",
+        )
+
+        # Create PFN
+        pfn = f"{sandbox_info.checksum_algorithm}:{sandbox_info.checksum}.{sandbox_info.format}"
+        logger.debug("Sandbox PFN is %s", pfn)
+
+        # Create sandbox in sandboxstore
+        SANDBOX_STORE_DIR.mkdir(exist_ok=True)
+        sandbox_path = SANDBOX_STORE_DIR / pfn
+        if not sandbox_path.exists():
+            with tarfile.open(sandbox_path, "w:gz") as tar:
+                for file in paths:
+                    if not file:
+                        break
+                    if isinstance(file, str):
+                        file = Path(file)
+                    tar.add(file, arcname=file.name)
+            logger.debug("Sandbox uploaded for %s", pfn)
+        else:
+            logger.debug("Sandbox already exists for %s", pfn)
+        return pfn
+
+
+def download_sandbox(pfn: str, destination: Path):
+    """Retrieve a sandbox from the sandboxstore and extract it to the given destination.
+
+    :param pfn: Sandbox PFN
+    :param destination: Destination directory
+    """
+    logger.debug("Retrieving sandbox for %s", pfn)
+    sandbox_archive = SANDBOX_STORE_DIR / pfn
+    with tarfile.open(sandbox_archive) as tf:
+        tf.extractall(path=Path(destination), filter="data")
+    logger.debug("Extracted %s to %s", pfn, destination)

dirac_cwl/execution_hooks/__init__.py

@@ -0,0 +1,40 @@
+"""Enhanced metadata registry for DIRAC CWL integration.
+
+This module provides a comprehensive plugin system for metadata models,
+supporting virtual organization-specific extensions and automatic discovery.
+
+The module maintains backward compatibility with the original API while
+providing enhanced functionality through the new plugin system.
+"""
+
+from .core import (
+    ExecutionHooksBasePlugin,
+    ExecutionHooksHint,
+    SchedulingHint,
+    TransformationExecutionHooksHint,
+)
+from .registry import (
+    ExecutionHooksPluginRegistry,
+    discover_plugins,
+    get_registry,
+)
+
+# Initialize the registry and discover core plugins
+_registry = get_registry()
+
+# Auto-discover plugins on import
+try:
+    discover_plugins()
+except Exception:
+    # Fail silently if plugin discovery fails during import
+    pass
+
+__all__ = [
+    # Core metadata and plugins
+    "ExecutionHooksHint",
+    "TransformationExecutionHooksHint",
+    "ExecutionHooksBasePlugin",
+    "SchedulingHint",
+    "ExecutionHooksPluginRegistry",
+    "get_registry",
+]