support custom TopCP images

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,
     },
+    "docker_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,14 +26,34 @@
 # 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,
     GeneratedFileResult,
     GenerateCodeException,
 )
 
+def validate_custom_docker_image(image_name: str) -> bool:
+    allowed_images_json = os.environ.get("TOPCP_ALLOWED_IMAGES")
+
+    if not allowed_images_json:
+        raise GenerateCodeException("Custom Docker images are not allowed.")
+
+    try:
+        allowed_prefixes = json.loads(allowed_images_json)
+        if not isinstance(allowed_prefixes, list):
+            raise GenerateCodeException("TopCP allowed images are improperly configured.")
+        for prefix in allowed_prefixes:
+            if image_name.startswith(prefix):
+                return True
+
+        raise GenerateCodeException(f"Custom Docker image '{image_name}' not allowed.")
+
+    except json.JSONDecodeError:
+        raise GenerateCodeException("TopCP allowed images are improperly configured.")
 
 class TopCPTranslator(CodeGenerator):
     # Generate the code. Ignoring caching for now
@@ -62,11 +82,22 @@ 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 "docker_image" in jquery:
+            docker_image = jquery["docker_image"]
+
+            validate_custom_docker_image(docker_image)
+
+            results.image = docker_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()

docs/deployment/reference.md

