Merge pull request #263 from ssl-hep/UX_multiple_codegen

Merging the latest production release (2.6.2) changes into the master.

Author: talvandaalen

.vscode/settings.json

@@ -9,6 +9,7 @@
         "aiohttp",
         "AOD's",
         "asyncio",
+        "atlasxaod",
         "awks",
         "AZNLOCTEQ",
         "backend's",
@@ -19,6 +20,7 @@
         "caplog",
         "cernopendata",
         "codecov",
+        "codegen",
         "Comming",
         "coveragerc",
         "DAOD",

README.md

@@ -194,6 +194,10 @@ The file can contain an `api_endpoint` as mentioned earlier. In addition the oth
 
 All strings are expanded using python's [os.path.expand](https://docs.python.org/3/library/os.path.html#os.path.expandvars) method - so `$NAME` and `${NAME}` will work to expand existing environment variables.
 
+For non-standard use cases, the user can specify:
+- The code generator that is used by the backend. This is done by passing a `codegen` argument to ServiceXDataset. This argument is normally inherited from the backend type set in `servicex.yaml`, but can be overridden with any valid `codegen` contained in the default type listing. A `codegen` entry can also be added to a backend in the yaml file to use as default.
+- The type of backend, using the `backend_type` argument on ServiceXDataset. This overrides the backend type setting in the `servicex.yaml` file.
+
 ## Features
 
 Implemented:

servicex/config_default.yaml

@@ -31,7 +31,13 @@ default_return_data: parquet
 backend_types:
   - type: xaod
     return_data: root-file
+    codegen: atlasxaod
   - type: uproot
     return_data: parquet
+    codegen: uproot
   - type: cms_run1_aod
     return_data: root-file
+    codegen: cms
+  - type: python
+    return_data: parquet
+    codegen: python

servicex/servicex.py

