fix: pass attribute aether for AetherBridgeComponent (#35534)

Azure · May 9, 2024 · 7e70f4b · 7e70f4b
1 parent 3a2c666
commit 7e70f4b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/component.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/component.py
@@ -103,6 +103,7 @@ def __init__(
  ae365exepool: Optional[Dict] = None,
  launcher: Optional[Dict] = None,
  datatransfer: Optional[Dict] = None,
+ aether: Optional[Dict] = None,
  **kwargs,
  ):
  _type, self._type_label = parse_name_label(type)
@@ -141,6 +142,7 @@ def __init__(
  self.ae365exepool = ae365exepool
  self.launcher = launcher
  self.datatransfer = datatransfer
+ self.aether = aether
 
  @classmethod
  def _build_io(cls, io_dict: Union[Dict, Input, Output], is_input: bool):

diff --git a/sdk/ml/azure-ai-ml/tests/internal/unittests/test_component.py b/sdk/ml/azure-ai-ml/tests/internal/unittests/test_component.py
@@ -989,6 +989,7 @@ def test_load_from_internal_aether_bridge_component(self):
  "mock_param1": {"type": "AnyFile", "optional": False},
  "mock_param2": {"type": "AnyFile", "optional": False},
  },
+ "aether": {"module_type": "ScrapingCloud", "ref_id": "mock_ref_id"},
  "outputs": {"job_info": {"type": "AnyFile"}},
  "type": "AetherBridgeComponent",
  "command": "mock.exe {inputs.mock_param1} {inputs.mock_param2} {outputs.job_info}",