support custom science images

Procfile

@@ -1,4 +1,5 @@
-minikube-mount: minikube mount $LOCAL_DIR:/mnt/servicex
+minikube-mount-app: minikube mount $LOCAL_DIR:/mnt/servicex
+minikube-mount-topcp: minikube mount $LOCAL_DIR/code_generator_TopCPToolkit:/mnt/topcp
 helm-install: sleep 5; cd $CHART_DIR && helm install -f $VALUES_FILE servicex . && tail -f /dev/null
 port-forward-app: sleep 30 && cd $LOCAL_DIR && bash local/port-forward.sh app
 port-forward-minio: sleep 20 && cd $LOCAL_DIR && bash local/port-forward.sh minio

code_generator_TopCPToolkit/servicex/TopCP_code_generator/query_translate.py

@@ -1,4 +1,3 @@
-import json
 import os
 
 options = {
@@ -36,12 +35,15 @@
         "ifTrue": ["--no-filter"],
         "ifFalse": None,
     },
+    "image": {
+        "properType": str,
+        "properTypeString": "string",
+        "optional": True,
+    },
 }
 
 
-def generate_files_from_query(query, query_file_path):
-    jquery = json.loads(query)
-
+def generate_files_from_query(jquery: dict, query_file_path):
     runTopCommand = [
         "runTop_el.py",
         "-i",
@@ -52,16 +54,19 @@ def generate_files_from_query(query, query_file_path):
         "customConfig",
     ]
 
-    # ensure all keys are specified
+    # ensure all required keys are specified
     for key in options:
         if key not in jquery:
+            # Skip optional parameters
+            if options[key].get("optional", False):
+                continue
             raise ValueError(
                 key + " must be specified. May be type None or ",
                 options[key]["properTypeString"],
             )
 
     for key in jquery:
-        # ensure only aviable options are allowed
+        # ensure only available options are allowed
         if key not in options:
             raise KeyError(
                 key + " is not implemented. Available keys: " + str(options.keys())
@@ -78,7 +83,7 @@ def generate_files_from_query(query, query_file_path):
             )
 
         # check for reco.yaml, parton.yaml and particle.yaml files
-        if isinstance(jquery[key], str):
+        if isinstance(jquery[key], str) and "fileName" in options[key]:
             with open(
                 os.path.join(query_file_path, options[key]["fileName"]), "w"
             ) as file:

code_generator_TopCPToolkit/servicex/TopCP_code_generator/request_translator.py

@@ -26,7 +26,9 @@
 # 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 os
+import json
 import shutil
