Press "Enter" to skip to content

使用API计划Amazon SageMaker笔记本任务并管理多步骤笔记本工作流程

Amazon SageMaker Studio为数据科学家提供了完全托管的解决方案,可交互地构建、训练和部署机器学习(ML)模型。Amazon SageMaker笔记本任务允许数据科学家在SageMaker Studio中随需应变地运行他们的笔记本,或按计划定时运行。通过这次发布,您可以使用Amazon SageMaker Pipelines提供的API以编程方式运行笔记本作业,这是Amazon SageMaker的ML工作流编排功能。此外,您还可以使用这些API创建具有多个依赖笔记本的多步骤ML工作流。

SageMaker Pipelines是一个原生工作流编排工具,用于构建利用SageMaker直接集成的ML管道。每个SageMaker管道由步骤组成,其中对应于使用Amazon EMR进行处理、训练或数据处理的单个任务。SageMaker笔记本任务现在作为SageMaker管道中的内置步骤类型即可使用。您可以使用此笔记本任务步骤,仅需几行代码即可轻松运行笔记本作业,使用Amazon SageMaker Python SDK。此外,您还可以将多个依赖笔记本串联起来,以形成有向无环图(DAG)的工作流。然后,您可以运行这些笔记本作业或DAG,并使用SageMaker Studio进行管理和可视化。

数据科学家目前使用SageMaker Studio进行交互式开发Jupyter笔记本,然后使用SageMaker笔记本任务定期运行这些笔记本。这些作业可以立即运行或按照计划在定期时间表上运行,无需数据工作人员将代码重构为Python模块。执行此操作的一些常见用例包括:

  • 在后台运行长时间运行的笔记本
  • 定期运行模型推理以生成报告
  • 从准备小样本数据集扩展到操作PB级大数据
  • 按照一定的频率重新训练和部署模型
  • 为模型质量或数据漂移监控安排作业
  • 探索参数空间以获得更好的模型

尽管此功能使得数据工作人员能够轻松自动化独立的笔记本,但是ML工作流往往由几个笔记本组成,每个笔记本执行特定任务并具有复杂的依赖关系。例如,用于监控模型数据漂移的笔记本应具有允许提取、转换和加载(ETL)以及处理新数据的预处理步骤,并且如果注意到重大漂移,则应具有模型刷新和训练的后处理步骤。此外,数据科学家可能希望根据新数据定期计划触发整个工作流程,以更新模型。为了使您能够轻松自动化笔记本并创建此类复杂工作流程,SageMaker笔记本任务现在可用作SageMaker Pipelines中的一个步骤。在本文中,我们将展示如何使用几行代码解决以下用例:

  • 立即或按计划以编程方式运行独立笔记本
  • 将笔记本的多步骤工作流程作为DAGs创建用于持续集成和持续交付(CI/CD)目的,可通过SageMaker Studio UI进行管理

解决方案概览

下图展示了我们的解决方案架构。您可以使用SageMaker Python SDK来运行单个笔记本作业或工作流。这个功能会创建一个SageMaker训练作业来运行笔记本。

使用API计划Amazon SageMaker笔记本任务并管理多步骤笔记本工作流程 四海 第1张

在以下部分中,我们以一个示例机器学习用例为例,展示创建笔记本作业工作流的步骤,通过不同笔记本步骤之间传递参数,调度工作流程,并通过SageMaker Studio进行监控。

在这个示例中,我们的机器学习问题是构建情感分析模型,这是一种文本分类任务。情感分析的最常见应用包括社交媒体监控、客户支持管理和分析客户反馈。本示例中使用的数据集是斯坦福情感树库(SST2)数据集,包含电影评论以及一个整数(0或1),表示评论的正面或负面情感。

以下是与SST2数据集对应的data.csv文件的示例,显示了其前两列的值。请注意,文件不应包含任何标题。

列1 列2
0 隐藏父母单位的新分泌物
0 只含有费力的笑话,没有机智
1 它热爱它的角色,并传达了有关人类本性的某种美好
0 始终满意于保持不变
0 关于最糟糕的书呆子复仇陈词滥调,制片人可以找出的
0 那太悲惨了,不值得如此肤浅的对待
1 表明像《爱国者游戏》这样的好莱坞大片的导演仍然可以制作出一部带有情感冲击的小制作电影。

在这个机器学习示例中,我们需要执行以下几个任务:

  1. 进行特征工程,将数据集准备成模型可以理解的格式。
  2. 在特征工程之后,运行使用Transformers的训练步骤。
  3. 设置批量推断,使用微调模型来预测新评论的情感。
  4. 设置数据监控步骤,以便定期监视我们的新数据是否有任何质量漂移,可能需要重新训练模型权重。

