You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "Cloud-Architect-Emma (via GitHub)" <gi...@apache.org> on 2026/04/20 21:34:48 UTC
[PR] SageMakerNotebookJobTrigger [airflow]
Cloud-Architect-Emma opened a new pull request, #65571:
URL: https://github.com/apache/airflow/pull/65571
Adds async support to SageMakerNotebookJobTrigger, enabling non‑blocking status polling for SageMaker Notebook Jobs. This update aligns the trigger with Airflow’s async architecture, improves efficiency for long‑running jobs, and enhances the structure of emitted events. Tests have been updated to cover the new async behavior.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adds async support to SageMakerNotebookJobTrigger [airflow]
Posted by "Cloud-Architect-Emma (via GitHub)" <gi...@apache.org>.
Cloud-Architect-Emma commented on code in PR #65571:
URL: https://github.com/apache/airflow/pull/65571#discussion_r3125750002
##########
providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker_unified_studio.py:
##########
@@ -7,60 +7,100 @@
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""This module contains the Amazon SageMaker Unified Studio Notebook job trigger."""
from __future__ import annotations
-from airflow.triggers.base import BaseTrigger
+import asyncio
+from typing import Any, Dict
+
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+from airflow.providers.amazon.aws.hooks.sagemaker_unified_studio import (
+ SageMakerNotebookHook,
+)
class SageMakerNotebookJobTrigger(BaseTrigger):
"""
- Watches for a notebook job, triggers when it finishes.
-
- Examples:
- .. code-block:: python
-
- from airflow.providers.amazon.aws.triggers.sagemaker_unified_studio import SageMakerNotebookJobTrigger
+ Async trigger that polls a SageMaker Unified Studio notebook execution until it
+ reaches a terminal state.
- notebook_trigger = SageMakerNotebookJobTrigger(
- execution_id="notebook_job_1234",
- execution_name="notebook_task",
- waiter_delay=10,
- waiter_max_attempts=1440,
- )
-
- :param execution_id: A unique, meaningful id for the task.
- :param execution_name: A unique, meaningful name for the task.
- :param waiter_delay: Interval in seconds to check the notebook execution status.
- :param waiter_max_attempts: Number of attempts to wait before returning FAILED.
+ :param execution_id: The execution ID returned by start_notebook_execution().
+ :param execution_name: The name of the execution (task_id).
+ :param waiter_delay: Seconds between polling attempts.
+ :param waiter_max_attempts: Maximum polling attempts before timing out.
"""
- def __init__(self, execution_id, execution_name, waiter_delay, waiter_max_attempts, **kwargs):
+ def __init__(
+ self,
+ *,
+ execution_id: str,
+ execution_name: str,
+ waiter_delay: int,
+ waiter_max_attempts: int,
+ **kwargs,
+ ):
super().__init__(**kwargs)
self.execution_id = execution_id
self.execution_name = execution_name
self.waiter_delay = waiter_delay
self.waiter_max_attempts = waiter_max_attempts
- def serialize(self):
+ def serialize(self) -> tuple[str, Dict[str, Any]]:
return (
- # dynamically generate the fully qualified name of the class
- self.__class__.__module__ + "." + self.__class__.__qualname__,
+ "airflow.providers.amazon.aws.triggers.sagemaker_unified_studio.SageMakerNotebookJobTrigger",
{
"execution_id": self.execution_id,
"execution_name": self.execution_name,
- "poll_interval": self.poll_interval,
+ "waiter_delay": self.waiter_delay,
+ "waiter_max_attempts": self.waiter_max_attempts,
},
)
async def run(self):
- pass
+ hook = SageMakerNotebookHook(execution_name=self.execution_name)
+ loop = asyncio.get_running_loop()
+
+ attempts = 0
+
+ while attempts < self.waiter_max_attempts:
+ attempts += 1
+
+ response = await loop.run_in_executor(
+ None,
+ hook._sagemaker_studio.execution_client.get_execution,
Review Comment:
Wrapped the call in a public hook method get_notebook_execution.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adds async support to SageMakerNotebookJobTrigger [airflow]
Posted by "Cloud-Architect-Emma (via GitHub)" <gi...@apache.org>.
Cloud-Architect-Emma commented on code in PR #65571:
URL: https://github.com/apache/airflow/pull/65571#discussion_r3126468095
##########
providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker_unified_studio.py:
##########
@@ -7,60 +7,100 @@
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
Review Comment:
I have restored the full Apache license header.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adds async support to SageMakerNotebookJobTrigger [airflow]
Posted by "boring-cyborg[bot] (via GitHub)" <gi...@apache.org>.
boring-cyborg[bot] commented on PR #65571:
URL: https://github.com/apache/airflow/pull/65571#issuecomment-4426186704
Awesome work, congrats on your first merged pull request! You are invited to check our [Issue Tracker](https://github.com/apache/airflow/issues) for additional contributions.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adds async support to SageMakerNotebookJobTrigger [airflow]
Posted by "Cloud-Architect-Emma (via GitHub)" <gi...@apache.org>.
Cloud-Architect-Emma commented on PR #65571:
URL: https://github.com/apache/airflow/pull/65571#issuecomment-4304034473
@o-nikolas just wanted to follow up, all threads are resolved and the PR is marked ready for review. Please let me know if anything else needs attention!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adds async support to SageMakerNotebookJobTrigger [airflow]
Posted by "Cloud-Architect-Emma (via GitHub)" <gi...@apache.org>.
Cloud-Architect-Emma commented on PR #65571:
URL: https://github.com/apache/airflow/pull/65571#issuecomment-4366197110
@shahar1 @o-nikolas reverted all unrelated changes, PR now only touches the 3 SageMaker trigger files with 1 clean commit. Ready for review!
Thank you.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] SageMakerNotebookJobTrigger [airflow]
Posted by "boring-cyborg[bot] (via GitHub)" <gi...@apache.org>.
boring-cyborg[bot] commented on PR #65571:
URL: https://github.com/apache/airflow/pull/65571#issuecomment-4284436623
Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our [Contributors' Guide](https://github.com/apache/airflow/blob/main/contributing-docs/README.rst)
Here are some useful points:
- Pay attention to the quality of your code (ruff, mypy and type annotations). Our [prek-hooks]( https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst#prerequisites-for-prek-hooks) will help you with that.
- In case of a new feature add useful documentation (in docstrings or in `docs/` directory). Adding a new operator? Check this short [guide](https://github.com/apache/airflow/blob/main/airflow-core/docs/howto/custom-operator.rst) Consider adding an example DAG that shows how users should use it.
- Consider using [Breeze environment](https://github.com/apache/airflow/blob/main/dev/breeze/doc/README.rst) for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
- Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
- Please follow [ASF Code of Conduct](https://www.apache.org/foundation/policies/conduct) for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
- Be sure to read the [Airflow Coding style]( https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#coding-style-and-best-practices).
- Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
Apache Airflow is a community-driven project and together we are making it better 🚀.
In case of doubts contact the developers at:
Mailing List: dev@airflow.apache.org
Slack: https://s.apache.org/airflow-slack
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adds async support to SageMakerNotebookJobTrigger [airflow]
Posted by "Cloud-Architect-Emma (via GitHub)" <gi...@apache.org>.
Cloud-Architect-Emma commented on code in PR #65571:
URL: https://github.com/apache/airflow/pull/65571#discussion_r3126470944
##########
providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker_unified_studio.py:
##########
@@ -7,60 +7,100 @@
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""This module contains the Amazon SageMaker Unified Studio Notebook job trigger."""
from __future__ import annotations
-from airflow.triggers.base import BaseTrigger
+import asyncio
+from typing import Any, Dict
Review Comment:
Removed Dict, now using builtin dict directly.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adds async support to SageMakerNotebookJobTrigger [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #65571:
URL: https://github.com/apache/airflow/pull/65571#issuecomment-4299547367
@Cloud-Architect-Emma This PR has been converted to **draft** because it does not yet meet our [Pull Request quality criteria](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-quality-criteria).
**Issues found:**
- :x: **Unresolved review comments** (1 thread from a maintainer): please walk through each unresolved review thread. **Even if a suggestion looks incorrect or irrelevant** — and some of them will be, especially any comments left by automated reviewers like GitHub Copilot — it is still the author's responsibility to respond: apply the fix, reply in-thread with a brief explanation of why the suggestion does not apply, or resolve the thread if the feedback is no longer relevant. Leaving threads unaddressed for weeks blocks the PR from moving forward.
**What to do next:**
- Walk through each unresolved review thread and respond as described above.
- Make sure static checks pass locally: `prek run --from-ref main --stage pre-commit`.
- Mark the PR as "Ready for review" when you're done.
Converting a PR to draft is **not** a rejection — it is an invitation to bring the PR up to the project's standards so that maintainer review time is spent productively. There is no rush — take your time and work at your own pace. We appreciate your contribution and are happy to wait for updates. If you have questions, feel free to ask on the [Airflow Slack](https://s.apache.org/airflow-slack).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adds async support to SageMakerNotebookJobTrigger [airflow]
Posted by "shahar1 (via GitHub)" <gi...@apache.org>.
shahar1 commented on PR #65571:
URL: https://github.com/apache/airflow/pull/65571#issuecomment-4362856138
Could you please revert changes to unrelated files? (for example, added newlines in README.md)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adds async support to SageMakerNotebookJobTrigger [airflow]
Posted by "Cloud-Architect-Emma (via GitHub)" <gi...@apache.org>.
Cloud-Architect-Emma commented on PR #65571:
URL: https://github.com/apache/airflow/pull/65571#issuecomment-4416307706
@potiuk thank you for the detailed review! Fixed both issues:
1. Removed the literal \nescape characters from_set_xcom_filesand_set_xcom_s3_path, they are now proper method definitions
2. Fixedtest_trigger_running_then_timeout, removed the extragen.asend(None)` call since the trigger only yields once on timeout
The notes about terminal state constants and kwargs in serialize() are valid, happy to address those in a follow-up if you'd prefer to keep this PR focused.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adds async support to SageMakerNotebookJobTrigger [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #65571:
URL: https://github.com/apache/airflow/pull/65571#issuecomment-4300312149
Quick follow-up to the triage comment above — one clarification on the "Unresolved review comments" item:
Once you believe a thread has been addressed — whether by pushing a fix, or by replying in-thread with an explanation of why the suggestion doesn't apply — **please mark the thread as resolved yourself** by clicking the "Resolve conversation" button at the bottom of each thread. Reviewers don't auto-close their own threads, so an addressed-but-unresolved thread reads as "still waiting on the author" and keeps the PR from moving forward. The author doing the resolve-click is the expected convention on this project.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adds async support to SageMakerNotebookJobTrigger [airflow]
Posted by "Cloud-Architect-Emma (via GitHub)" <gi...@apache.org>.
Cloud-Architect-Emma commented on PR #65571:
URL: https://github.com/apache/airflow/pull/65571#issuecomment-4300825579
@potiuk Thanks for the review and guidance. I’ve addressed the feedback and resolved the outstanding thread. I’ve also verified the changes locally and marked the PR as ready for review. I appreciate the detailed review process.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adds async support to SageMakerNotebookJobTrigger [airflow]
Posted by "Cloud-Architect-Emma (via GitHub)" <gi...@apache.org>.
Cloud-Architect-Emma commented on PR #65571:
URL: https://github.com/apache/airflow/pull/65571#issuecomment-4289274707
@shahar1 I have added async unit tests covering success, failure, running, timeout, and serialization. All tests passing locally.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adds async support to SageMakerNotebookJobTrigger [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #65571:
URL: https://github.com/apache/airflow/pull/65571#discussion_r3125699625
##########
providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker_unified_studio.py:
##########
@@ -7,60 +7,100 @@
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
Review Comment:
Why did you remove part of the license? This will fail pre checks
##########
providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker_unified_studio.py:
##########
@@ -7,60 +7,100 @@
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""This module contains the Amazon SageMaker Unified Studio Notebook job trigger."""
from __future__ import annotations
-from airflow.triggers.base import BaseTrigger
+import asyncio
+from typing import Any, Dict
Review Comment:
You don't need to import Dict for typing, just use the regular builtin
##########
providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker_unified_studio.py:
##########
@@ -7,60 +7,100 @@
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""This module contains the Amazon SageMaker Unified Studio Notebook job trigger."""
from __future__ import annotations
-from airflow.triggers.base import BaseTrigger
+import asyncio
+from typing import Any, Dict
+
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+from airflow.providers.amazon.aws.hooks.sagemaker_unified_studio import (
+ SageMakerNotebookHook,
+)
class SageMakerNotebookJobTrigger(BaseTrigger):
"""
- Watches for a notebook job, triggers when it finishes.
-
- Examples:
- .. code-block:: python
-
- from airflow.providers.amazon.aws.triggers.sagemaker_unified_studio import SageMakerNotebookJobTrigger
+ Async trigger that polls a SageMaker Unified Studio notebook execution until it
+ reaches a terminal state.
- notebook_trigger = SageMakerNotebookJobTrigger(
- execution_id="notebook_job_1234",
- execution_name="notebook_task",
- waiter_delay=10,
- waiter_max_attempts=1440,
- )
-
- :param execution_id: A unique, meaningful id for the task.
- :param execution_name: A unique, meaningful name for the task.
- :param waiter_delay: Interval in seconds to check the notebook execution status.
- :param waiter_max_attempts: Number of attempts to wait before returning FAILED.
+ :param execution_id: The execution ID returned by start_notebook_execution().
+ :param execution_name: The name of the execution (task_id).
+ :param waiter_delay: Seconds between polling attempts.
+ :param waiter_max_attempts: Maximum polling attempts before timing out.
"""
- def __init__(self, execution_id, execution_name, waiter_delay, waiter_max_attempts, **kwargs):
+ def __init__(
+ self,
+ *,
+ execution_id: str,
+ execution_name: str,
+ waiter_delay: int,
+ waiter_max_attempts: int,
+ **kwargs,
+ ):
super().__init__(**kwargs)
self.execution_id = execution_id
self.execution_name = execution_name
self.waiter_delay = waiter_delay
self.waiter_max_attempts = waiter_max_attempts
- def serialize(self):
+ def serialize(self) -> tuple[str, Dict[str, Any]]:
return (
- # dynamically generate the fully qualified name of the class
- self.__class__.__module__ + "." + self.__class__.__qualname__,
+ "airflow.providers.amazon.aws.triggers.sagemaker_unified_studio.SageMakerNotebookJobTrigger",
{
"execution_id": self.execution_id,
"execution_name": self.execution_name,
- "poll_interval": self.poll_interval,
+ "waiter_delay": self.waiter_delay,
+ "waiter_max_attempts": self.waiter_max_attempts,
},
)
async def run(self):
- pass
+ hook = SageMakerNotebookHook(execution_name=self.execution_name)
+ loop = asyncio.get_running_loop()
+
+ attempts = 0
+
+ while attempts < self.waiter_max_attempts:
+ attempts += 1
+
+ response = await loop.run_in_executor(
+ None,
+ hook._sagemaker_studio.execution_client.get_execution,
Review Comment:
This is accessing the "private" client attribute. Maybe add a little helper function to the hook class that encapsulates this.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] SageMakerNotebookJobTrigger [airflow]
Posted by "Cloud-Architect-Emma (via GitHub)" <gi...@apache.org>.
Cloud-Architect-Emma commented on PR #65571:
URL: https://github.com/apache/airflow/pull/65571#issuecomment-4284525913
Hi @o-nikolas,
When you have a moment, could you please take a look at this PR? I’ve added async support to SageMakerNotebookJobTrigger and updated the related tests. Everything is passing on my side, but I’d appreciate your review to confirm the approach aligns with the provider’s patterns.
Thanks for your time!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adds async support to SageMakerNotebookJobTrigger [airflow]
Posted by "Cloud-Architect-Emma (via GitHub)" <gi...@apache.org>.
Cloud-Architect-Emma commented on code in PR #65571:
URL: https://github.com/apache/airflow/pull/65571#discussion_r3125750002
##########
providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker_unified_studio.py:
##########
@@ -7,60 +7,100 @@
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""This module contains the Amazon SageMaker Unified Studio Notebook job trigger."""
from __future__ import annotations
-from airflow.triggers.base import BaseTrigger
+import asyncio
+from typing import Any, Dict
+
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+from airflow.providers.amazon.aws.hooks.sagemaker_unified_studio import (
+ SageMakerNotebookHook,
+)
class SageMakerNotebookJobTrigger(BaseTrigger):
"""
- Watches for a notebook job, triggers when it finishes.
-
- Examples:
- .. code-block:: python
-
- from airflow.providers.amazon.aws.triggers.sagemaker_unified_studio import SageMakerNotebookJobTrigger
+ Async trigger that polls a SageMaker Unified Studio notebook execution until it
+ reaches a terminal state.
- notebook_trigger = SageMakerNotebookJobTrigger(
- execution_id="notebook_job_1234",
- execution_name="notebook_task",
- waiter_delay=10,
- waiter_max_attempts=1440,
- )
-
- :param execution_id: A unique, meaningful id for the task.
- :param execution_name: A unique, meaningful name for the task.
- :param waiter_delay: Interval in seconds to check the notebook execution status.
- :param waiter_max_attempts: Number of attempts to wait before returning FAILED.
+ :param execution_id: The execution ID returned by start_notebook_execution().
+ :param execution_name: The name of the execution (task_id).
+ :param waiter_delay: Seconds between polling attempts.
+ :param waiter_max_attempts: Maximum polling attempts before timing out.
"""
- def __init__(self, execution_id, execution_name, waiter_delay, waiter_max_attempts, **kwargs):
+ def __init__(
+ self,
+ *,
+ execution_id: str,
+ execution_name: str,
+ waiter_delay: int,
+ waiter_max_attempts: int,
+ **kwargs,
+ ):
super().__init__(**kwargs)
self.execution_id = execution_id
self.execution_name = execution_name
self.waiter_delay = waiter_delay
self.waiter_max_attempts = waiter_max_attempts
- def serialize(self):
+ def serialize(self) -> tuple[str, Dict[str, Any]]:
return (
- # dynamically generate the fully qualified name of the class
- self.__class__.__module__ + "." + self.__class__.__qualname__,
+ "airflow.providers.amazon.aws.triggers.sagemaker_unified_studio.SageMakerNotebookJobTrigger",
{
"execution_id": self.execution_id,
"execution_name": self.execution_name,
- "poll_interval": self.poll_interval,
+ "waiter_delay": self.waiter_delay,
+ "waiter_max_attempts": self.waiter_max_attempts,
},
)
async def run(self):
- pass
+ hook = SageMakerNotebookHook(execution_name=self.execution_name)
+ loop = asyncio.get_running_loop()
+
+ attempts = 0
+
+ while attempts < self.waiter_max_attempts:
+ attempts += 1
+
+ response = await loop.run_in_executor(
+ None,
+ hook._sagemaker_studio.execution_client.get_execution,
Review Comment:
okay.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adds async support to SageMakerNotebookJobTrigger [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk merged PR #65571:
URL: https://github.com/apache/airflow/pull/65571
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Adds async support to SageMakerNotebookJobTrigger [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on PR #65571:
URL: https://github.com/apache/airflow/pull/65571#issuecomment-4330078035
> @o-nikolas just wanted to follow up, all threads are resolved and the PR is marked ready for review. Please let me know if anything else needs attention!
There are many tests failing, please ensure these are green before marking ready for review, thanks!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org