jroachgolf84 opened a new issue, #72279: URL: https://github.com/apache/airflow/issues/72279
## Description > This is an issue that is reserved for the Airflow Summit "Contributors" Workshop. This is denoted with the label `contributors-workshop`. Out of respect for the organizers and participants of this workshop, **please do not implement a PR that addresses this issue.** > > If this issue is still open following Airflow Summit, the label will be removed and the issue can be picked up. Sub-issue of #72144, covering one of the three defer sites set aside for the Contributor's Workshop in [this comment](https://github.com/apache/airflow/issues/72144#issuecomment-5443098667). | | | | --- | --- | | Call site | `providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker_unified_studio_notebook.py`, in `SageMakerUnifiedStudioNotebookOperator.execute` | | Trigger | `SageMakerUnifiedStudioNotebookTrigger` in `providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker_unified_studio_notebook.py` | | Shape | Call site only, the trigger already accepts the parameters | `SageMakerUnifiedStudioNotebookOperator` is an `AwsBaseOperator`, so it always carries `region_name`, `verify` and `botocore_config`. When it defers it passes none of the three, so the triggerer builds its hook from `aws_conn_id` alone and can end up polling a different region than the one the notebook run was started in. ### What needs to change? Nothing on the trigger side is in the way. `SageMakerUnifiedStudioNotebookTrigger.__init__` takes `**kwargs` and forwards them to `AwsBaseWaiterTrigger`, and its `hook()` already reads all three values off `self`: ```python def hook(self) -> AwsGenericHook: return SageMakerUnifiedStudioNotebookHook( aws_conn_id=self.aws_conn_id, region_name=self.region_name, verify=self.verify, config=self.botocore_config, ) ``` They are never populated because the operator does not send them. The fix is to add the three arguments at the defer site. One thing to be careful of while editing that call. The trigger deliberately computes `waiter_max_attempts` from `timeout_configuration` and pops any `waiter_max_attempts` that arrives through `kwargs`, so do not add one while you are there. ### How to verify it The parent issue's reproduction applies: construct the operator with a non-default `region_name`, `verify` and `botocore_config`, trigger the defer, and assert those values survive into the serialized payload. ```python with pytest.raises(TaskDeferred) as deferred: operator.execute(None) assert deferred.value.trigger.serialize()[1] ``` `AwsBaseWaiterTrigger.serialize` prunes empty values, so use values that are actually distinguishable from the defaults, and assert on the serialized dict rather than on attributes of the trigger object. ### Definition of done 1. Pass `region_name`, `verify` and `botocore_config` through to `SageMakerUnifiedStudioNotebookTrigger` at the defer site. 2. Add or extend a unit test asserting all three survive into the serialized trigger payload, and that the derived `waiter_max_attempts` behaviour is unchanged. 3. Once #72171 has landed, remove the `("operators/sagemaker_unified_studio_notebook.py", "SageMakerUnifiedStudioNotebookTrigger")` entry from the `PENDING_MIGRATION` allowlist that PR introduces. That allowlist does not exist on `main` yet, so this step only applies after it merges. The invariant test asserts each entry is still needed, so a stale line fails the suite. 4. These should pass: ```bash breeze testing providers-tests providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_unified_studio_notebook.py breeze testing providers-tests providers/amazon/tests/unit/amazon/aws/triggers/test_sagemaker_unified_studio_notebook.py ``` Note on sequencing. The code change here is independent of #72171 and can be made and reviewed straight away. Only the allowlist deletion in step 3 has to wait for that PR. --- Drafted-by: Claude Code (Opus 5); reviewed and edited by @jroachgolf84 before posting -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