@@ -175,6 +175,8 @@ def __init__(
         self,
         dataset: DatasetType,
         backend_name: Optional[str] = None,
+        backend_type: Optional[str] = None,
+        codegen: Optional[str] = None,
         image: Optional[str] = None,
         max_workers: int = 20,
         result_destination: str = "object-store",
@@ -201,6 +203,13 @@ def __init__(
                                         will default to xaod, unless you have any endpoint listed
                                         in your servicex file. It will default to best match there,
                                         or fail if a name has been given.
+            backend_type                The type of backend. Overrides the `type` in the `yaml`
+                                        config file.
+            codegen                     The type of code generator passed to the backend to
+                                        generate the code that powers the requested transform.
+                                        Don't use unless you know what you are doing - sometimes
+                                        the return filetype is also needed. Better to use
+                                        `backend_type`!
             image                       Name of transformer image to use to transform the data. If
                                         left as default, `None`, then the default image for the
                                         ServiceX backend will be used.
@@ -246,6 +255,7 @@ def __init__(
         ServiceXABC.__init__(
             self,
             dataset,
+            codegen,
             image,
             max_workers,
             result_destination,
@@ -272,7 +282,9 @@ def __init__(
 
         if not servicex_adaptor:
             # Given servicex adaptor is none, this should be ok. Fixes type checkers
-            end_point, token = config.get_servicex_adaptor_config(backend_name)
+            end_point, token = config.get_servicex_adaptor_config(
+                backend_name, backend_type=backend_type
+            )
             servicex_adaptor = ServiceXAdaptor(end_point, token)
         self._servicex_adaptor = servicex_adaptor
 
@@ -292,13 +304,24 @@ def __init__(
             else default_client_session
         )
 
-        self._return_types = [config.get_default_returned_datatype(backend_name)]
+        self._return_types = [
+            config.get_default_returned_datatype(
+                backend_name, backend_type=backend_type
+            )
+        ]
         self._converter = (
             data_convert_adaptor
             if data_convert_adaptor is not None
             else DataConverterAdaptor(self._return_types[0])
         )
 
+        # TODO - this should not be in the ABC backend since we have to have some intelligence
+        # in setting it, it turns out.
+        if self._codegen is None:
+            self._codegen = config.get_backend_info(
+                backend_name, "codegen", backend_type=backend_type
+            )
+
     def first_supported_datatype(
         self, datatypes: Union[List[str], str]
     ) -> Optional[str]:
@@ -1057,12 +1080,15 @@ def _build_json_query(
         assert data_format in g_allowed_formats
 
         # Items that must always be present
+        codegen = self._codegen if self._codegen is not None else "default"
+
         json_query: Dict[str, Union[str, Iterable[str]]] = {
             "selection": selection_query,
             "result-destination": self._result_destination,
             "result-format": data_format,
             "chunk-size": "1000",
             "workers": str(self._max_workers),
+            "codegen": codegen,
         }
 
         # Add the appropriate did.

servicex/servicex_config.py

@@ -16,7 +16,7 @@ class ServiceXConfigAdaptor:
     def __init__(self, config: Optional[ConfigView] = None):
         """The config needed for the app.
 
-        Note: The config is held onto and only queired when the information is required.
+        Note: The config is held onto and only queried when the information is required.
 
         Args:
             config (ConfigView): The config information for the app. If null, then we just use the
@@ -37,8 +37,10 @@ def settings(self) -> ConfigView:
         """
         return self._settings
 
-    def get_default_returned_datatype(self, backend_name: Optional[str]) -> str:
-        """Return the default return data type, given the backend is a certian type.
+    def get_default_returned_datatype(
+        self, backend_name: Optional[str], backend_type: Optional[str] = None
+    ) -> str:
+        """Return the default return data type, given the backend is a certain type.
 
         Args:
             backend_name (Optional[str]): The backend type string (`xaod`, `uproot`, etc)
@@ -47,7 +49,9 @@ def get_default_returned_datatype(self, backend_name: Optional[str]) -> str:
             str: The backend datatype, like `root` or `parquet`.
         """
         # Check to see if we know about the backend info
-        r = self.get_backend_info(backend_name, "return_data")
+        r = self.get_backend_info(
+            backend_name, "return_data", backend_type=backend_type
+        )
 
         if r is None:
             raise ServiceXException(
@@ -57,22 +61,27 @@ def get_default_returned_datatype(self, backend_name: Optional[str]) -> str:
 
         return r
 
-    def get_backend_info(self, backend_name: Optional[str], key: str) -> Optional[str]:
+    def get_backend_info(
+        self, backend_name: Optional[str], key: str, backend_type: Optional[str] = None
+    ) -> Optional[str]:
         """Find an item in the backend info, searching first for the backend
         name/type and then the defaults with a given type.
 
         Args:
             backend_name (str): Backend name
             key (str): The key for the info we are after
+            backend_type (str): The backend type, if we are overriding normal lookup
 
         Returns:
             Optional[str]: Return a string for the info we are after, or return None if we can't
                            find it.
         """
-        config = self._get_backend_info(backend_name)
+        config = self._get_backend_info(backend_name, backend_type)
         return config[key] if key in config else None
 
-    def _get_backend_info(self, backend_name: Optional[str]) -> Dict[str, str]:
+    def _get_backend_info(
+        self, backend_name: Optional[str], backend_type: Optional[str]
+    ) -> Dict[str, str]:
         """Returns all the info for a backend name/type.
 
         Search algorithm is non-trivial:
@@ -89,6 +98,8 @@ def _get_backend_info(self, backend_name: Optional[str]) -> Dict[str, str]:
 
         Args:
             backend_name (str): Name or type of the api end point we are going to look up.
+            backend_type (str): Override the backend type, if we are looking for a specific
+                                type behavior.
 
         Returns:
             Dict[str, str]: Attributes for this backend's configuration
@@ -159,8 +170,14 @@ def _get_backend_info(self, backend_name: Optional[str]) -> Dict[str, str]:
             )
 
         # Now, extract the type and see if we can figure out any defaults from the
-        # `backend_types` info.
-        type_lookup = config["type"] if "type" in config else backend_name
+        # `backend_types` info. Skip this if we have a type we are passed in.
+        type_lookup = (
+            backend_type
+            if backend_type is not None
+            else config["type"]
+            if "type" in config
+            else backend_name
+        )
         if type_lookup is None:
             return config
 
@@ -181,7 +198,9 @@ def _get_backend_info(self, backend_name: Optional[str]) -> Dict[str, str]:
         return config
 
     def get_servicex_adaptor_config(
-        self, backend_name: Optional[str] = None
+        self,
+        backend_name: Optional[str] = None,
+        backend_type: Optional[str] = None,
     ) -> Tuple[str, Optional[str]]:
         """Return the servicex (endpoint, token) from a given backend configuration.
 
@@ -193,7 +212,7 @@ def get_servicex_adaptor_config(
             Tuple[str, str]: The tuple of info to create a `ServiceXAdaptor`: end point,
             token (optionally).
         """
-        config = self._get_backend_info(backend_name)
+        config = self._get_backend_info(backend_name, backend_type)
 
         endpoint = config["endpoint"]
         token = config["token"] if "token" in config else None

servicex/servicexabc.py

@@ -29,6 +29,7 @@ class ServiceXABC(ABC):
     def __init__(
         self,
         dataset: DatasetType,
+        codegen: Optional[str] = None,
         image: Optional[str] = None,
         max_workers: int = 20,
         result_destination: str = "object-store",
@@ -40,6 +41,8 @@ def __init__(
         Arguments
 
             dataset                     Name of a dataset from which queries will be selected.
+            codegen                     The type of code generator passed to the backend to
+                                        generate the code that powers the requested transform.
             image                       Name of transformer image to use to transform the data. If
                                         None the default implementation is used.
             cache_adaptor               Runs the caching for data and queries that are sent up and
@@ -65,6 +68,7 @@ def __init__(
                be processed (and some files can even be completed before that is known).
         """
         self._dataset = dataset
+        self._codegen = codegen
         self._image = image
         self._max_workers = max_workers
         self._result_destination = result_destination

tests/test_servicex.py

@@ -49,8 +49,12 @@ def test_default_ctor(mocker):
 
     fe.ServiceXDataset("localds://dude", "uproot-ftw", config_adaptor=config)
 
-    config.get_servicex_adaptor_config.assert_called_with("uproot-ftw")
-    config.get_default_returned_datatype.assert_called_with("uproot-ftw")
+    config.get_servicex_adaptor_config.assert_called_with(
+        "uproot-ftw", backend_type=None
+    )
+    config.get_default_returned_datatype.assert_called_with(
+        "uproot-ftw", backend_type=None
+    )
 
 
 def test_sx_name(mocker):
@@ -79,8 +83,8 @@ def test_default_ctor_no_type(mocker):
 
     fe.ServiceXDataset("localds://dude", config_adaptor=config)
 
-    config.get_servicex_adaptor_config.assert_called_with(None)
-    config.get_default_returned_datatype.assert_called_with(None)
+    config.get_servicex_adaptor_config.assert_called_with(None, backend_type=None)
+    config.get_default_returned_datatype.assert_called_with(None, backend_type=None)
 
 
 def test_default_ctor_cache(mocker):
@@ -166,6 +170,21 @@ def test_get_datatypes_good(mocker):
     assert r.first_supported_datatype(["root", "parquet"]) == "root"
 
 
+def test_get_datatypes_background_type(mocker):
+    "Test that we return a good datatype"
+
+    # config = mocker.MagicMock(spec=ServiceXConfigAdaptor)
+    # config.settings = Configuration("servicex", "servicex")
+    # config.get_servicex_adaptor_config.return_value = (
+    #     "http://no-way.dude",
+    #     "no_spoon_there_is",
+    # )
+
+    r = fe.ServiceXDataset("localds://dude", "uproot", backend_type="xaod")
+
+    assert r.first_supported_datatype(["root-file", "parquet"]) == "root-file"
+
+
 def test_get_datatypes_single(mocker):
     "Test that we return a good datatype"
 
@@ -716,8 +735,8 @@ async def test_stream_root_files_from_minio(mocker):
     assert r.file == "one_minio_entry"
     assert r.url == "http://the.url.com"
 
-    assert mock_servicex_adaptor.query_json["result-format"] == "root-file"
     assert mock_minio_adaptor.access_called_with == ("123-456", "one_minio_entry")
+    assert mock_servicex_adaptor.query_json["result-format"] == "root-file"
 
 
 @pytest.mark.asyncio
@@ -1300,6 +1319,77 @@ async def test_no_title_spec(mocker, good_awkward_file_data):
     assert "title" not in called
 
 
+@pytest.mark.asyncio
+async def test_codegen_override(mocker, good_awkward_file_data):
+    mock_cache = build_cache_mock(mocker)
+    mock_logger = mocker.MagicMock(spec=log_adaptor)
+    mock_servicex_adaptor = MockServiceXAdaptor(mocker, "123-456")
+    mock_minio_adaptor = MockMinioAdaptor(mocker, files=["one_minio_entry"])
+    data_adaptor = mocker.MagicMock(spec=DataConverterAdaptor)
+
+    ds = fe.ServiceXDataset(
+        "localds://mc16_tev:13",
+        servicex_adaptor=mock_servicex_adaptor,  # type: ignore
+        minio_adaptor=mock_minio_adaptor,  # type: ignore
+        cache_adaptor=mock_cache,
+        data_convert_adaptor=data_adaptor,
+        local_log=mock_logger,
+        max_workers=50,
+        codegen="good_codegen",
+    )
+    await ds.get_data_rootfiles_async("(valid qastle string)")
+
+    called = mock_servicex_adaptor.query_json
+    assert called["codegen"] == "good_codegen"
+
+
+@pytest.mark.asyncio
+async def test_codegen_backend_type(mocker, good_awkward_file_data):
+    mock_cache = build_cache_mock(mocker)
+    mock_logger = mocker.MagicMock(spec=log_adaptor)
+    mock_servicex_adaptor = MockServiceXAdaptor(mocker, "123-456")
+    mock_minio_adaptor = MockMinioAdaptor(mocker, files=["one_minio_entry"])
+    data_adaptor = mocker.MagicMock(spec=DataConverterAdaptor)
+
+    ds = fe.ServiceXDataset(
+        "localds://mc16_tev:13",
+        servicex_adaptor=mock_servicex_adaptor,  # type: ignore
+        minio_adaptor=mock_minio_adaptor,  # type: ignore
+        cache_adaptor=mock_cache,
+        data_convert_adaptor=data_adaptor,
+        local_log=mock_logger,
+        max_workers=50,
+        backend_type="uproot",
+    )
+    await ds.get_data_rootfiles_async("(valid qastle string)")
+
+    called = mock_servicex_adaptor.query_json
+    assert called["codegen"] == "uproot"
+
+
+@pytest.mark.asyncio
+async def test_codegen_default_by_backend(mocker, good_awkward_file_data):
+    mock_cache = build_cache_mock(mocker)
+    mock_logger = mocker.MagicMock(spec=log_adaptor)
+    mock_servicex_adaptor = MockServiceXAdaptor(mocker, "123-456")
+    mock_minio_adaptor = MockMinioAdaptor(mocker, files=["one_minio_entry"])
+    data_adaptor = mocker.MagicMock(spec=DataConverterAdaptor)
+
+    ds = fe.ServiceXDataset(
+        "localds://mc16_tev:13",
+        servicex_adaptor=mock_servicex_adaptor,  # type: ignore
+        minio_adaptor=mock_minio_adaptor,  # type: ignore
+        cache_adaptor=mock_cache,
+        data_convert_adaptor=data_adaptor,
+        local_log=mock_logger,
+        max_workers=50,
+    )
+    await ds.get_data_rootfiles_async("(valid qastle string)")
+
+    called = mock_servicex_adaptor.query_json
+    assert called["codegen"] == "atlasxaod"
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "n_ds, n_query", [(1, 4), (4, 1), (1, 100), (100, 1), (4, 4), (20, 20)]