Convert network_config dict to NetworkConfig object in SageMaker orchestrator

src/zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py

@@ -20,6 +20,7 @@
 
 import boto3
 import sagemaker
+from sagemaker.network import NetworkConfig
 from sagemaker.processing import ProcessingInput, ProcessingOutput
 from sagemaker.workflow.execution_variables import ExecutionVariables
 from sagemaker.workflow.pipeline import Pipeline
@@ -256,6 +257,24 @@ def prepare_or_run_pipeline(
             processor_args_for_step["base_job_name"] = orchestrator_run_name
             processor_args_for_step["env"] = environment
 
+            # Convert network_config to sagemaker.network.NetworkConfig if present
+            network_config = processor_args_for_step.get("network_config")
+            if network_config and isinstance(network_config, dict):
+                try:
+                    processor_args_for_step["network_config"] = NetworkConfig(**network_config)
+                except TypeError:
+                    # If the network_config passed is not compatible with the NetworkConfig class,
+                    # raise a more informative error.
+                    raise TypeError(
+                        "Expected a sagemaker.network.NetworkConfig compatible object for the network_config argument, "
+                        "but the network_config processor argument is invalid."
+                        "See https://sagemaker.readthedocs.io/en/stable/api/utility/network.html#sagemaker.network.NetworkConfig "
+                        "for more information about the NetworkConfig class."
+                    )
+                except Exception as e:
+                    # Reraise any other error
+                    raise e
+
             # Construct S3 inputs to container for step
             inputs = None