@@ -25,8 +25,8 @@ parameters for the [rabbitMQ](https://github.com/bitnami/charts/tree/master/bitn
 | `app.replicas`                                | Number of App pods to start. Experimental!                                                                                                                                                                       | 1                                              |
 | `app.auth`                                    | Enable authentication or allow unfettered access (Python boolean string)                                                                                                                                         | `false`                                        |
 | `app.oauthMetadataURL`                        | OpenID Connect provider well known configuration endpoint URL                                                                                                                                                    | -                                              |
-| `app.oauthClientID`                          | OpenID Connect application Client ID                                                                                                                                                                             | -                                              |
-| `app.oauthClientSecret`                      | OpenID Connect application Client Secret                                                                                                                                                                         | -                                              |
+| `app.oauthClientID`                           | OpenID Connect application Client ID                                                                                                                                                                             | -                                              |
+| `app.oauthClientSecret`                       | OpenID Connect application Client Secret                                                                                                                                                                         | -                                              |
 | `app.adminEmail`                              | Email address for initial admin user                                                                                                                                                                             | <admin@example.com>                            |
 | `app.tokenExpires`                            | Seconds until the ServiceX API tokens (JWT refresh tokens) expire                                                                                                                                                | False (never)                                  |
 | `app.authExpires`                             | Seconds until the JWT access tokens expire                                                                                                                                                                       | 21600 (six hours)                              |
@@ -75,6 +75,13 @@ parameters for the [rabbitMQ](https://github.com/bitnami/charts/tree/master/bitn
 | `codegen.atlasr21.enabled`                    | Deploy the ATLAS FuncADL Release 21 code generator? - also all of the code gen settings above are available                                                                                                      | true                                           |
 | `codegen.atlasr22.enabled`                    | Deploy the ATLAS FuncADL Release 22 code generator? - also all of the code gen settings above are available                                                                                                      | true                                           |
 | `codegen.python.enabled`                      | Deploy the python uproot code generator? - also all of the code gen settings, above are available                                                                                                                | true                                           |
+| `codegen.topcp.enabled`                       | Deploy the TopCP code generator?                                                                                                                                                                                 | true                                           |
+| `codegen.topcp.image`                         | TopCP ode generator image                                                                                                                                                                                        | `sslhep/servicex_code_gen_topcp`               |
+| `codegen.topcp.pullPolicy`                    | TopCP code generator image pull policy                                                                                                                                                                           | Always                                         |
+| `codegen.topcp.tag`                           | TopCP code generator image tag                                                                                                                                                                                   | develop                                        |
+| `codegen.topcp.defaultScienceContainerImage`  | The default image used by a TopCP transformer container                                                                                                                                                          | sslhep/servicex_science_image_topcp            |
+| `codegen.topcp.defaultScienceContainerTag`    | The default tag used by a TopCP transformer container                                                                                                                                                            | 2.17.0-25.2.45                                 |
+| `codegen.topcp.allowedImages`                 | A list of strings, of which one must be a valid prefix of a submitted custom TopCP docker image                                                                                                                  | - "sslhep/servicex_science_image_topcp:"       |
 | `x509Secrets.image`                           | X509 Secret Service image name                                                                                                                                                                                   | `sslhep/x509-secrets`                          |
 | `x509Secrets.tag`                             | X509 Secret Service image tag                                                                                                                                                                                    | `latest`                                       |
 | `x509Secrets.pullPolicy`                      | X509 Secret Service image pull policy                                                                                                                                                                            | `Always`                                       |
@@ -94,9 +101,9 @@ parameters for the [rabbitMQ](https://github.com/bitnami/charts/tree/master/bitn
 | `minio.apiIngress.enabled`                    | Should minio chart deploy an ingress to the service?                                                                                                                                                             | false                                          |
 | `minio.apiIngress.hostname`                   | Hostname associate with ingress controller                                                                                                                                                                       | nil                                            |
 | `transformer.cachePrefix`                     | Prefix string to stick in front of file paths. Useful for XCache. If `transformer.cacheVPSSite` is also set, this will be ignored                                                                                | nil                                            |
-| `transformer.cacheVPSSite`                     | Specify a Virtual Placement Service site whose XCaches we should use. Will update automatically if list changes. If set, takes priority over `transformer.cachePrefix`                                           | nil                                            |
-| `transformer.cacheVPSCheckInterval`                     | How frequency should the Virtual Placement Service be consulted for the list of XCaches (in seconds)                                                                                                             | 1800                                           |
-| `transformer.cacheVPSLivenessURL`                     | URL from which Virtual Placement Service site info can be obtained                                                                                                                                               | https://vps.cern.ch/liveness                   |
+| `transformer.cacheVPSSite`                    | Specify a Virtual Placement Service site whose XCaches we should use. Will update automatically if list changes. If set, takes priority over `transformer.cachePrefix`                                           | nil                                            |
+| `transformer.cacheVPSCheckInterval`           | How frequency should the Virtual Placement Service be consulted for the list of XCaches (in seconds)                                                                                                             | 1800                                           |
+| `transformer.cacheVPSLivenessURL`             | URL from which Virtual Placement Service site info can be obtained                                                                                                                                               | https://vps.cern.ch/liveness                   |
 | `transformer.autoscaler.enabled`              | Enable/disable horizontal pod autoscaler for transformers                                                                                                                                                        | True                                           |
 | `transformer.autoscaler.cpuScaleThreshold`    | CPU percentage threshold for pod scaling                                                                                                                                                                         | 30                                             |
 | `transformer.autoscaler.minReplicas`          | Minimum number of transformer pods per request                                                                                                                                                                   | 1                                              |

helm/servicex/templates/app/configmap.yaml

@@ -183,6 +183,11 @@ data:
     CODE_GEN_IMAGES = { {{ join "," $code_gen_images }} }
     {{- end }}
 
+    {{- if .Values.codeGen.topcp.enabled }}
+    # TopCP custom image configuration
+    TOPCP_ALLOWED_IMAGES = {{ .Values.codeGen.topcp.allowedImages | toJson | quote }}
+    {{- end }}
+
     {{- $didFinders := list }}
     {{- if .Values.didFinder.CERNOpenData.enabled }}
       {{- $didFinders = append $didFinders "cernopendata" }}

helm/servicex/templates/codegen/deployment.yaml

@@ -22,6 +22,10 @@ spec:
         env:
         - name: INSTANCE_NAME
           value: {{ $.Release.Name }}
+        {{- if eq $codeGenName "topcp" }}
+        - name: TOPCP_ALLOWED_IMAGES
+          value: {{ $.Values.codeGen.topcp.allowedImages | toJson | quote }}
+        {{- end }}
         - name: TRANSFORMER_SCIENCE_IMAGE
           value: {{ .defaultScienceContainerImage }}:{{ .defaultScienceContainerTag }}
         {{- $compressionAlgorithm := .compressionAlgorithm | default "ZSTD" }}

helm/servicex/values.yaml

@@ -108,6 +108,8 @@ codeGen:
     tag: develop
     defaultScienceContainerImage: sslhep/servicex_science_image_topcp
     defaultScienceContainerTag: 2.17.0-25.2.45
+    allowedImages:
+      - "sslhep/servicex_science_image_topcp:"
 
 didFinder:
   CERNOpenData:
@@ -175,7 +177,6 @@ postgresql:
     image:
       repository: bitnamilegacy/os-shell
       tag: 12-debian-12-r51
-
   global:
     postgresql:
       auth:
@@ -200,7 +201,6 @@ rabbitmq:
     image:
       repository: bitnamilegacy/os-shell
       tag: 12-debian-12-r51
-
   extraConfiguration: |-
     consumer_timeout = 3600000
 

servicex_app/servicex_app/code_gen_adapter.py

@@ -25,6 +25,8 @@
 # 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 Optional
+
 import requests
 from requests_toolbelt.multipart import decoder
 
@@ -43,7 +45,11 @@ def post_request(self, post_url, post_obj):
         return result
 
     def generate_code_for_selection(
-        self, request_record: TransformRequest, namespace: str, user_codegen_name: str
+        self,
+        request_record: TransformRequest,
+        namespace: str,
+        user_codegen_name: str,
+        custom_image: Optional[str] = None,
     ) -> tuple[str, str, str, str]:
         """
         Generates the C++ code for a request's selection string.