Merge pull request #406 from ssl-hep/3.0_develop_add_dataset

Add dataset keyword

Author: kyungeonchoi

pyproject.toml

@@ -101,6 +101,11 @@ FuncADL_CMS = "servicex.func_adl.func_adl_dataset:FuncADLQuery_CMS"
 PythonFunction = "servicex.python_dataset:PythonQuery"
 UprootRaw = "servicex.uproot_raw.uproot_raw:UprootRawQuery"
 
+[project.entry-points.'servicex.dataset']
+Rucio = "servicex.dataset_identifier:RucioDatasetIdentifier"
+FileList = "servicex.dataset_identifier:FileListDataset"
+CERNOpenData = "servicex.dataset_identifier:CERNOpenDataDatasetIdentifier"
+
 [tool.hatch.build.targets.sdist]
 # hatchling always includes:
 # pyproject.toml, .gitignore, any README, any LICENSE, AUTHORS

servicex/__init__.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2022, IRIS-HEP
+# Copyright (c) 2024, IRIS-HEP
 # All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
@@ -38,7 +38,8 @@
 from .models import ResultFormat, ResultDestination
 from .dataset_group import DatasetGroup
 from .dataset_identifier import RucioDatasetIdentifier, FileListDataset
-from . import query
+import servicex.dataset as dataset
+import servicex.query as query
 
 __all__ = [
     "ServiceXClient",
@@ -60,5 +61,6 @@
     "DefinitionList",
     "ServiceXSpec",
     "deliver",
+    "dataset",
     "query"
 ]

servicex/databinder_models.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2022, IRIS-HEP
+# Copyright (c) 2024, IRIS-HEP
 # All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
@@ -33,7 +33,8 @@
     model_validator,
 )
 
-from servicex.dataset_identifier import RucioDatasetIdentifier, FileListDataset
+from servicex.dataset_identifier import (DataSetIdentifier, RucioDatasetIdentifier,
+                                         FileListDataset)
 from servicex.query_core import Query as SXQuery, QueryStringGenerator
 from servicex.models import ResultFormat
 
@@ -43,31 +44,37 @@ class Sample(BaseModel):
     Codegen: Optional[str] = None
     RucioDID: Optional[str] = None
     XRootDFiles: Optional[Union[str, List[str]]] = None
+    Dataset: Optional[DataSetIdentifier] = None
     NFiles: Optional[int] = Field(default=None)
     Query: Optional[Union[str, SXQuery, QueryStringGenerator]] = Field(default=None)
     IgnoreLocalCache: bool = False
 
     model_config = {"arbitrary_types_allowed": True}
 
     @property
-    def dataset_identifier(self):
-        if self.RucioDID:
-            return RucioDatasetIdentifier(self.RucioDID, num_files=self.NFiles or 0)
+    def dataset_identifier(self) -> DataSetIdentifier:
+        if self.Dataset:
+            return self.Dataset
+        elif self.RucioDID:
+            return RucioDatasetIdentifier(self.RucioDID, num_files=self.NFiles)
         elif self.XRootDFiles:
             return FileListDataset(self.XRootDFiles)
+        else:  # pragma: no cover
+            raise RuntimeError("No valid dataset found, somehow validation failed")
 
     @model_validator(mode="before")
     @classmethod
     def validate_did_xor_file(cls, values):
         """
-        Ensure that only one of RootFile or RucioDID is specified.
+        Ensure that only one of Dataset, RootFile, or RucioDID is specified.
         :param values:
         :return:
         """
-        if "XRootDFiles" in values and "RucioDID" in values:
-            raise ValueError("Only specify one of XRootDFiles or RucioDID, not both.")
-        if "XRootDFiles" not in values and "RucioDID" not in values:
-            raise ValueError("Must specify one of XRootDFiles or RucioDID.")
+        count = sum(["RucioDID" in values, "XRootDFiles" in values, "Dataset" in values])
+        if count > 1:
+            raise ValueError("Only specify one of Dataset, XRootDFiles, or RucioDID.")
+        if count == 0:
+            raise ValueError("Must specify one of Dataset, XRootDFiles, or RucioDID.")
         return values
 
 

