fix: support is_control output for internal components

Author: elliotzh

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/_schema/component.py

@@ -15,6 +15,7 @@
     InternalInputPortSchema,
     InternalOutputPortSchema,
     InternalParameterSchema,
+    InternalPrimitiveOutputSchema,
 )
 
 
@@ -60,7 +61,16 @@ class InternalBaseComponentSchema(ComponentSchema):
             ]
         ),
     )
-    outputs = fields.Dict(keys=fields.Str(), values=NestedField(InternalOutputPortSchema))
+    # support primitive output for all internal components for now
+    outputs = fields.Dict(
+        keys=fields.Str(),
+        values=UnionField(
+            [
+                NestedField(InternalPrimitiveOutputSchema),
+                NestedField(InternalOutputPortSchema),
+            ]
+        ),
+    )
 
     # type field is required for registration
     type = StringTransformedEnum(

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/_schema/input_output.py

@@ -6,9 +6,20 @@
 
 from azure.ai.ml._schema import StringTransformedEnum, UnionField
 from azure.ai.ml._schema.component.input_output import InputPortSchema, OutputPortSchema, ParameterSchema
-from azure.ai.ml._schema.core.fields import DumpableFloatField, DumpableIntegerField
+from azure.ai.ml._schema.core.fields import DumpableFloatField, DumpableIntegerField, DumpableEnumField
 
 
+SUPPORTED_INTERNAL_PARAM_TYPES = [
+    "integer",
+    "Integer",
+    "boolean",
+    "Boolean",
+    "string",
+    "String",
+    "float",
+    "Float",
+]
+
 class InternalInputPortSchema(InputPortSchema):
     # skip client-side validate for type enum & support list
     type = UnionField(
@@ -39,19 +50,17 @@ class InternalOutputPortSchema(OutputPortSchema):
     datastore_mode = fields.Str()
 
 
+class InternalPrimitiveOutputSchema(OutputPortSchema):
+    type = DumpableEnumField(
+        allowed_values=SUPPORTED_INTERNAL_PARAM_TYPES,
+        required=True,
+    )
+    is_control = fields.Bool()
+
+
 class InternalParameterSchema(ParameterSchema):
-    type = StringTransformedEnum(
-        allowed_values=[
-            "integer",
-            "Integer",
-            "boolean",
-            "Boolean",
-            "string",
-            "String",
-            "float",
-            "Float",
-        ],
-        casing_transform=lambda x: x,
+    type = DumpableEnumField(
+        allowed_values=SUPPORTED_INTERNAL_PARAM_TYPES,
         required=True,
         data_key="type",
     )

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/component/command_component.py

@@ -38,7 +38,7 @@ class Meta:
         ],
         metadata={"description": "Provides the configuration for a distributed run."},
     )
-    # primitive output is only supported for command component
+    # primitive output is only supported for command component & pipeline component
     outputs = fields.Dict(
         keys=fields.Str(),
         values=UnionField(

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/component/input_output.py

@@ -51,16 +51,12 @@ class OutputPortSchema(metaclass=PatchedSchemaMeta):
     )
 
 
-class PrimitiveOutputSchema(metaclass=PatchedSchemaMeta):
+class PrimitiveOutputSchema(OutputPortSchema):
     type = DumpableEnumField(
         allowed_values=SUPPORTED_PARAM_TYPES,
         required=True,
     )
-    description = fields.Str()
     is_control = fields.Bool()
-    mode = DumpableEnumField(
-        allowed_values=SUPPORTED_INPUT_OUTPUT_MODES,
-    )
 
 
 class ParameterSchema(metaclass=PatchedSchemaMeta):

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/pipeline/pipeline_component.py

@@ -136,7 +136,7 @@ class PipelineComponentSchema(ComponentSchema):
     type = StringTransformedEnum(allowed_values=[NodeType.PIPELINE])
     jobs = PipelineJobsField()
 
-    # primitive output is only supported for command component
+    # primitive output is only supported for command component & pipeline component
     outputs = fields.Dict(
         keys=fields.Str(),
         values=UnionField(

sdk/ml/azure-ai-ml/tests/internal/unittests/test_component.py

@@ -377,7 +377,7 @@ def test_invalid_additional_includes(self, yaml_path: str, expected_error_msg_pr
         assert validation_result.error_messages["*"].startswith(expected_error_msg_prefix)
 
     def test_component_input_types(self) -> None:
-        yaml_path = "./tests/test_configs/internal/component_with_input_types/component_spec.yaml"
+        yaml_path = "./tests/test_configs/internal/component_with_input_outputs/component_spec.yaml"
         component: InternalComponent = load_component(yaml_path)
         component.code = "scope:1"
 
@@ -403,7 +403,7 @@ def test_component_input_types(self) -> None:
         assert component._validate().passed is True, repr(component._validate())
 
     def test_component_input_with_attrs(self) -> None:
-        yaml_path = "./tests/test_configs/internal/component_with_input_types/component_spec_with_attrs.yaml"
+        yaml_path = "./tests/test_configs/internal/component_with_input_outputs/component_spec_with_attrs.yaml"
         component: InternalComponent = load_component(source=yaml_path)
 
         expected_inputs = {
@@ -433,6 +433,24 @@ def test_component_input_with_attrs(self) -> None:
         assert regenerated_component._to_rest_object().properties.component_spec["inputs"] == expected_inputs["inputs"]
         assert component._validate().passed is True, repr(component._validate())
 
+    def test_component_output_with_attrs(self) -> None:
+        yaml_path = "./tests/test_configs/internal/component_with_input_outputs/component_spec_with_outputs.yaml"
+        component: InternalComponent = load_component(source=yaml_path)
+        assert component
+
+        expected_outputs = {
+            "primitive_is_control": {
+                "is_control": True,
+                "type": "boolean",
+            }
+        }
+        assert component._to_rest_object().properties.component_spec["outputs"] == expected_outputs
+        assert component._validate().passed is True, repr(component._validate())
+
+        regenerated_component = Component._from_rest_object(component._to_rest_object())
+        assert regenerated_component._to_rest_object().properties.component_spec["outputs"] == expected_outputs
+        assert component._validate().passed is True, repr(component._validate())
+
     def test_component_input_list_type(self) -> None:
         yaml_path = "./tests/test_configs/internal/scope-component/component_spec.yaml"
         component: InternalComponent = load_component(yaml_path)

sdk/ml/azure-ai-ml/tests/test_configs/internal/component_with_input_types/component_spec.yaml renamed to sdk/ml/azure-ai-ml/tests/test_configs/internal/component_with_input_outputs/component_spec.yaml

@@ -0,0 +1,22 @@
+$schema: https://componentsdk.azureedge.net/jsonschema/CommandComponent.json
+name: convert_file_to_boolean
+version: 0.0.1
+display_name: Convert File to Boolean
+type: CommandComponent
+description: Convert input file to boolean type output.
+tags:
+  codegenBy: dsl.component
+inputs:
+  input:
+    optional: false
+    type:
+    - AnyFile
+    - AnyDirectory
+outputs:
+  primitive_is_control:
+    is_control: true
+    type: boolean
+command: python -m azure.ml.component.dsl.executor --file run.py --name convert_file_to_boolean
+  --params --input {inputs.input} --output {outputs.output}
+environment:
+  name: AzureML-Component

sdk/ml/azure-ai-ml/tests/test_configs/internal/component_with_input_types/component_spec_with_attrs.yaml renamed to sdk/ml/azure-ai-ml/tests/test_configs/internal/component_with_input_outputs/component_spec_with_attrs.yaml

@@ -0,0 +1,7 @@
+#DECLARE Output_stream string = @@Output_SSPath@@;
+#DECLARE In_Data string =@"@@Input_TextData@@";
+
+RawData = EXTRACT @@ExtractionClause@@ FROM @In_Data
+USING DefaultTextExtractor();
+
+OUTPUT RawData TO SSTREAM @Output_stream;