Input/Output Sandbox/Data Management

src/dirac_cwl_proto/data_management_mocks/data_manager.py

@@ -0,0 +1,100 @@
+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_proto.data_management_mocks.file_catalog import LocalFileCatalog
+
+
+class MockDataManager(DataManager):
+    """
+    Mock DIRAC DataManager for local file storage
+    """
+
+    def __init__(self):
+        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"]
+            else:
+                success[lfn] = 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
+
+        'lfn' is the file LFN
+        'file' is the full path to the local file
+        'diracSE' is the Storage Element to which to put the file
+        'guid' is the guid with which the file is to be registered (if not provided will be generated)
+        'path' is the path on the storage where the file will be put (if not provided the LFN will be used)
+        '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 self: self reference
+        :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)
+
+        """
+        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": {}})

src/dirac_cwl_proto/data_management_mocks/file_catalog.py

@@ -0,0 +1,79 @@
+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):
+        self._eligibleCatalogs = {}
+        self._files = {}  # store metadata and logical file names
+        super(FileCatalog, self).__init__()
+
+    def _getEligibleCatalogs(self):
+        self._eligibleCatalogs = {
+            "MyMockCatalog": {"Type": "MockFileCatalog", "Backend": "Memory"}
+        }
+        return S_OK(self._eligibleCatalogs)
+
+    def findFile(self, lfn):
+        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):
+        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):
+    def __init__(self, catalogs=None, vo=None):
+        self._eligibleCatalogs = {
+            "MyMockCatalog": {"Type": "MockFileCatalog", "Backend": "LocalFileSystem"}
+        }
+        self._metadataPath = "filecatalog/metadata.json"
+        super(FileCatalog, self).__init__()
+
+    def _getEligibleCatalogs(self):
+        return S_OK(self._eligibleCatalogs)
+
+    def getFileMetadata(self, lfn):
+        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):
+        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):
+        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):
+        try:
+            with open(self._metadataPath, "r") as file:
+                meta = json.load(file)
+        except Exception:
+            meta = {}
+        return meta
+
+    def _setAllMetadata(self, metadata):
+        Path(self._metadataPath).parent.mkdir(parents=True, exist_ok=True)
+        with open(self._metadataPath, "w+") as file:
+            json.dump(metadata, file)

src/dirac_cwl_proto/data_management_mocks/sandbox.py

@@ -0,0 +1,91 @@
+import logging
+import random
+import tarfile
+from pathlib import Path
+from typing import Optional, Sequence
+
+from DIRAC.WorkloadManagementSystem.Client.SandboxStoreClient import SandboxStoreClient  # type: ignore[import-untyped]
+from DIRACCommon.Core.Utilities.ReturnValues import S_OK  # type: ignore[import-untyped]
+
+logger = logging.getLogger(__name__)
+
+
+class MockSandboxStoreClient(SandboxStoreClient):
+    """
+    Local mock for Dirac's SandboxStore Client.
+    """
+
+    def __init__(self):
+        pass
+
+    def uploadFilesAsSandbox(
+        self,
+        fileList: Sequence[Path | str],
+        sizeLimit: int = 0,
+        assignTo: Optional[dict] = None,
+    ):
+        """Create and upload a sandbox archive from a list of files.
+
+        Packages the provided files into a compressed tar archive and stores
+        it under the local sandbox directory.
+
+        :param Sequence[Path | str] fileList: Files to be included in the sandbox.
+        :param int sizeLimit: Maximum allowed archive size in bytes. Currently unused.
+        :param Optional[dict] assignTo: Mapping of job identifiers to sandbox types (e.g. { 'Job:<id>': '<type>' }).
+
+        :return S_OK(sandbox_path): Path to the created sandbox archive, or `None` if no files were provided.
+        """
+        if len(fileList) == 0:
+            return S_OK()
+        sandbox_id = random.randint(1000, 9999)
+        sandbox_path = Path("sandboxstore") / f"sandbox_{str(sandbox_id)}.tar.gz"
+        sandbox_path.parent.mkdir(exist_ok=True, parents=True)
+        with tarfile.open(sandbox_path, "w:gz") as tar:
+            for file in fileList:
+                if not file:
+                    break
+                if isinstance(file, str):
+                    file = Path(file)
+                tar.add(file, arcname=file.name)
+        res = S_OK(str(sandbox_path))
+        res["SandboxFileName"] = f"sandbox_{str(sandbox_id)}.tar.gz"
+        return res
+
+    def downloadSandbox(
+        self,
+        sbLocation: str | Path,
+        destinationDir: str = "",
+        inMemory: bool = False,
+        unpack: bool = True,
+    ) -> list[Path]:
+        """Download and extract files from a sandbox archive.
+
+        Opens the given sandbox archive and extracts its contents to the specified
+        directory.
+
+        :param str|Path sbLocation: Path to the sandbox archive file.
+        :param str destinationDir: Directory to extract the files into. Defaults to the current directory.
+        :param bool inMemory: Placeholder for in-memory extraction.
+        :param bool unpack: Whether to unpack the archive. Only unpacking is currently supported.
+
+        :return S_OK({list[Path]}): List of paths to the extracted files.
+        """
+        if not unpack or inMemory:
+            raise NotImplementedError
+        else:
+            sandbox_path = Path("sandboxstore") / f"{sbLocation}.tar.gz"
+            with tarfile.open(sandbox_path, "r:gz") as tar:
+                tar.extractall(destinationDir, filter="data")
+                files = tar.getnames()
+            logger.info("Files downloaded successfully!")
+            return S_OK([str(Path(destinationDir) / file) for file in files])
+
+    def downloadSandboxForJob(
+        self, jobId, sbType, destinationPath="", inMemory=False, unpack=True
+    ) -> None:
+        """
+        Download sandbox contents for a specific job.
+
+        Placeholder for future implementation of job-based sandbox retrieval.
+        """
+        raise NotImplementedError

src/dirac_cwl_proto/execution_hooks/__init__.py

@@ -8,7 +8,6 @@
 """
 
 from .core import (
-    DataCatalogInterface,
     ExecutionHooksBasePlugin,
     ExecutionHooksHint,
     SchedulingHint,
@@ -36,7 +35,6 @@
     "TransformationExecutionHooksHint",
     "ExecutionHooksBasePlugin",
     "SchedulingHint",
-    "DataCatalogInterface",
     "ExecutionHooksPluginRegistry",
     "get_registry",
 ]