servicex/dataset/__init__.py

@@ -0,0 +1,35 @@
+# Copyright (c) 2024, IRIS-HEP
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+from ..dataset_identifier import (RucioDatasetIdentifier as Rucio,  # noqa: F401
+                                  FileListDataset as FileList,
+                                  CERNOpenDataDatasetIdentifier as CERNOpenData,
+                                  DataSetIdentifier as GenericDataSet)
+
+
+__any__ = ['Rucio', 'FileList', 'CERNOpenData', 'GenericDataSet']

servicex/dataset_identifier.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2022, IRIS-HEP
+# Copyright (c) 2024, IRIS-HEP
 # All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
@@ -25,7 +25,7 @@
 # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-from typing import List, Union
+from typing import List, Union, Optional
 
 from servicex.models import TransformRequest
 
@@ -35,14 +35,14 @@ class DataSetIdentifier:
     Base class for specifying the dataset to transform. This can either be a list of
     xRootD URIs or a rucio DID
     """
-    def __init__(self, scheme: str, dataset: str, num_files: int = None):
+    def __init__(self, scheme: str, dataset: str, num_files: Optional[int] = None):
         self.scheme = scheme
         self.dataset = dataset
         self.num_files = num_files
 
     @property
     def did(self):
-        num_files_arg = f"?files={self.num_files}" if self.num_files else ""
+        num_files_arg = f"?files={self.num_files}" if self.num_files is not None else ""
         return f"{self.scheme}://{self.dataset}{num_files_arg}"
 
     def populate_transform_request(self, transform_request: TransformRequest) -> None:
@@ -51,7 +51,7 @@ def populate_transform_request(self, transform_request: TransformRequest) -> Non
 
 
 class RucioDatasetIdentifier(DataSetIdentifier):
-    def __init__(self, dataset: str, num_files: int = None):
+    def __init__(self, dataset: str, num_files: Optional[int] = None):
         r"""
         Rucio Dataset - this will be looked up using the Rucio data management
         service.
