Adjustments for Asset Library in Batch Processing (#2399)

* Add first draft of changes to make workflows processing command accept additional images metadata for each image, based on metadata file associated to each image

* Bump version

* Add CLI extension + fix problem with metadata keys

* Add ability to inject workflows ee init params

* Add inference-cli extensions that let user creating robofql batches

* Fix bugs spotted while testing

* Bump version

* Add tests

* Add additional requirements to CLI tests

* Add additional requirements to CLI tests

* Apply suggestions from code review

* Make linters happy

* Remove dummy dump

Author: PawelPeczek-Roboflow

.github/workflows/unit_tests_inference_cli_x86.yml

@@ -45,7 +45,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install --upgrade setuptools
-          pip install -r requirements/requirements.cli.txt -r requirements/requirements.test.unit.txt -r requirements/requirements.sdk.http.txt
+          pip install -r requirements/requirements.cli.txt -r requirements/_requirements.txt -r requirements/requirements.cpu.txt -r requirements/requirements.test.unit.txt -r requirements/requirements.sdk.http.txt
       - name: 🧪 Unit Tests of Inference CLI
         timeout-minutes: 30
         run: python -m pytest tests/inference_cli/unit_tests

inference/core/version.py

@@ -1,4 +1,4 @@
-__version__ = "1.3.0"
+__version__ = "1.3.1"
 
 
 if __name__ == "__main__":

inference_cli/lib/roboflow_cloud/batch_processing/api_operations.py

@@ -4,7 +4,7 @@
 import string
 from collections import Counter
 from datetime import datetime, timezone
-from typing import Generator, List, Optional, Set, Union
+from typing import Dict, Generator, List, Optional, Set, Union
 
 import backoff
 import requests
@@ -402,6 +402,8 @@ def trigger_job_with_workflows_images_processing(
     inference_backend: Optional[InferenceBackend] = None,
     job_name: Optional[str] = None,
     max_image_failure_rate: Optional[float] = None,
+    images_metadata_part_name: Optional[str] = None,
+    image_metadata_mapping: Optional[Dict[str, str]] = None,
 ) -> str:
     workspace = get_workspace(api_key=api_key)
     compute_configuration = ComputeConfigurationV2(
@@ -411,6 +413,7 @@ def trigger_job_with_workflows_images_processing(
     input_configuration = StagingBatchInputV1(
         batch_id=batch_id,
         part_name=part_name,
+        images_metadata_part=images_metadata_part_name,
     )
     workflow_parameters = None
     if workflow_parameters_path:
@@ -423,6 +426,7 @@ def trigger_job_with_workflows_images_processing(
         persist_images_outputs=save_image_outputs,
         images_outputs_to_be_persisted=image_outputs_to_save,
         aggregation_format=aggregation_format,
+        images_metadata_inputs_mapping=image_metadata_mapping,
     )
     if not job_id:
         job_id = f"job-{_generate_random_string(length=12)}"

inference_cli/lib/roboflow_cloud/batch_processing/core.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 import typer
 from typing_extensions import Annotated
@@ -112,6 +112,24 @@ def show_job_details(
         raise typer.Exit(code=1)
 
 
+def parse_key_value(
+    values: List[str],
+    require_non_empty_keys: bool = True,
+    require_non_empty_values: bool = True,
+) -> Dict[str, str]:
+    result: Dict[str, str] = {}
+    for item in values:
+        if "=" not in item:
+            raise typer.BadParameter(f"'{item}' must be in key=value format")
+        key, value = item.split("=", 1)  # split once, so values may contain '='
+        if require_non_empty_keys and not key:
+            raise typer.BadParameter(f"Empty key in item '{item}'")
+        if require_non_empty_values and not value:
+            raise typer.BadParameter(f"Empty value in item '{item}'")
+        result[key] = value
+    return result
+
+
 @batch_processing_app.command(help="Trigger batch job to process images with Workflow")
 def process_images_with_workflow(
     batch_id: Annotated[
@@ -156,7 +174,15 @@ def process_images_with_workflow(
         typer.Option(
             "--part-name",
             "-p",
-            help="Name of the batch part " "(relevant for multipart batches",
+            help="Name of the batch part with images (relevant for multipart batches)",
+        ),
+    ] = None,
+    images_metadata_part_name: Annotated[
+        Optional[str],
+        typer.Option(
+            "--images-metadata-part-name",
+            "-imp",
+            help="Name of batch part bringing images metadata (relevant for multipart batches)",
         ),
     ] = None,
     machine_type: Annotated[
@@ -243,7 +269,17 @@ def process_images_with_workflow(
     job_name: Annotated[
         Optional[str], typer.Option("--job-name", "-jn", help="Name of your job")
     ] = None,
+    image_metadata_mapping: Annotated[
+        Optional[List[str]],
+        typer.Option(
+            "--metadata-mapping",
+            "-mm",
+            help="Key-value mapping of workflow input to metadata key as workflow_input=metadata_key. Repeatable.",
+        ),
+    ] = None,
 ) -> None:
+    if image_metadata_mapping is not None:
+        image_metadata_mapping = parse_key_value(image_metadata_mapping)
     if api_key is None:
         api_key = ROBOFLOW_API_KEY
     if workers_per_machine is None and machine_size is not None:
@@ -275,6 +311,8 @@ def process_images_with_workflow(
             inference_backend=inference_backend,
             job_name=job_name,
             max_image_failure_rate=max_image_failure_rate,
+            images_metadata_part_name=images_metadata_part_name,
+            image_metadata_mapping=image_metadata_mapping,
         )
         print(f"Triggered job with ID: {job_id}")
     except KeyboardInterrupt:
@@ -331,7 +369,7 @@ def process_videos_with_workflow(
         typer.Option(
             "--part-name",
             "-p",
-            help="Name of the batch part " "(relevant for multipart batches",
+            help="Name of the batch part relevant for multipart batches",
         ),
     ] = None,
     machine_type: Annotated[

inference_cli/lib/roboflow_cloud/batch_processing/entities.py

@@ -99,6 +99,9 @@ class StagingBatchInputV1(BaseModel):
     type: Literal["staging-batch-input-v1"] = Field(default="staging-batch-input-v1")
     batch_id: str = Field(serialization_alias="batchId")
     part_name: Optional[str] = Field(serialization_alias="partName", default=None)
+    images_metadata_part: Optional[str] = Field(
+        serialization_alias="imagesMetadataPart", default=None
+    )
 
 
 class AggregationFormat(str, Enum):
@@ -130,6 +133,10 @@ class WorkflowsProcessingSpecificationV1(BaseModel):
     max_video_fps: Optional[Union[int, float]] = Field(
         serialization_alias="maxVideoFPS", default=None
     )
+    images_metadata_inputs_mapping: Optional[Dict[str, str]] = Field(
+        serialization_alias="imagesMetadataInputsMapping",
+        default=None,
+    )
 
 
 class WorkflowProcessingJobV1(BaseModel):

inference_cli/lib/roboflow_cloud/data_staging/api_operations.py

@@ -8,7 +8,6 @@
 from multiprocessing.pool import Pool, ThreadPool
 from threading import Lock
 from typing import Dict, Generator, List, Optional, Set, TextIO, Tuple, Union
-from urllib.parse import urlparse
 from uuid import uuid4
 
 import backoff
@@ -645,6 +644,86 @@ def trigger_images_references_ingest(
         raise RFAPICallError("Could not decode Roboflow API response.") from error
 
 
+def create_images_batch_from_roboql_query(
+    query: str,
+    batch_id: str,
+    api_key: str,
+    ingest_id: Optional[str] = None,
+    batch_name: Optional[str] = None,
+    notifications_url: Optional[str] = None,
+) -> None:
+    workspace = get_workspace(api_key=api_key)
+    trigger_roboql_batch_ingest(
+        workspace=workspace,
+        batch_id=batch_id,
+        query=query,
+        api_key=api_key,
+        ingest_id=ingest_id,
+        batch_name=batch_name,
+        notifications_url=notifications_url,
+    )
+    print(
+        f"RoboQL ingest accepted for batch '{batch_id}'. The batch is materialised "
+        "asynchronously by the data-staging worker — initial counts will be zero "
+        "until the worker registers shards."
+    )
+    if notifications_url:
+        print(f"Monitor updates that will be sent to: {notifications_url}")
+    print(
+        f"Use `inference rf-cloud data-staging show-batch-details --batch-id {batch_id}` "
+        "to inspect the batch once it appears, and "
+        f"`inference rf-cloud data-staging list-batch-content --batch-id {batch_id}` "
+        "to list its files once the worker has staged them."
+    )
+
+
+@backoff.on_exception(
+    backoff.constant,
+    exception=RetryError,
+    max_tries=3,
+    interval=1,
+)
+def trigger_roboql_batch_ingest(
+    workspace: str,
+    batch_id: str,
+    query: str,
+    api_key: str,
+    ingest_id: Optional[str] = None,
+    batch_name: Optional[str] = None,
+    notifications_url: Optional[str] = None,
+) -> None:
+    params = {}
+    if api_key is not None:
+        params["api_key"] = api_key
+    payload: Dict[str, Union[str, None]] = {"query": query}
+    if batch_name is not None:
+        payload["displayName"] = batch_name
+    if ingest_id is not None:
+        payload["ingestId"] = ingest_id
+    if notifications_url is not None:
+        payload["notificationsURL"] = notifications_url
+    try:
+        response = requests.post(
+            f"{API_BASE_URL}/data-staging/v1/external/{workspace}/batches/{batch_id}/roboql-ingest",
+            params=params,
+            timeout=REQUEST_TIMEOUT,
+            json=payload,
+        )
+    except (ConnectionError, requests.exceptions.ConnectionError) as error:
+        raise RetryError(
+            f"Connectivity error. Try reaching Roboflow API in browser: {API_BASE_URL}"
+        ) from error
+    except Timeout as e:
+        raise RetryError(
+            f"Timeout error. Could not complete `trigger_roboql_batch_ingest(...)` operation "
+            f"within {REQUEST_TIMEOUT} seconds. Verify connectivity and try again. "
+            f"You may also set `ROBOFLOW_API_REQUEST_TIMEOUT` env variable to a larger "
+            f"value (in seconds), e.g. `export ROBOFLOW_API_REQUEST_TIMEOUT=120`. "
+            f"If the problem persists - contact Roboflow support."
+        ) from e
+    handle_response_errors(response=response, operation_name="trigger RoboQL ingest")
+
+
 def create_videos_batch_from_directory(
     directory: str,
     batch_id: str,
@@ -1567,21 +1646,21 @@ def pull_batch_element_to_directory(
         override_existing=override_existing,
     ):
         return None
-    parsed_url = urlparse(file_metadata.download_url)
-    file_name = os.path.basename(parsed_url.path)
-    is_archive = file_name.endswith(".tar.gz") or file_name.endswith(".tar")
+    is_archive = file_metadata.file_name.endswith(
+        ".tar.gz"
+    ) or file_metadata.file_name.endswith(".tar")
     if not is_archive:
         result_path = pull_file_to_directory(
             url=file_metadata.download_url,
             target_directory=target_directory,
-            file_name=file_name,
+            file_name=file_metadata.file_name,
         )
         export_log.denote_export(file_metadata=file_metadata, local_path=result_path)
         return None
     with tempfile.TemporaryDirectory() as tmp_dir:
         download_and_unpack_archive(
             url=file_metadata.download_url,
-            file_name=file_name,
+            file_name=file_metadata.file_name,
             pull_dir=tmp_dir,
             target_dir=target_directory,
         )

inference_cli/lib/roboflow_cloud/data_staging/core.py

@@ -157,7 +157,6 @@ def create_batch_of_images(
         Optional[str],
         typer.Option(
             "--images-dir",
-            "-i",
             help="Path to your images directory to upload (required if data source is 'local-directory')",
         ),
     ] = None,
@@ -178,6 +177,15 @@ def create_batch_of_images(
             "Required for cloud-storage source. Supports S3, GCS, and Azure.",
         ),
     ] = None,
+    query: Annotated[
+        Optional[str],
+        typer.Option(
+            "--query",
+            "-q",
+            help="RoboQL query selecting images from the Roboflow platform (required if data source is 'roboql'). "
+            "The batch is materialised asynchronously by the data-staging worker.",
+        ),
+    ] = None,
     ingest_id: Annotated[
         Optional[str],
         typer.Option(
@@ -275,6 +283,24 @@ def create_batch_of_images(
                 notifications_url=notifications_url,
                 notification_categories=notification_categories,
             )
+        elif source is DataSource.ROBOQL:
+            if query is None:
+                raise ValueError(
+                    "`query` not provided when `roboql` specified as a data source"
+                )
+            if notification_categories:
+                print(
+                    "--notification-category option is not supported for RoboQL ingests - system emits a "
+                    "webhook at the end of ingest."
+                )
+            api_operations.create_images_batch_from_roboql_query(
+                query=query,
+                batch_id=batch_id,
+                api_key=api_key,
+                ingest_id=ingest_id,
+                batch_name=batch_name,
+                notifications_url=notifications_url,
+            )
         else:
             if references is None:
                 raise ValueError(

inference_cli/lib/roboflow_cloud/data_staging/entities.py

@@ -9,6 +9,7 @@ class DataSource(str, Enum):
     LOCAL_DIRECTORY = "local-directory"
     REFERENCES_FILE = "references-file"
     CLOUD_STORAGE = "cloud-storage"
+    ROBOQL = "roboql"
 
 
 class ShardDetails(BaseModel):

inference_cli/lib/workflows/common.py

@@ -4,6 +4,7 @@
 from collections import defaultdict
 from copy import copy
 from datetime import datetime
+from pathlib import Path
 from threading import Lock
 from typing import Any, Dict, List, Optional, Set, TextIO, Tuple