feat: chunk CLI parameter

aldbr · aldbr · commit c94c15867d65 · 2026-04-08T17:59:36.000+02:00
diff --git a/src/dirac_cwl/execution_hooks/core.py b/src/dirac_cwl/execution_hooks/core.py
@@ -350,9 +350,11 @@ def update_cwl(cls, cwl_object: Any, descriptor: Self) -> None:
 class TransformationExecutionHooksHint(ExecutionHooksHint):
     """Extended data manager for transformations."""
 
-    group_size: Optional[int] = Field(default=None, description="Input grouping configuration for transformation jobs")
-    input_data: Optional[Dict[str, List[str]]] = Field(default=None, description="Input data for transformation jobs")
-    input_query: Optional[Dict] = Field(default=None, description="Input query 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
diff --git a/src/dirac_cwl/production/__init__.py b/src/dirac_cwl/production/__init__.py
@@ -76,6 +76,7 @@
 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"),
 ):
@@ -88,15 +89,48 @@ 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 if existing (1st Transformation inputs)
-        input_data = None
-        if inputs_file:
-            input_data = load_inputfile(task.cwlVersion, inputs_file).get("input-data")
+        # 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}")
@@ -107,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, input_data=input_data)
+    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):
@@ -167,12 +200,9 @@ def _get_transformations(
     # Create a subworkflow and a transformation for each step
     transformations = []
 
-    for index, step in enumerate(production.task.steps):
+    for step in production.task.steps:
         step_task = _create_subworkflow(step, str(production.task.cwlVersion), production.task.inputs)
-
-        transformations.append(
-            TransformationSubmissionModel(task=step_task, input_data=production.input_data if index == 0 else None)
-        )
+        transformations.append(TransformationSubmissionModel(task=step_task))
     return transformations
 
 
diff --git a/src/dirac_cwl/transformation/__init__.py b/src/dirac_cwl/transformation/__init__.py
@@ -9,8 +9,8 @@
 
 import typer
 from cwl_utils.pack import pack
-from cwl_utils.parser import load_document, save
-from cwl_utils.parser.cwl_v1_2 import File
+from cwl_utils.parser import File, load_document, save
+from cwl_utils.parser.cwl_v1_2 import File as CWLFile
 from cwl_utils.parser.utils import load_inputfile
 from rich import print_json
 from rich.console import Console
@@ -35,10 +35,32 @@
 # -----------------------------------------------------------------------------
 
 
+def _parse_chunk(chunk_str: str) -> tuple[str, int]:
+    """Parse a --chunk value of the form PARAM=SIZE.
+
+    :param chunk_str: The chunk string to parse.
+    :return: Tuple of (parameter name, chunk size).
+    :raises typer.BadParameter: If the format is invalid or size is not a positive integer.
+    """
+    if "=" not in chunk_str:
+        raise typer.BadParameter(f"Invalid --chunk format '{chunk_str}'. Expected PARAM=SIZE (e.g., input-data=3).")
+    param, size_str = chunk_str.split("=", 1)
+    if not param:
+        raise typer.BadParameter("Parameter name cannot be empty in --chunk.")
+    try:
+        size = int(size_str)
+    except ValueError as err:
+        raise typer.BadParameter(f"Chunk size must be an integer, got '{size_str}'.") from err
+    if size <= 0:
+        raise typer.BadParameter(f"Chunk size must be > 0, got {size}.")
+    return param, size
+
+
 @app.command("submit")
 def submit_transformation_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 jobs locally instead of submitting them to the router"),
 ):
