Beat the limitations of EE in terms of singular elements pushed into batch inputs

inference/core/workflows/execution_engine/constants.py

@@ -2,6 +2,7 @@
 PARSED_NODE_INPUT_SELECTORS_PROPERTY = "parsed_node_input_selectors"
 STEP_DEFINITION_PROPERTY = "definition"
 WORKFLOW_INPUT_BATCH_LINEAGE_ID = "<workflow_input>"
+TOP_LEVEL_LINEAGE_KEY = "top_level_lineage"
 IMAGE_TYPE_KEY = "type"
 IMAGE_VALUE_KEY = "value"
 ROOT_PARENT_ID_KEY = "root_parent_id"

inference/core/workflows/execution_engine/introspection/schema_parser.py

@@ -64,12 +64,16 @@ def parse_block_manifest(
     inputs_accepting_batches_and_scalars = set(
         manifest_type.get_parameters_accepting_batches_and_scalars()
     )
+    inputs_enforcing_auto_batch_casting = set(
+        manifest_type.get_parameters_enforcing_auto_batch_casting()
+    )
     return parse_block_manifest_schema(
         schema=schema,
         inputs_dimensionality_offsets=inputs_dimensionality_offsets,
         dimensionality_reference_property=dimensionality_reference_property,
         inputs_accepting_batches=inputs_accepting_batches,
         inputs_accepting_batches_and_scalars=inputs_accepting_batches_and_scalars,
+        inputs_enforcing_auto_batch_casting=inputs_enforcing_auto_batch_casting,
     )
 
 
@@ -79,6 +83,7 @@ def parse_block_manifest_schema(
     dimensionality_reference_property: Optional[str],
     inputs_accepting_batches: Set[str],
     inputs_accepting_batches_and_scalars: Set[str],
+    inputs_enforcing_auto_batch_casting: Set[str],
 ) -> BlockManifestMetadata:
     primitive_types = retrieve_primitives_from_schema(
         schema=schema,
@@ -89,6 +94,7 @@ def parse_block_manifest_schema(
         dimensionality_reference_property=dimensionality_reference_property,
         inputs_accepting_batches=inputs_accepting_batches,
         inputs_accepting_batches_and_scalars=inputs_accepting_batches_and_scalars,
+        inputs_enforcing_auto_batch_casting=inputs_enforcing_auto_batch_casting,
     )
     return BlockManifestMetadata(
         primitive_types=primitive_types,
@@ -255,6 +261,7 @@ def retrieve_selectors_from_schema(
     dimensionality_reference_property: Optional[str],
     inputs_accepting_batches: Set[str],
     inputs_accepting_batches_and_scalars: Set[str],
+    inputs_enforcing_auto_batch_casting: Set[str],
 ) -> Dict[str, SelectorDefinition]:
     result = []
     for property_name, property_definition in schema[PROPERTIES_KEY].items():
@@ -277,6 +284,7 @@ def retrieve_selectors_from_schema(
                 is_list_element=True,
                 inputs_accepting_batches=inputs_accepting_batches,
                 inputs_accepting_batches_and_scalars=inputs_accepting_batches_and_scalars,
+                inputs_enforcing_auto_batch_casting=inputs_enforcing_auto_batch_casting,
             )
         elif property_definition.get(TYPE_KEY) == OBJECT_TYPE and isinstance(
             property_definition.get(ADDITIONAL_PROPERTIES_KEY), dict
@@ -290,6 +298,7 @@ def retrieve_selectors_from_schema(
                 is_dict_element=True,
                 inputs_accepting_batches=inputs_accepting_batches,
                 inputs_accepting_batches_and_scalars=inputs_accepting_batches_and_scalars,
+                inputs_enforcing_auto_batch_casting=inputs_enforcing_auto_batch_casting,
             )
         else:
             selector = retrieve_selectors_from_simple_property(
@@ -300,6 +309,7 @@ def retrieve_selectors_from_schema(
                 is_dimensionality_reference_property=is_dimensionality_reference_property,
                 inputs_accepting_batches=inputs_accepting_batches,
                 inputs_accepting_batches_and_scalars=inputs_accepting_batches_and_scalars,
+                inputs_enforcing_auto_batch_casting=inputs_enforcing_auto_batch_casting,
             )
         if selector is not None:
             result.append(selector)
@@ -314,6 +324,7 @@ def retrieve_selectors_from_simple_property(
     is_dimensionality_reference_property: bool,
     inputs_accepting_batches: Set[str],
     inputs_accepting_batches_and_scalars: Set[str],
+    inputs_enforcing_auto_batch_casting: Set[str],
     is_list_element: bool = False,
     is_dict_element: bool = False,
 ) -> Optional[SelectorDefinition]:
@@ -325,7 +336,10 @@ def retrieve_selectors_from_simple_property(
             if property_name in inputs_accepting_batches_and_scalars:
                 points_to_batch = {True, False}
             else:
-                points_to_batch = {property_name in inputs_accepting_batches}
+                points_to_batch = {
+                    property_name in inputs_accepting_batches
+                    or property_name in inputs_enforcing_auto_batch_casting
+                }
         else:
             points_to_batch = {declared_points_to_batch}
         allowed_references = [
@@ -359,6 +373,7 @@ def retrieve_selectors_from_simple_property(
             is_dimensionality_reference_property=is_dimensionality_reference_property,
             inputs_accepting_batches=inputs_accepting_batches,
             inputs_accepting_batches_and_scalars=inputs_accepting_batches_and_scalars,
+            inputs_enforcing_auto_batch_casting=inputs_enforcing_auto_batch_casting,
             is_list_element=True,
         )
     if property_defines_union(property_definition=property_definition):
@@ -372,6 +387,7 @@ def retrieve_selectors_from_simple_property(
             is_dimensionality_reference_property=is_dimensionality_reference_property,
             inputs_accepting_batches=inputs_accepting_batches,
             inputs_accepting_batches_and_scalars=inputs_accepting_batches_and_scalars,
+            inputs_enforcing_auto_batch_casting=inputs_enforcing_auto_batch_casting,
         )
     return None
 
@@ -394,6 +410,7 @@ def retrieve_selectors_from_union_definition(
     is_dimensionality_reference_property: bool,
     inputs_accepting_batches: Set[str],
     inputs_accepting_batches_and_scalars: Set[str],
+    inputs_enforcing_auto_batch_casting: Set[str],
 ) -> Optional[SelectorDefinition]:
     union_types = (
         union_definition.get(ANY_OF_KEY, [])
@@ -410,6 +427,7 @@ def retrieve_selectors_from_union_definition(
             is_dimensionality_reference_property=is_dimensionality_reference_property,
             inputs_accepting_batches=inputs_accepting_batches,
             inputs_accepting_batches_and_scalars=inputs_accepting_batches_and_scalars,
+            inputs_enforcing_auto_batch_casting=inputs_enforcing_auto_batch_casting,
             is_list_element=is_list_element,
         )
         if result is None:

inference/core/workflows/execution_engine/v1/compiler/entities.py

@@ -216,6 +216,12 @@ def iterate_through_definitions(self) -> Generator[StepInputDefinition, None, No
 StepInputData = Dict[str, Union[StepInputDefinition, CompoundStepInputDefinition]]
 
 
+@dataclass
+class AutoBatchCastingConfig:
+    casted_dimensionality: int
+    lineage_support: Optional[List[str]]
+
+
 @dataclass
 class StepNode(ExecutionGraphNode):
     step_manifest: WorkflowBlockManifest
@@ -224,6 +230,9 @@ class StepNode(ExecutionGraphNode):
     child_execution_branches: Dict[str, str] = field(default_factory=dict)
     execution_branches_impacting_inputs: Set[str] = field(default_factory=set)
     batch_oriented_parameters: Set[str] = field(default_factory=set)
+    auto_batch_casting_lineage_supports: Dict[str, AutoBatchCastingConfig] = field(
+        default_factory=dict
+    )
     step_execution_dimensionality: int = 0
 
     def controls_flow(self) -> bool: