Fix marshmallow 4.x context parameter compatibility issues - Phase 1

Co-authored-by: kshitij-microsoft <[email protected]>

Author: Copilot

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

@@ -123,9 +123,9 @@ def _deserialize(self, value, attr, data, **kwargs):
         # Update base_path to parent path of component file.
         component_schema_context = deepcopy(self.context)
         component_schema_context[BASE_PATH_CONTEXT_KEY] = source_path.parent
-        component = AnonymousCommandComponentSchema(context=component_schema_context).load(
+        component = AnonymousCommandComponentSchema().load(
             component_dict
-        )
+        , context=component_schema_context)
         component._source_path = source_path
         component._source = ComponentSource.YAML_COMPONENT
         return component

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

@@ -206,9 +206,9 @@ def _deserialize(self, value, attr, data, **kwargs):
         # Update base_path to parent path of component file.
         component_schema_context = deepcopy(self.context)
         component_schema_context[BASE_PATH_CONTEXT_KEY] = source_path.parent
-        component = AnonymousDataTransferCopyComponentSchema(context=component_schema_context).load(
+        component = AnonymousDataTransferCopyComponentSchema().load(
             component_dict
-        )
+        , context=component_schema_context)
         component._source_path = source_path
         component._source = ComponentSource.YAML_COMPONENT
         return component
@@ -224,9 +224,9 @@ def _deserialize(self, value, attr, data, **kwargs):
         # Update base_path to parent path of component file.
         component_schema_context = deepcopy(self.context)
         component_schema_context[BASE_PATH_CONTEXT_KEY] = source_path.parent
-        component = AnonymousDataTransferImportComponentSchema(context=component_schema_context).load(
+        component = AnonymousDataTransferImportComponentSchema().load(
             component_dict
-        )
+        , context=component_schema_context)
         component._source_path = source_path
         component._source = ComponentSource.YAML_COMPONENT
         return component
@@ -242,9 +242,9 @@ def _deserialize(self, value, attr, data, **kwargs):
         # Update base_path to parent path of component file.
         component_schema_context = deepcopy(self.context)
         component_schema_context[BASE_PATH_CONTEXT_KEY] = source_path.parent
-        component = AnonymousDataTransferExportComponentSchema(context=component_schema_context).load(
+        component = AnonymousDataTransferExportComponentSchema().load(
             component_dict
-        )
+        , context=component_schema_context)
         component._source_path = source_path
         component._source = ComponentSource.YAML_COMPONENT
         return component

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

@@ -64,9 +64,9 @@ def _deserialize(self, value, attr, data, **kwargs):
         # Update base_path to parent path of component file.
         component_schema_context = deepcopy(self.context)
         component_schema_context[BASE_PATH_CONTEXT_KEY] = source_path.parent
-        component = AnonymousImportComponentSchema(context=component_schema_context).load(
+        component = AnonymousImportComponentSchema().load(
             component_dict
-        )
+        , context=component_schema_context)
         component._source_path = source_path
         component._source = ComponentSource.YAML_COMPONENT
         return component

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

@@ -94,9 +94,9 @@ def _deserialize(self, value, attr, data, **kwargs):
         # Update base_path to parent path of component file.
         component_schema_context = deepcopy(self.context)
         component_schema_context[BASE_PATH_CONTEXT_KEY] = source_path.parent
-        component = AnonymousParallelComponentSchema(context=component_schema_context).load(
+        component = AnonymousParallelComponentSchema().load(
             component_dict
-        )
+        , context=component_schema_context)
         component._source_path = source_path
         component._source = ComponentSource.YAML_COMPONENT
         return component

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

@@ -70,9 +70,9 @@ def _deserialize(self, value, attr, data, **kwargs):
         # Update base_path to parent path of component file.
         component_schema_context = deepcopy(self.context)
         component_schema_context[BASE_PATH_CONTEXT_KEY] = source_path.parent
-        component = AnonymousSparkComponentSchema(context=component_schema_context).load(
+        component = AnonymousSparkComponentSchema().load(
             component_dict
-        )
+        , context=component_schema_context)
         component._source_path = source_path
         component._source = ComponentSource.YAML_COMPONENT
         return component

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/core/schema.py

@@ -17,6 +17,29 @@
 from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, FILE_PREFIX, PARAMS_OVERRIDE_KEY
 from azure.ai.ml.exceptions import MlException
 
