airflow.providers.amazon.aws.hooks.sagemaker_unified_studio¶

此模块包含 Amazon SageMaker Unified Studio Notebook 的 Hook。

类¶

SageMakerNotebookHook

与 Sagemaker Unified Studio 工作流交互。

模块内容¶

class airflow.providers.amazon.aws.hooks.sagemaker_unified_studio.SageMakerNotebookHook(execution_name, input_config=None, output_config=None, compute=None, termination_condition=None, tags=None, waiter_delay=10, waiter_max_attempts=1440, *args, **kwargs)[源]¶

基类: airflow.hooks.base.BaseHook

与 Sagemaker Unified Studio 工作流交互。

此 Hook 封装了 Sagemaker 工作流 Notebook 执行 API。

示例

from airflow.providers.amazon.aws.hooks.sagemaker_unified_studio import SageMakerNotebookHook

notebook_hook = SageMakerNotebookHook(
    input_config={"input_path": "path/to/notebook.ipynb", "input_params": {"param1": "value1"}},
    output_config={"output_uri": "folder/output/location/prefix", "output_formats": "NOTEBOOK"},
    execution_name="notebook_execution",
    waiter_delay=10,
    waiter_max_attempts=1440,
)

参数:

execution_name (str) – 要执行的 notebook 作业名称，与 task_id 相同。
input_config (dict | None) – 输入文件的配置。示例：{'input_path': 'folder/input/notebook.ipynb', 'input_params': {'param1': 'value1'}}
output_config (dict | None) – 输出格式的配置。应包含一个 output_formats 参数来指定输出格式。示例：{'output_formats': ['NOTEBOOK']}
compute (dict | None) – 用于 notebook 执行的计算配置。如果在远程计算上执行，这是必需属性。示例：{ "instance_type": "ml.m5.large", "volume_size_in_gb": 30, "volume_kms_key_id": "", "image_uri": "string", "container_entrypoint": [ "string" ]}
termination_condition (dict | None) – 用于匹配以终止远程执行的条件。示例：{ "MaxRuntimeInSeconds": 3600 }
tags (dict | None) – 与远程执行运行关联的标签。示例：{ "md_analytics": "logs" }
waiter_delay (int) – 检查任务执行状态的时间间隔（秒）。
waiter_max_attempts (int) – 在返回 FAILED 之前等待的尝试次数。

execution_name[源]¶

input_config[源]¶

output_config[源]¶

compute = None[源]¶

termination_condition[源]¶

tags[源]¶

waiter_delay = 10[源]¶

waiter_max_attempts = 1440[源]¶

start_notebook_execution()[源]¶

wait_for_execution_completion(execution_id, context)[源]¶