通过将笔记本作业作为SageMaker管道中的一个步骤来启动,我们可以编排这个由三个独立步骤组成的工作流。工作流的每个步骤都在不同的笔记本中开发,然后转换为独立的笔记本作业步骤,并连接为一个流水线:

  • 预处理 – 从Amazon Simple Storage Service(Amazon S3)下载公共SST2数据集,并为第2步的笔记本创建CSV文件。SST2数据集是一个文本分类数据集,具有两个标签(0和1)和一列用于分类的文本。
  • 训练 – 使用BERT进行文本分类的微调,利用Transformers库对形成的CSV文件进行微调。我们在这个步骤中使用了一个测试数据准备笔记本作为依赖项,它是微调和批量推断步骤的一个依赖项。当微调完成后,使用run magic运行该笔记本,并为带有微调模型的示例推断准备测试数据集。
  • 转换和监控 – 执行批量推断,并设置数据质量与模型监控,以获得基准数据集建议。

使用API计划Amazon SageMaker笔记本任务并管理多步骤笔记本工作流程 四海 第2张

运行笔记本

该解决方案的示例代码可在GitHub上获取。

创建SageMaker笔记本作业步骤与创建其他SageMaker Pipeline步骤类似。在这个笔记本示例中,我们使用SageMaker Python SDK来编排工作流程。要在SageMaker Pipelines中创建笔记本步骤,您可以定义以下参数:

  • 输入笔记本 – 此笔记本步骤将进行编排的笔记本的名称。您可以在此处传递输入笔记本的本地路径。如果此笔记本正在运行其他笔记本,您可以在笔记本作业步骤的AdditionalDependencies参数中传递这些笔记本。
  • 镜像 URI – 笔记本作业步骤背后的Docker镜像。这可以是SageMaker已经提供的预定义镜像,也可以是您定义并推送到Amazon弹性容器注册表(Amazon ECR)的自定义镜像。有关受支持的镜像,请参阅本文末尾的注意事项部分。
  • 内核名称 – 您在SageMaker Studio上使用的内核的名称。此内核规范在您提供的镜像中注册。
  • 实例类型(可选) – 笔记本作业的Amazon弹性计算云(Amazon EC2)实例类型。
  • 参数(可选) – 要传递的参数,这些参数可以在您的笔记本中访问。这些可以定义为键值对。此外,这些参数可以在各种笔记本作业运行或流水线运行之间进行修改。

我们的示例共有五个笔记本:

  • nb-job-pipeline.ipynb – 这是我们定义流水线和工作流程的主要笔记本。
  • preprocess.ipynb – 这个笔记本是我们工作流程的第一步,包含将下载公共AWS数据集并将其创建为CSV文件的代码。
  • training.ipynb – 这个笔记本是我们工作流程的第二步,包含用于对先前步骤中的CSV进行本地训练和微调的代码。此步骤还具有从prepare-test-set.ipynb笔记本中拉取用于使用Fine-tuned模型进行样本推理的测试数据集的依赖关系。
  • prepare-test-set.ipynb – 这个笔记本创建了一个测试数据集,我们的训练笔记本将在第二个流水线步骤中使用并用于使用Fine-tuned模型进行样本推理。
  • transform-monitor.ipynb – 这个笔记本是我们工作流程的第三步,它使用基础BERT模型运行SageMaker批量转换作业,同时设置模型监控的数据质量。

接下来,我们介绍主笔记本nb-job-pipeline.ipynb,它将所有子笔记本组合成一个流水线,并运行端到端的工作流程。请注意,尽管以下示例只运行一次笔记本,但您也可以调度流水线以重复运行笔记本。有关详细说明,请参阅SageMaker文档

对于我们的第一个笔记本作业步骤,我们通过一个默认S3存储桶传递一个参数。我们可以使用此存储桶来存储我们希望在其他流水线步骤中使用的任何工件。对于第一个笔记本(preprocess.ipynb),我们下载AWS公共SST2训练数据集,并将其创建为训练CSV文件,并将其推送到此S3存储桶。请参阅以下代码:

# Parametersprint(default_s3_bucket)!aws s3 cp s3://sagemaker-sample-files/datasets/text/SST2/sst2.train sst2.train# 将只读取前500行,以加快执行速度with open('sst2.train', 'r') as f:    lines = f.readlines()[:500] data = []for line in lines:    label, text = line.strip().split(' ', 1)    data.append((int(label), text))df = pd.DataFrame(data, columns=['label', 'text'])df.to_csv("train.csv", index=False) # 使用较小的数据集创建csv文件!aws s3 cp "train.csv" {default_s3_bucket}

我们可以将这个笔记本转换为一个NotebookJobStep,在主笔记本中使用以下代码实现:

# 提供S3存储桶以转储多媒体处理工件innb_job_params = {"default_s3_bucket": notebook_artifacts}preprocess_nb_step = NotebookJobStep(name=preprocess_step_name,description=preprocess_description,notebook_job_name=preprocess_job_name,image_uri=image_uri,kernel_name=kernel_name,display_name=display_name,role=role,input_notebook=preprocess_notebook,instance_type="ml.m5.4xlarge",parameters=nb_job_params,)