+
+class ContextAwareSchema(PatchedBaseSchema, metaclass=PatchedSchemaMeta):
+    """A marshmallow 4.x compatible schema that can store and use context"""
+    
+    def __init__(self, *args, **kwargs):
+        # Store context as instance variable (marshmallow 4.x doesn't accept it in constructor)
+        self.context = kwargs.pop("context", None)
+        
+        # In marshmallow 4.x, filter out unsupported constructor parameters  
+        # Valid parameters for Schema constructor: only, exclude, many, load_only, dump_only, partial
+        valid_schema_params = {'only', 'exclude', 'many', 'load_only', 'dump_only', 'partial'}
+        filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_schema_params}
+        
+        super().__init__(*args, **filtered_kwargs)
+
+    def load(self, json_data, *, many=None, partial=None, unknown=None):
+        """Override load to use stored context"""
+        return super().load(json_data, many=many, partial=partial, unknown=unknown, context=self.context)
+
+    def dump(self, obj, *, many=None):
+        """Override dump to use stored context"""
+        return super().dump(obj, many=many, context=self.context)
+
 module_logger = logging.getLogger(__name__)
 
 
@@ -34,12 +57,21 @@ def __init__(self, *args, **kwargs):
         self.old_base_path = self.context.get(BASE_PATH_CONTEXT_KEY)
 
         # In marshmallow 4.x, filter out unsupported constructor parameters
-        # Valid parameters for Schema constructor: only, exclude, many, context, load_only, dump_only, partial
-        valid_schema_params = {'only', 'exclude', 'many', 'context', 'load_only', 'dump_only', 'partial'}
+        # Valid parameters for Schema constructor: only, exclude, many, load_only, dump_only, partial
+        # Note: context is NOT a valid constructor parameter in marshmallow 4.x
+        valid_schema_params = {'only', 'exclude', 'many', 'load_only', 'dump_only', 'partial'}
         filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_schema_params}
 
         super().__init__(*args, **filtered_kwargs)
 
+    def load(self, json_data, *, many=None, partial=None, unknown=None):
+        """Override load to use stored context"""
+        return super().load(json_data, many=many, partial=partial, unknown=unknown, context=self.context)
+
+    def dump(self, obj, *, many=None):
+        """Override dump to use stored context"""
+        return super().dump(obj, many=many, context=self.context)
+
     @pre_load
     def add_param_overrides(self, data, **kwargs):
         # Removing params override from context so that overriding is done once on the yaml

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

@@ -247,9 +247,9 @@ def _deserialize(self, value, attr, data, **kwargs):
         # Update base_path to parent path of component file.
         component_schema_context = deepcopy(self.context)
         component_schema_context[BASE_PATH_CONTEXT_KEY] = source_path.parent
-        component = _AnonymousPipelineComponentSchema(context=component_schema_context).load(
+        component = _AnonymousPipelineComponentSchema().load(
             component_dict
-        )
+        , context=component_schema_context)
         component._source_path = source_path
         component._source = ComponentSource.YAML_COMPONENT
         return component

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/_artifacts/_package/base_environment_source.py

@@ -42,7 +42,7 @@ def _from_rest_object(cls, rest_obj: RestBaseEnvironmentId) -> "RestBaseEnvironm
         return BaseEnvironment(type=rest_obj.base_environment_source_type, resource_id=rest_obj.resource_id)
 
     def _to_dict(self) -> Dict:
-        return dict(BaseEnvironmentSourceSchema(context={BASE_PATH_CONTEXT_KEY: "./"}).dump(self))
+        return dict(BaseEnvironmentSourceSchema().dump(self, context={BASE_PATH_CONTEXT_KEY: "./"}))
 
     def _to_rest_object(self) -> RestBaseEnvironmentId:
         return RestBaseEnvironmentId(base_environment_source_type=self.type, resource_id=self.resource_id)

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/_artifacts/_package/model_package.py

@@ -296,7 +296,7 @@ def dump(
         dump_yaml_to_file(dest, yaml_serialized, default_flow_style=False)
 
     def _to_dict(self) -> Dict:
-        return dict(ModelPackageSchema(context={BASE_PATH_CONTEXT_KEY: "./"}).dump(self))
+        return dict(ModelPackageSchema().dump(self, context={BASE_PATH_CONTEXT_KEY: "./"}))
 
     @classmethod
     def _from_rest_object(cls, model_package_rest_object: PackageResponse) -> Any:

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/_artifacts/code.py

@@ -85,7 +85,7 @@ def _load(
         return res
 
     def _to_dict(self) -> Dict:
-        res: dict = CodeAssetSchema(context={BASE_PATH_CONTEXT_KEY: "./"}).dump(self)
+        res: dict = CodeAssetSchema().dump(self, context={BASE_PATH_CONTEXT_KEY: "./"})
         return res
 
     @classmethod