+
 from . import query_translate
 from servicex_codegen.code_generator import (
     CodeGenerator,
@@ -62,11 +64,18 @@ def generate_code(self, query, cache_path: str):
         capabilities_path = os.environ.get(
             "CAPABILITIES_PATH", "/home/servicex/transformer_capabilities.json"
         )
+
+        jquery = json.loads(query)
+        query_translate.generate_files_from_query(jquery, query_file_path)
+
         shutil.copyfile(
             capabilities_path,
             os.path.join(query_file_path, "transformer_capabilities.json"),
         )
 
-        query_translate.generate_files_from_query(query, query_file_path)
+        results = GeneratedFileResult(_hash, query_file_path)
+
+        if jquery.get("image"):
+            results.image = jquery["image"]
 
-        return GeneratedFileResult(_hash, query_file_path)
+        return results

code_generator_TopCPToolkit/servicex/templates/transform_single_file.py

@@ -12,21 +12,23 @@ def transform_single_file(file_path: str, output_path: Path, output_format: str)
     # create input.txt file for event loop and insert file_path as only line
     with open("input.txt", "w") as f:
         f.write(file_path)
-    # move reco.yaml, parton.yaml and particle.yaml if they exit to CONFIG_LOC loacation
+
+    # move reco.yaml, parton.yaml and particle.yaml if they exist to CONFIG_LOC location
+    config_loc = os.environ.get("CONFIG_LOC", os.getcwd())
     if os.path.exists("/generated/reco.yaml"):
         shutil.copyfile(
             "/generated/reco.yaml",
-            os.path.join(os.environ.get("CONFIG_LOC"), "reco.yaml"),
+            os.path.join(config_loc, "reco.yaml"),
         )
     if os.path.exists("/generated/parton.yaml"):
         shutil.copyfile(
             "/generated/parton.yaml",
-            os.path.join(os.environ.get("CONFIG_LOC"), "parton.yaml"),
+            os.path.join(config_loc, "parton.yaml"),
         )
     if os.path.exists("/generated/particle.yaml"):
         shutil.copyfile(
             "/generated/particle.yaml",
-            os.path.join(os.environ.get("CONFIG_LOC"), "particle.yaml"),
+            os.path.join(config_loc, "particle.yaml"),
         )
 
     generated_transformer.runTop_el()

code_generator_TopCPToolkit/tests/test_src.py

@@ -88,6 +88,98 @@ def test_generate_code():
             translator.generate_code(query, tmpdirname)
 
 
+def test_generate_code_fails_with_unknown_selection_key():
+    os.environ["TEMPLATE_PATH"] = "servicex/templates/transform_single_file.py"
+    os.environ["CAPABILITIES_PATH"] = "transformer_capabilities.json"
+
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        translator = TopCPTranslator()
+        query = (
+            '{"reco": "CommonServices:\\n  systematicsHistogram: \'listOfSystematics\'\\n\\n'
+            "PileupReweighting: {}\\n\\nEventCleaning:\\n    runEventCleaning: False\\n"
+            "    runGRL: False\\n\\nElectrons:\\n  - containerName: 'AnaElectrons'\\n"
+            "    crackVeto: True\\n    IFFClassification: {}\\n    WorkingPoint:\\n"
+            "      - selectionName: 'loose'\\n        identificationWP: 'TightLH'\\n"
+            "        isolationWP: 'NonIso'\\n        noEffSF: True\\n"
+            "      - selectionName: 'tight'\\n        identificationWP: 'TightLH'\\n"
+            "        isolationWP: 'Tight_VarRad'\\n        noEffSF: True\\n"
+            "    PtEtaSelection:\\n        minPt: 25000.0\\n        maxEta: 2.47\\n"
+            "        useClusterEta: True\\n\\n"
+            "# After configuring each container, many variables will be saved automatically.\\n"
+            "Output:\\n  treeName: 'reco'\\n  vars: []\\n  metVars: []\\n  containers:\\n"
+            "      # Format should follow: '<suffix>:<output container>'\\n"
+            "      el_: 'AnaElectrons'\\n      '': 'EventInfo'\\n  commands:\\n"
+            "    # Turn output branches on and off with 'enable' and 'disable'\\n\\n"
+            'AddConfigBlocks: []\\n", "parton": null, "particle": null, "max_events": 100, '
+            '"no_systematics": true, "no_filter": false, "unknown_key": "unknown_value"}'
+        )
+
+        with pytest.raises(KeyError):
+            translator.generate_code(query, tmpdirname)
+
+
+def test_generate_code_fails_with_missing_required_selection_key():
+    os.environ["TEMPLATE_PATH"] = "servicex/templates/transform_single_file.py"
+    os.environ["CAPABILITIES_PATH"] = "transformer_capabilities.json"
+
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        translator = TopCPTranslator()
+        query = (
+            '{"reco": "CommonServices:\\n  systematicsHistogram: \'listOfSystematics\'\\n\\n'
+            "PileupReweighting: {}\\n\\nEventCleaning:\\n    runEventCleaning: False\\n"
+            "    runGRL: False\\n\\nElectrons:\\n  - containerName: 'AnaElectrons'\\n"
+            "    crackVeto: True\\n    IFFClassification: {}\\n    WorkingPoint:\\n"
+            "      - selectionName: 'loose'\\n        identificationWP: 'TightLH'\\n"
+            "        isolationWP: 'NonIso'\\n        noEffSF: True\\n"
+            "      - selectionName: 'tight'\\n        identificationWP: 'TightLH'\\n"
+            "        isolationWP: 'Tight_VarRad'\\n        noEffSF: True\\n"
+            "    PtEtaSelection:\\n        minPt: 25000.0\\n        maxEta: 2.47\\n"
+            "        useClusterEta: True\\n\\n"
+            "# After configuring each container, many variables will be saved automatically.\\n"
+            "Output:\\n  treeName: 'reco'\\n  vars: []\\n  metVars: []\\n  containers:\\n"
+            "      # Format should follow: '<suffix>:<output container>'\\n"
+            "      el_: 'AnaElectrons'\\n      '': 'EventInfo'\\n  commands:\\n"
+            "    # Turn output branches on and off with 'enable' and 'disable'\\n\\n"
+            'AddConfigBlocks: []\\n", "parton": null, "particle": null, "max_events": 100, '
+            '"no_systematics": true}'
+        )
+
+        with pytest.raises(ValueError):
+            translator.generate_code(query, tmpdirname)
+
+
+def test_generate_code_with_custom_image():
+    """Test that custom image is properly set in the result"""
+    os.environ["TEMPLATE_PATH"] = "servicex/templates/transform_single_file.py"
+    os.environ["CAPABILITIES_PATH"] = "transformer_capabilities.json"
+
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        translator = TopCPTranslator()
+        query = (
+            '{"reco": "CommonServices:\\n  systematicsHistogram: \'listOfSystematics\'\\n\\n'
+            "PileupReweighting: {}\\n\\nEventCleaning:\\n    runEventCleaning: False\\n"
+            "    runGRL: False\\n\\nElectrons:\\n  - containerName: 'AnaElectrons'\\n"
+            "    crackVeto: True\\n    IFFClassification: {}\\n    WorkingPoint:\\n"
+            "      - selectionName: 'loose'\\n        identificationWP: 'TightLH'\\n"
+            "        isolationWP: 'NonIso'\\n        noEffSF: True\\n"
+            "      - selectionName: 'tight'\\n        identificationWP: 'TightLH'\\n"
+            "        isolationWP: 'Tight_VarRad'\\n        noEffSF: True\\n"
+            "    PtEtaSelection:\\n        minPt: 25000.0\\n        maxEta: 2.47\\n"
+            "        useClusterEta: True\\n\\n"
+            "# After configuring each container, many variables will be saved automatically.\\n"
+            "Output:\\n  treeName: 'reco'\\n  vars: []\\n  metVars: []\\n  containers:\\n"
+            "      # Format should follow: '<suffix>:<output container>'\\n"
+            "      el_: 'AnaElectrons'\\n      '': 'EventInfo'\\n  commands:\\n"
+            "    # Turn output branches on and off with 'enable' and 'disable'\\n\\n"
+            'AddConfigBlocks: []\\n", "parton": null, "particle": null, "max_events": 100, '
+            '"no_systematics": true, "no_filter": false, '
+            '"image": "custom/topcp:v1.0"}'
+        )
+
+        result = translator.generate_code(query, tmpdirname)
+        assert result.image == "custom/topcp:v1.0"
+
+
 def test_app():
     import servicex.TopCP_code_generator