现在我们有了一个样本CSV文件,我们可以在培训笔记本中开始训练我们的模型。我们的训练笔记本使用与S3存储桶相同的参数,并从该位置下载培训数据集。然后,我们使用以下代码片段使用Transformers训练器对象进行微调:

from transformers import TrainingArguments, Trainertraining_args = TrainingArguments(output_dir="test_trainer", evaluation_strategy="epoch")trainer = Trainer(    model=model,    args=training_args,    train_dataset=small_train_dataset,    eval_dataset=small_eval_dataset,    compute_metrics=compute_metrics,)trainer.train()

在微调后,我们想要运行一些批量推断,以查看模型的表现如何。这是使用一个独立的笔记本(prepare-test-set.ipynb)在相同的本地路径下完成的,该笔记本创建一个用我们训练过的模型进行推断的测试数据集。我们可以在培训笔记本中使用以下魔术单元运行附加的笔记本:

%run 'prepare-test-set.ipynb'

我们在第二个笔记本作业步骤中的AdditionalDependencies参数中定义了这个额外的笔记本依赖项:

train_nb_step = NotebookJobStep(name=training_step_name,description=training_description,notebook_job_name=training_job_name,input_notebook=training_notebook,additional_dependencies=[test_data_prep_notebook],image_uri=image_uri,kernel_name=kernel_name,display_name=display_name,instance_type="ml.m5.12xlarge",role=role,parameters=nb_job_params,)

我们还必须指定训练笔记本作业步骤(步骤2)依赖于预处理笔记本作业步骤(步骤1),通过使用add_depends_on API调用,如下所示:

train_nb_step.add_depends_on([preprocess_nb_step])

我们的最后一步是拿到BERT模型并运行SageMaker批量转换,同时使用SageMaker模型监控设置数据捕获和质量。请注意,这与通过管道使用内置的TransformCapture步骤不同。我们的这一步骤将执行相同的API,但会作为笔记本作业步骤进行跟踪。该步骤依赖于我们之前定义的训练作业步骤,因此我们也使用depends_on标志进行捕获。

batch_monitor_step = NotebookJobStep(name=batch_monitor_step_name,description=batch_monitor_description,notebook_job_name=batch_monitor_job_name,input_notebook=batch_monitor_notebook,image_uri=image_uri,kernel_name=kernel_name,display_name=display_name,instance_type="ml.m5.12xlarge",role=role,parameters=nb_job_params,)batch_monitor_step.add_depends_on([train_nb_step])

在定义了工作流程的各个步骤之后,我们可以创建并运行端到端的管道:

# 创建管道pipeline = Pipeline(name=pipeline_name,steps=[preprocess_nb_step, train_nb_step, batch_monitor_step],)# 执行管道pipeline.create(session.get_execution_role())execution = pipeline.start(parameters={})execution.wait(delay=30, max_attempts=60)execution_steps = execution.list_steps()print(execution_steps)

监控管道运行

您可以通过SageMaker Pipelines DAG跟踪和监控笔记本步骤运行情况,如下面的截图所示。

使用API计划Amazon SageMaker笔记本任务并管理多步骤笔记本工作流程 四海 第3张

您还可以通过SageMaker Studio UI在笔记本作业仪表板上选择性地监视已创建的输出文件。在SageMaker Studio之外使用此功能时,您可以使用标签定义可以在笔记本作业仪表板上跟踪运行状态的用户。有关要包含的标签的更多详细信息,请参见在Studio UI仪表板中查看笔记本作业并下载输出

使用API计划Amazon SageMaker笔记本任务并管理多步骤笔记本工作流程 四海 第4张

对于此示例,我们将生成的笔记本作业输出到本地路径中名为outputs的目录中,与您的流水线运行代码一起。如下图所示,您可以在此处看到输入笔记本的输出以及为该步骤定义的任何参数。

使用API计划Amazon SageMaker笔记本任务并管理多步骤笔记本工作流程 四海 第5张

清理

如果您按照我们的示例进行操作,请务必删除创建的流水线、笔记本作业和示例笔记本下载的s3数据。

注意事项

此功能的一些重要注意事项如下:

结论

通过此发布,数据工作者现在可以使用几行代码通过SageMaker Python SDK以编程方式运行他们的笔记本。此外,您可以使用笔记本创建复杂的多步骤工作流,大大减少从笔记本到CI/CD流水线的转换所需的时间。创建流水线后,您可以使用SageMaker Studio查看和运行流程图,管理和比较运行结果。无论您是在安排端到端的ML工作流程还是其中一部分,我们鼓励您尝试基于笔记本的工作流程

Leave a Reply

Your email address will not be published. Required fields are marked *