@@ -64,19 +64,56 @@ def __init__(self, dataset: str, num_files: int = None):
         """
         super().__init__("rucio", dataset, num_files=num_files)
 
+    yaml_tag = '!Rucio'
 
-class FileListDataset:
+    @classmethod
+    def from_yaml(cls, _, node):
+        return cls(node.value)
+
+
+class FileListDataset(DataSetIdentifier):
     def __init__(self, files: Union[List[str], str]):
         r"""
         Dataset specified as a list of XRootD URIs.
 
         :param files: Either a list of URIs or a single URI string
         """
-        if type(files) is str:
+        self.files: List[str]
+        if isinstance(files, str):
             self.files = [files]
         else:
             self.files = files
 
     def populate_transform_request(self, transform_request: TransformRequest) -> None:
         transform_request.file_list = self.files
         transform_request.did = None
+
+    @property
+    def did(self):
+        return None
+
+    yaml_tag = '!FileList'
+
+    @classmethod
+    def from_yaml(cls, constructor, node):
+        return cls(constructor.construct_sequence(node))
+
+
+class CERNOpenDataDatasetIdentifier(DataSetIdentifier):
+    def __init__(self, dataset: int, num_files: Optional[int] = None):
+        r"""
+        CERN Open Data Dataset - this will be looked up using the CERN Open Data DID finder.
+
+        :param dataset: The dataset ID - this is an integer.
+        :param num_files: Maximum number of files to return. This is useful during development
+            to perform quick runs. ServiceX is careful to make sure it always
+            returns the same subset of files.
+
+        """
+        super().__init__("cernopendata", f'{dataset}', num_files=num_files)
+
+    yaml_tag = '!CERNOpenData'
+
+    @classmethod
+    def from_yaml(cls, _, node):
+        return cls(int(node.value))

servicex/query/__init__.py

@@ -26,7 +26,6 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-
 import sys
 if sys.version_info < (3, 10):
     from importlib_metadata import entry_points

servicex/servicex_client.py

@@ -77,6 +77,9 @@ def _load_ServiceXSpec(
             from importlib.metadata import entry_points
 
         plugins = entry_points(group="servicex.query")
+        for _ in plugins:
+            yaml.register_class(_.load())
+        plugins = entry_points(group="servicex.dataset")
         for _ in plugins:
             yaml.register_class(_.load())
 

tests/test_databinder.py

@@ -2,14 +2,11 @@
 from unittest.mock import patch
 from pydantic import ValidationError
 
-from servicex import ServiceXSpec, FileListDataset, RucioDatasetIdentifier
+from servicex import ServiceXSpec, FileListDataset, RucioDatasetIdentifier, dataset
 
 
 def basic_spec(samples=None):
     return {
-        "General": {
-            "Codegen": "python",
-        },
         "Sample": samples
         or [{"Name": "sampleA", "XRootDFiles": "root://a.root", "Query": "a"}],
     }
@@ -35,7 +32,6 @@ def test_load_config():
 
 
 def test_single_root_file():
-
     spec = ServiceXSpec.model_validate(
         basic_spec(
             samples=[
@@ -52,6 +48,7 @@ def test_single_root_file():
     assert spec.Sample[0].dataset_identifier.files == [
         "root://eospublic.cern.ch//file1.root"
     ]
+    assert spec.Sample[0].dataset_identifier.did is None
 
 
 def test_list_of_root_files():
@@ -118,6 +115,19 @@ def test_rucio_did_numfiles():
     )
 
 
+def test_cernopendata():
+    spec = ServiceXSpec.model_validate({
+        "Sample": [
+            {
+                "Name": "sampleA",
+                "Dataset": dataset.CERNOpenData(1507),
+                "Function": "a"
+            }
+        ]
+    })
+    assert spec.Sample[0].dataset_identifier.did == "cernopendata://1507"
+
+
 def test_invalid_dataset_identifier():
     with pytest.raises(ValidationError):
         ServiceXSpec.model_validate(
@@ -172,6 +182,7 @@ def test_submit_mapping(transformed_result, codegen_list):
 
 def test_yaml(tmp_path):
     from servicex.servicex_client import _load_ServiceXSpec
+    from servicex.dataset import FileList, Rucio, CERNOpenData
     # Nominal paths
     with open(path := (tmp_path / "python.yaml"), "w") as f:
         f.write("""
@@ -194,12 +205,29 @@ def run_query(input_filenames=None):
     RucioDID: user.kchoi:user.kchoi.fcnc_tHq_ML.ttH.v11
     Query: !UprootRaw |
                 [{"treename": "nominal"}]
+  - Name: ttH4
+    Dataset: !Rucio user.kchoi:user.kchoi.fcnc_tHq_ML.ttH.v11
+    Query: !UprootRaw '[{"treename": "nominal"}]'
+  - Name: ttH5
+    Dataset: !FileList ["/path/to/file1.root", "/path/to/file2.root"]
+    Query: !UprootRaw '[{"treename": "nominal"}]'
+  - Name: ttH6
+    Dataset: !CERNOpenData 1507
+    Query: !UprootRaw '[{"treename": "nominal"}]'
 """)
         f.flush()
         result = _load_ServiceXSpec(path)
         assert type(result.Sample[0].Query).__name__ == 'PythonQuery'
         assert type(result.Sample[1].Query).__name__ == 'FuncADLQuery_Uproot'
         assert type(result.Sample[2].Query).__name__ == 'UprootRawQuery'
+        assert isinstance(result.Sample[3].dataset_identifier, Rucio)
+        assert (result.Sample[3].dataset_identifier.did
+                == 'rucio://user.kchoi:user.kchoi.fcnc_tHq_ML.ttH.v11')
+        assert isinstance(result.Sample[4].dataset_identifier, FileList)
+        assert (result.Sample[4].dataset_identifier.files
+                == ["/path/to/file1.root", "/path/to/file2.root"])
+        assert isinstance(result.Sample[5].dataset_identifier, CERNOpenData)
+        assert result.Sample[5].dataset_identifier.did == 'cernopendata://1507'
 
     # Path from string
     result2 = _load_ServiceXSpec(str(path))