feat: automatic job grouping (#95)

* feat: add input_data to TransformationSubmissionModel, updated CLI

* chore: add test workflows for transformation grouping

* feat: added Job Grouping

* chore: rename

* chore: added job grouping test

* fix: updated transformation hoooks group_size

* refactor: rename

* refactor: updated tests

* refactor: updated group_size definition and values

* fix: fixed JobWrapper tests where task.cwl was not deleted

* refactor: removed input_name from get_input_query

* fix: file location

* feat: add input_data to ProductionSubmissionModel

* fix: fixed tests and jobwrapper imports

* fix: fixed imports

* fix: removed pytest warnings

* fix: fixed test

* fix: fixed mypy error (?)

* fix: removed fixture and fixed imports

* fix: updated fixture and tests

* fix: added input_data to ExecHints instead of SubModels

* docs: updated README.md

* feat: chunk CLI parameter

---------

Co-authored-by: aldbr <aldbr@outlook.com>

Author: Stellatsuu

README.md

@@ -70,17 +70,17 @@ Inside the Pixi environment:
 pixi shell
 
 # Submit
-dirac-cwl job submit <workflow_path> [--parameter-path <input_path>] [--metadata-path <metadata_path>]
+dirac-cwl job submit <workflow_path> [--input-files <input_path>]
 
-dirac-cwl transformation submit <workflow_path> [--metadata-path <metadata_path>]
+dirac-cwl transformation submit <workflow_path> [--inputs-file <input_path>]
 
-dirac-cwl production submit <workflow_path> [--steps-metadata-path <steps_metadata_path>]
+dirac-cwl production submit <workflow_path> [--inputs-file <input_path>]
 ```
 
 Or prefix individual commands:
 
 ```bash
-pixi run dirac-cwl job submit <workflow_path> --parameter-path <input_path>
+pixi run dirac-cwl job submit <workflow_path> --input-files <input_path>
 ```
 
 Common tasks are defined in `pyproject.toml` and can be run with Pixi:

pixi.lock

@@ -22,18 +22,19 @@
 import json
 import logging
 from pathlib import Path
-from typing import Any, Dict, List, cast
+from typing import Any, Dict, List, Type, cast
 
 import yaml
+from pydantic import BaseModel
 
 # Configure logging
 logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
 logger = logging.getLogger(__name__)
 
 
-def collect_pydantic_models() -> Dict[str, Any]:
+def collect_pydantic_models() -> Dict[str, Type[BaseModel]]:
     """Collect all Pydantic models from the metadata system."""
-    models = {}
+    models: Dict[str, Type[BaseModel]] = {}
 
     # Import core models
     try:
@@ -44,14 +45,10 @@ def collect_pydantic_models() -> Dict[str, Any]:
             TransformationExecutionHooksHint,
         )
 
-        models.update(
-            {
-                "ExecutionHooksBasePlugin": ExecutionHooksBasePlugin,
-                "ExecutionHooks": ExecutionHooksHint,
-                "Scheduling": SchedulingHint,
-                "TransformationExecutionHooks": TransformationExecutionHooksHint,
-            }
-        )
+        models["ExecutionHooksBasePlugin"] = ExecutionHooksBasePlugin
+        models["ExecutionHooks"] = ExecutionHooksHint
+        models["Scheduling"] = SchedulingHint
+        models["TransformationExecutionHooks"] = TransformationExecutionHooksHint
         logger.info("Collected core metadata models")
     except ImportError as e:
         logger.error("Failed to import core models: %s", e)
@@ -65,14 +62,10 @@ def collect_pydantic_models() -> Dict[str, Any]:
             TransformationSubmissionModel,
         )
 
-        models.update(
-            {
-                "JobInputModel": JobInputModel,
-                "JobSubmissionModel": JobSubmissionModel,
-                "TransformationSubmissionModel": TransformationSubmissionModel,
-                "ProductionSubmissionModel": ProductionSubmissionModel,
-            }
-        )
+        models["JobInputModel"] = JobInputModel
+        models["JobSubmissionModel"] = JobSubmissionModel
+        models["TransformationSubmissionModel"] = TransformationSubmissionModel
+        models["ProductionSubmissionModel"] = ProductionSubmissionModel
         logger.info("Collected submission models")
     except ImportError as e:
         logger.error("Failed to import submission models: %s", e)

scripts/generate_schemas.py

@@ -23,13 +23,14 @@
     Union,
 )
 
+from cwl_utils.parser import File, save
 from DIRAC.DataManagementSystem.Client.DataManager import (  # type: ignore[import-untyped]
     DataManager,
 )
 from DIRACCommon.Core.Utilities.ReturnValues import (  # type: ignore[import-untyped]
     returnSingleResult,
 )
-from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, field_validator
 
 from dirac_cwl.commands import PostProcessCommand, PreProcessCommand
 from dirac_cwl.mocks.data_manager import MockDataManager
@@ -139,12 +140,12 @@ async def store_output(
                     if res and not res["OK"]:
                         raise RuntimeError(f"Could not save file {src} with LFN {str(lfn)} : {res['Message']}")
 
-    def get_input_query(self, input_name: str, **kwargs: Any) -> Union[Path, List[Path], None]:
+    def get_input_query(self, **kwargs: Any) -> Union[Path, List[Path], None]:
         """Generate LFN-based input query path.
 
         Accepts and ignores extra kwargs for interface compatibility.
         """
-        # Build LFN: /query_root/vo/campaign/site/data_type/input_name
+        # Build LFN: /query_root/vo/campaign/site/data_type
         pass
 
     @classmethod
@@ -332,10 +333,40 @@ def from_cwl(cls, cwl_object: Any) -> Self:
                 descriptor = descriptor.model_copy(update=hint_data)
         return descriptor
 
+    @classmethod
+    def update_cwl(cls, cwl_object: Any, descriptor: Self) -> None:
+        """Update CWL object with metadata descriptor."""
+        hints = getattr(cwl_object, "hints", []) or []
+        for hint in hints:
+            if hint.get("class") == "dirac:ExecutionHooks":
+                hint.update(descriptor.model_dump())
+                return
+
+        # Create the ExecutionHooks hint if it doesn't exist
+        hints.append({"class": "dirac:ExecutionHooks", **descriptor.model_dump()})
+        cwl_object.hints = hints
+
 
 class TransformationExecutionHooksHint(ExecutionHooksHint):
     """Extended data manager for transformations."""
 
-    group_size: Optional[Dict[str, int]] = Field(
-        default=None, description="Input grouping configuration for transformation jobs"
+    group_size: Optional[int] = Field(default=None, description="Number of input files per job")
+    input_data: Optional[Dict[str, List[str]]] = Field(
+        default=None, description="Static input file lists, keyed by CWL input parameter name"
     )
+    input_query: Optional[Dict] = Field(default=None, description="Dynamic input query for transformation jobs")
+
+    @field_validator("input_data", mode="before")
+    @classmethod
+    def convert_input_data(cls, value):
+        """Convert an input data dict containing a list of Files to a list of strings.
+
+        :param value: Input data dict to convert.
+        :return: Converted input data dict.
+        """
+        if value:
+            return {
+                key: [save(item)["location"] if isinstance(item, File) else item for item in items]
+                for key, items in value.items()
+            }
+        return None

src/dirac_cwl/execution_hooks/core.py

@@ -34,12 +34,12 @@ class QueryBasedPlugin(ExecutionHooksBasePlugin):
     campaign: Optional[str] = Field(default=None, description="Campaign name for LFN path")
     data_type: Optional[str] = Field(default=None, description="Data type classification")
 
-    def get_input_query(self, input_name: str, **kwargs: Any) -> Union[Path, List[Path], None]:
+    def get_input_query(self, **kwargs: Any) -> Union[Path, List[Path], None]:
         """Generate LFN-based input query path.
 
         Accepts and ignores extra kwargs for interface compatibility.
         """
-        # Build LFN: /query_root/vo/campaign/site/data_type/input_name
+        # Build LFN: /query_root/vo/campaign/site/data_type
         path_parts = []
 
         if self.vo:
@@ -53,6 +53,6 @@ def get_input_query(self, input_name: str, **kwargs: Any) -> Union[Path, List[Pa
             path_parts.append(self.data_type)
 
         if len(path_parts) > 0:  # More than just VO
-            return Path(self.query_root) / Path(*path_parts) / Path(input_name)
+            return Path(self.query_root) / Path(*path_parts)
 
-        return Path(self.query_root) / Path(input_name)
+        return Path(self.query_root)

src/dirac_cwl/execution_hooks/plugins/core.py

@@ -13,7 +13,7 @@
 from cwl_utils.parser.cwl_v1_2 import (
     File,
 )
-from cwl_utils.parser.cwl_v1_2_utils import load_inputfile
+from cwl_utils.parser.utils import load_inputfile
 from diracx.cli.utils import AsyncTyper
 from rich import print_json
 from rich.console import Console
@@ -40,7 +40,7 @@
 @app.async_command("submit")
 async def submit_job_client(
     task_path: str = typer.Argument(..., help="Path to the CWL file"),
-    parameter_path: list[str] | None = typer.Option(None, help="Path to the files containing the metadata"),
+    input_files: list[str] | None = typer.Option(None, help="Paths to the CWL input files"),
     # Specific parameter for the purpose of the prototype
     local: bool | None = typer.Option(True, help="Run the job locally instead of submitting it to the router"),
 ):
@@ -70,35 +70,35 @@ async def submit_job_client(
     console.print(f"\t[green]:heavy_check_mark:[/green] Task {task_path}")
     console.print("\t[green]:heavy_check_mark:[/green] Hints")
 
-    # Extract parameters if any
-    parameters = []
-    if parameter_path:
-        for parameter_p in parameter_path:
+    # Extract inputs if any
+    inputs = []
+    if input_files:
+        for file in input_files:
             try:
-                parameter = load_inputfile(parameter_p)
+                input_file = load_inputfile(task.cwlVersion, file)
             except Exception as ex:
                 console.print(
-                    f"[red]:heavy_multiplication_x:[/red] [bold]CLI:[/bold] Failed to validate the parameter:\n{ex}"
+                    f"[red]:heavy_multiplication_x:[/red] [bold]CLI:[/bold] Failed to validate the input file:\n{ex}"
                 )
                 return typer.Exit(code=1)
 
             # Prepare files for the ISB
-            isb_file_paths = prepare_input_sandbox(parameter)
+            isb_file_paths = prepare_input_sandbox(input_file)
 
-            # Upload parameter sandbox
+            # Upload input file sandbox
             sandbox_id = await submission_client.create_sandbox(isb_file_paths)
 
-            parameters.append(
+            inputs.append(
                 JobInputModel(
                     sandbox=[sandbox_id] if sandbox_id else None,
-                    cwl=parameter,
+                    cwl=input_file,
                 )
             )
-            console.print(f"\t[green]:heavy_check_mark:[/green] Parameter {parameter_p}")
+            console.print(f"\t[green]:heavy_check_mark:[/green] File {file}")
 
     job = JobSubmissionModel(
         task=task,
-        inputs=parameters,
+        inputs=inputs,
     )
     console.print("[green]:heavy_check_mark:[/green] [bold]CLI:[/bold] Job(s) validated.")
 

src/dirac_cwl/job/__init__.py

@@ -10,7 +10,7 @@
 
 import DIRAC  # type: ignore[import-untyped]
 from cwl_utils.parser import load_document_by_uri
-from cwl_utils.parser.cwl_v1_2_utils import load_inputfile
+from cwl_utils.parser.utils import load_inputfile
 from ruamel.yaml import YAML
 
 if os.getenv("DIRAC_PROTO_LOCAL") != "1":
@@ -41,7 +41,7 @@ async def main():
         task_obj = load_document_by_uri(f.name)
 
     if job_model_dict["input"]:
-        cwl_inputs_obj = load_inputfile(job_model_dict["input"]["cwl"])
+        cwl_inputs_obj = load_inputfile(task_obj.cwlVersion, job_model_dict["input"]["cwl"])
         job_model_dict["input"]["cwl"] = cwl_inputs_obj
     job_model_dict["task"] = task_obj
 

src/dirac_cwl/job/job_wrapper_template.py

@@ -20,6 +20,7 @@
     WorkflowInputParameter,
     WorkflowStep,
 )
+from cwl_utils.parser.utils import load_inputfile
 from rich import print_json
 from rich.console import Console
 from schema_salad.exceptions import ValidationException
@@ -74,6 +75,8 @@
 @app.command("submit")
 def submit_production_client(
     task_path: str = typer.Argument(..., help="Path to the CWL file"),
+    inputs_file: str | None = typer.Option(None, help="Path to the CWL inputs file"),
+    chunk: str | None = typer.Option(None, help="Split an array input into jobs: PARAM=SIZE (e.g., input-data=3)"),
     # Specific parameter for the purpose of the prototype
     local: Optional[bool] = typer.Option(True, help="Run the job locally instead of submitting it to the router"),
 ):
@@ -86,10 +89,49 @@ def submit_production_client(
     """
     os.environ["DIRAC_PROTO_LOCAL"] = "0"
 
+    # --chunk and --inputs-file must be used together
+    if chunk and not inputs_file:
+        console.print("[red]:heavy_multiplication_x:[/red] [bold]CLI:[/bold] --chunk requires --inputs-file.")
+        return typer.Exit(code=1)
+    if inputs_file and not chunk:
+        console.print("[red]:heavy_multiplication_x:[/red] [bold]CLI:[/bold] --inputs-file requires --chunk.")
+        return typer.Exit(code=1)
+
     # Validate the workflow
     console.print("[blue]:information_source:[/blue] [bold]CLI:[/bold] Validating the production...")
     try:
         task = load_document(pack(task_path))
+
+        # Load Production inputs and inject into the first step's hint
+        if inputs_file and chunk:
+            from dirac_cwl.transformation import _parse_chunk
+
+            all_inputs = load_inputfile(task.cwlVersion, inputs_file)
+            chunk_param, chunk_size = _parse_chunk(chunk)
+
+            if chunk_param not in all_inputs:
+                console.print(
+                    f"[red]:heavy_multiplication_x:[/red] [bold]CLI:[/bold] "
+                    f"Parameter '{chunk_param}' not found in inputs file. "
+                    f"Available parameters: {list(all_inputs.keys())}"
+                )
+                return typer.Exit(code=1)
+            if not isinstance(all_inputs[chunk_param], list):
+                console.print(
+                    f"[red]:heavy_multiplication_x:[/red] [bold]CLI:[/bold] "
+                    f"Parameter '{chunk_param}' must be an array type for --chunk."
+                )
+                return typer.Exit(code=1)
+
+            input_data = {chunk_param: all_inputs[chunk_param]}
+
+            # Inject into the first step's hint
+            if task.steps:
+                from dirac_cwl.execution_hooks import TransformationExecutionHooksHint
+
+                hint_update = TransformationExecutionHooksHint(group_size=chunk_size, input_data=input_data)
+                TransformationExecutionHooksHint.update_cwl(task.steps[0].run, hint_update)
+
     except FileNotFoundError as ex:
         console.print(f"[red]:heavy_multiplication_x:[/red] [bold]CLI:[/bold] Failed to load the task:\n{ex}")
         return typer.Exit(code=1)
@@ -99,11 +141,10 @@ def submit_production_client(
     console.print(f"\t[green]:heavy_check_mark:[/green] Task {task_path}")
     console.print("\t[green]:heavy_check_mark:[/green] Metadata")
 
-    # Create the production
     production = ProductionSubmissionModel(task=task)
     console.print("[green]:heavy_check_mark:[/green] [bold]CLI:[/bold] Production validated.")
 
-    # Submit the tranaformation
+    # Submit the transformation
     console.print("[blue]:information_source:[/blue] [bold]CLI:[/bold] Submitting the production...")
     print_json(production.model_dump_json(indent=4))
     if not submit_production_router(production):
@@ -161,12 +202,7 @@ def _get_transformations(
 
     for step in production.task.steps:
         step_task = _create_subworkflow(step, str(production.task.cwlVersion), production.task.inputs)
-
-        transformations.append(
-            TransformationSubmissionModel(
-                task=step_task,
-            )
-        )
+        transformations.append(TransformationSubmissionModel(task=step_task))
     return transformations
 
 

src/dirac_cwl/production/__init__.py

@@ -102,8 +102,6 @@ class JobModel(BaseJobModel):
 # -----------------------------------------------------------------------------
 # Transformation models
 # -----------------------------------------------------------------------------
-
-
 class TransformationSubmissionModel(BaseModel):
     """Transformation definition sent to the router."""