@@ -50,15 +72,55 @@ def submit_transformation_client(
     - Start the transformation
     """
     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 transformation...")
     try:
         task = load_document(pack(task_path))
 
-        # Load Transformation inputs if existing
-        input_data = None
-        if inputs_file:
-            input_data = load_inputfile(task.cwlVersion, inputs_file).get("input-data")
+        # Warn if the hint already has input_data — CLI overrides it
+        existing_hint = TransformationExecutionHooksHint.from_cwl(task)
+        if inputs_file and existing_hint.input_data:
+            console.print(
+                "[yellow]:warning:[/yellow] [bold]CLI:[/bold] "
+                "The workflow hint already contains input_data. "
+                "Overriding with --inputs-file/--chunk values."
+            )
+
+        # Load and validate inputs, inject into hint
+        if inputs_file and chunk:
+            all_inputs = load_inputfile(task.cwlVersion, inputs_file)
+            chunk_param, chunk_size = _parse_chunk(chunk)
+
+            # Validate the chunk parameter exists and is a list
+            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, "
+                    f"got {type(all_inputs[chunk_param]).__name__}."
+                )
+                return typer.Exit(code=1)
+
+            input_data = {chunk_param: all_inputs[chunk_param]}
+
+            # Inject input_data and group_size into the task's ExecutionHooks hint
+            hint_update = TransformationExecutionHooksHint(group_size=chunk_size, input_data=input_data)
+            TransformationExecutionHooksHint.update_cwl(task, hint_update)
 
     except FileNotFoundError as ex:
         console.print(f"[red]:heavy_multiplication_x:[/red] [bold]CLI:[/bold] Failed to load the task:\n{ex}")
@@ -68,7 +130,7 @@ def submit_transformation_client(
         return typer.Exit(code=1)
     console.print(f"\t[green]:heavy_check_mark:[/green] Task {task_path}")
 
-    transformation = TransformationSubmissionModel(task=task, input_data=input_data)
+    transformation = TransformationSubmissionModel(task=task)
     console.print("[green]:heavy_check_mark:[/green] [bold]CLI:[/bold] Transformation validated.")
 
     # Submit the transformation
@@ -111,26 +173,40 @@ def submit_transformation_router(transformation: TransformationSubmissionModel)
     except Exception as exc:
         raise ValueError(f"Invalid DIRAC hints:\n{exc}") from exc
 
-    # Inputs from Transformation inputs_file
-    if transformation.input_data:
-        nb_files = len(transformation.input_data)
-        group_size = transformation_execution_hooks.group_size or 1
-        nb_groups = (nb_files + group_size - 1) // group_size
-        logger.info("Creating %s jobs for the transformation with group size of %s...", nb_groups, group_size)
-
-        input_data = transformation.input_data
-        for i, start in enumerate(range(0, len(input_data), group_size)):
-            files_chunk = input_data[start : start + group_size]
+    # Inputs from static input_data (populated by --chunk on the client)
+    if transformation_execution_hooks.input_data:
+        if transformation_execution_hooks.configuration:
+            raise ValueError(
+                "Cannot specify both static input_data and dynamic input query (configuration). "
+                "Use --chunk/--inputs-file for standalone transformations with known files. "
+                "Use configuration for transformations that discover inputs from upstream outputs."
+            )
 
+        group_size = transformation_execution_hooks.group_size or 1
+        for param_name, file_list in transformation_execution_hooks.input_data.items():
+            nb_files = len(file_list)
+            nb_groups = (nb_files + group_size - 1) // group_size
             logger.info(
-                "Group %i files: %s",
-                i + 1,
-                [save(file).get("path") if isinstance(file, File) else file for file in files_chunk],
+                "Chunking '%s': %s files into %s jobs (group_size=%s)",
+                param_name,
+                nb_files,
+                nb_groups,
+                group_size,
             )
-            job_model_params.append(JobInputModel(sandbox=None, cwl={"input-data": files_chunk}))
 
-    # Inputs from DataCatalog/Bookkeeping service
-    if transformation_execution_hooks.configuration and transformation_execution_hooks.group_size:
+            for i, start in enumerate(range(0, nb_files, group_size)):
+                files_chunk = file_list[start : start + group_size]
+                logger.info(
+                    "Group %i files: %s",
+                    i + 1,
+                    [save(file).get("path") if isinstance(file, File) else file for file in files_chunk],
+                )
+                job_model_params.append(JobInputModel(sandbox=None, cwl={param_name: files_chunk}))
+
+    # Inputs from DataCatalog/Bookkeeping service (dynamic query)
+    elif transformation_execution_hooks.configuration:
+        group_size = transformation_execution_hooks.group_size or 1
+
         # Get the metadata class
         transformation_metadata = transformation_execution_hooks.to_runtime(transformation)
 
@@ -143,9 +219,9 @@ def submit_transformation_router(transformation: TransformationSubmissionModel)
         # Wait for the input to be available
         logger.info("\t- Waiting for input data...")
         logger.debug("\t\t- Query: %s", input_query)
-        logger.debug("\t\t- Group Size: %s", transformation_execution_hooks.group_size)
+        logger.debug("\t\t- Group Size: %s", group_size)
 
-        while not (inputs := _get_inputs(input_query, transformation_execution_hooks.group_size)):
+        while not (inputs := _get_inputs(input_query, group_size)):
             logger.debug("\t\t- Result: %s", inputs)
             time.sleep(5)
 
@@ -215,7 +291,7 @@ def _generate_job_model_parameter(
     for group in grouped_input_data:
         cwl_inputs = {}
         for input_name, input_data in group.items():
-            cwl_inputs[input_name] = [File(location=str(Path("lfn:") / path)) for path in input_data]
+            cwl_inputs[input_name] = [CWLFile(location=str(Path("lfn:") / path)) for path in input_data]
 
         job_model_params.append(JobInputModel(sandbox=None, cwl=cwl_inputs))
 
diff --git a/test/schemas/dirac-metadata.json b/test/schemas/dirac-metadata.json
@@ -318,7 +318,7 @@
             }
           ],
           "default": null,
-          "description": "Input grouping configuration for transformation jobs",
+          "description": "Number of input files per job",
           "title": "Group Size"
         },
         "hook_plugin": {
@@ -343,7 +343,7 @@
             }
           ],
           "default": null,
-          "description": "Input data for transformation jobs",
+          "description": "Static input file lists, keyed by CWL input parameter name",
           "title": "Input Data"
         },
         "input_query": {
@@ -357,7 +357,7 @@
             }
           ],
           "default": null,
-          "description": "Input query for transformation jobs",
+          "description": "Dynamic input query for transformation jobs",
           "title": "Input Query"
         },
         "output_paths": {
diff --git a/test/test_workflows.py b/test/test_workflows.py
@@ -507,27 +507,32 @@ def test_run_transformation_validation_failure(cli_runner, cwl_file, cleanup, ex
 
 
 @pytest.mark.parametrize(
-    "cwl_file, inputs_file, destination_source_input_data",
+    "cwl_file, inputs_file, chunk, destination_source_input_data",
     [
         # --- Job Grouping ---
-        # Count and list files in inputs_file
+        # 5 files with chunk size 2 → 3 jobs (2/2/1)
         (
             "test/workflows/job_grouping/job_grouping.cwl",
             "test/workflows/pi/type_dependencies/job/inputs-pi_gather_catalog.yaml",
+            "input-data=2",
             {"filecatalog/pi/100": [f"test/workflows/pi/type_dependencies/job/result_{i}.sim" for i in range(1, 6)]},
-        )
+        ),
+        # 5 files with chunk size 3 → 2 jobs (3/2)
+        (
+            "test/workflows/job_grouping/job_grouping.cwl",
+            "test/workflows/pi/type_dependencies/job/inputs-pi_gather_catalog.yaml",
+            "input-data=3",
+            {"filecatalog/pi/100": [f"test/workflows/pi/type_dependencies/job/result_{i}.sim" for i in range(1, 6)]},
+        ),
     ],
 )
-def test_run_transformation_with_inputs_file(
-    cli_runner, cleanup, pi_test_files, cwl_file, inputs_file, destination_source_input_data
+def test_run_transformation_with_chunk(
+    cli_runner, cleanup, pi_test_files, cwl_file, inputs_file, chunk, destination_source_input_data
 ):
-    """Test successful transformation submission and execution with inputs file."""
+    """Test successful transformation submission with --chunk flag."""
     create_test_input_datafiles(destination_source_input_data)
 
-    command = ["transformation", "submit", cwl_file]
-
-    if inputs_file:
-        command.extend(["--inputs-file", inputs_file])
+    command = ["transformation", "submit", cwl_file, "--inputs-file", inputs_file, "--chunk", chunk]
 
     result = cli_runner.invoke(app, command)
     clean_output = strip_ansi_codes(result.stdout)
diff --git a/test/workflows/job_grouping/job_grouping.cwl b/test/workflows/job_grouping/job_grouping.cwl
@@ -19,7 +19,6 @@ outputs:
 
 hints:
   - class: dirac:ExecutionHooks
-    group_size: 2 # will create nb_inputs // 2 jobs
     output_sandbox: ["result-file"]
 
 steps:
