Browse Source
* add sagemaker pipeline task plugin [DOC] add Sagemaker task plugin document change license and fix dependencies * [fix] Optimize the code * Update dolphinscheduler-task-plugin/dolphinscheduler-task-sagemaker/src/main/java/org/apache/dolphinscheduler/plugin/task/sagemaker/SagemakerParameters.java Optimization parameter judgment Co-authored-by: caishunfeng <caishunfeng2021@gmail.com> * [fix] nips Co-authored-by: caishunfeng <caishunfeng2021@gmail.com>3.1.0-release
JieguangZhou
2 years ago
committed by
GitHub
30 changed files with 1159 additions and 1 deletions
@ -0,0 +1,64 @@ |
|||||||
|
# SageMaker Node |
||||||
|
|
||||||
|
## Overview |
||||||
|
|
||||||
|
[Amazon SageMaker](https://docs.aws.amazon.com/sagemaker/index.html) is a fully managed machine learning service. With Amazon SageMaker, data scientists and developers can quickly build and train machine learning models, and then deploy them into a production-ready hosted environment. |
||||||
|
|
||||||
|
[Amazon SageMaker Model Building Pipelines](https://docs.aws.amazon.com/sagemaker/latest/dg/pipelines.html) is a tool for building machine learning pipelines that take advantage of direct SageMaker integration. |
||||||
|
|
||||||
|
|
||||||
|
For users using big data and machine learning, SageMaker task plugin help users connect big data workflows with SageMaker usage scenarios. |
||||||
|
|
||||||
|
DolphinScheduler SageMaker task plugin features are as follows: |
||||||
|
|
||||||
|
- Start a SageMaker pipeline execution. Continuously get the execution status until the pipeline completes execution. |
||||||
|
|
||||||
|
## Create Task |
||||||
|
|
||||||
|
- Click `Project -> Management-Project -> Name-Workflow Definition`, and click the "Create Workflow" button to enter the |
||||||
|
DAG editing page. |
||||||
|
- Drag from the toolbar <img src="../../../../img/tasks/icons/sagemaker.png" width="15"/> task node to canvas. |
||||||
|
|
||||||
|
## Task Example |
||||||
|
|
||||||
|
First, introduce some general parameters of DolphinScheduler: |
||||||
|
|
||||||
|
- **Node name**: The node name in a workflow definition is unique. |
||||||
|
- **Run flag**: Identifies whether this node schedules normally, if it does not need to execute, select |
||||||
|
the `prohibition execution`. |
||||||
|
- **Descriptive information**: Describe the function of the node. |
||||||
|
- **Task priority**: When the number of worker threads is insufficient, execute in the order of priority from high |
||||||
|
to low, and tasks with the same priority will execute in a first-in first-out order. |
||||||
|
- **Worker grouping**: Assign tasks to the machines of the worker group to execute. If `Default` is selected, |
||||||
|
randomly select a worker machine for execution. |
||||||
|
- **Environment Name**: Configure the environment name in which run the script. |
||||||
|
- **Times of failed retry attempts**: The number of times the task failed to resubmit. |
||||||
|
- **Failed retry interval**: The time interval (unit minute) for resubmitting the task after a failed task. |
||||||
|
- **Delayed execution time**: The time (unit minute) that a task delays in execution. |
||||||
|
- **Timeout alarm**: Check the timeout alarm and timeout failure. When the task runs exceed the "timeout", an alarm |
||||||
|
email will send and the task execution will fail. |
||||||
|
- **Predecessor task**: Selecting a predecessor task for the current task, will set the selected predecessor task as |
||||||
|
upstream of the current task. |
||||||
|
|
||||||
|
Here are some specific parameters for the SagaMaker plugin: |
||||||
|
|
||||||
|
- **SagemakerRequestJson**: Request parameters of StartPipelineExecution,see also [AWS API](https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_StartPipelineExecution.html) |
||||||
|
|
||||||
|
|
||||||
|
The task plugin are shown as follows: |
||||||
|
|
||||||
|
![sagemaker_pipeline](../../../../img/tasks/demo/sagemaker_pipeline.png) |
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
## Environment to prepare |
||||||
|
|
||||||
|
Some AWS configuration is required, modify a field in file `common.properties` |
||||||
|
```yaml |
||||||
|
# The AWS access key. if resource.storage.type=S3 or use EMR-Task, This configuration is required |
||||||
|
resource.aws.access.key.id=<YOUR AWS ACCESS KEY> |
||||||
|
# The AWS secret access key. if resource.storage.type=S3 or use EMR-Task, This configuration is required |
||||||
|
resource.aws.secret.access.key=<YOUR AWS SECRET KEY> |
||||||
|
# The AWS Region to use. if resource.storage.type=S3 or use EMR-Task, This configuration is required |
||||||
|
resource.aws.region=<AWS REGION> |
||||||
|
``` |
@ -0,0 +1,57 @@ |
|||||||
|
# SageMaker 节点 |
||||||
|
|
||||||
|
## 综述 |
||||||
|
|
||||||
|
[Amazon SageMaker](https://aws.amazon.com/cn/pm/sagemaker) 是一个云机器学习平台。 提供了完整的基础设施,工具和工作流来帮助用户可以创建、训练和发布机器学习模型。 |
||||||
|
|
||||||
|
[Amazon SageMaker Model Building Pipelines](https://docs.aws.amazon.com/sagemaker/latest/dg/pipelines.html) 是一个可以直接使用SageMaker各种集成的机器学习管道构建工具,用户可以使用使用 Amazon SageMaker Pipeline 来构建端到端的机器学习系统。 |
||||||
|
|
||||||
|
对于使用大数据与人工智能的用户,SageMaker 任务组件帮助用户可以串联起大数据工作流与SagaMaker的使用场景。 |
||||||
|
|
||||||
|
DolphinScheduler SageMaker 组件的功能: |
||||||
|
- 启动 SageMaker Pipeline Execution,并持续获取状态,直至Pipeline执行完成。 |
||||||
|
|
||||||
|
## 创建任务 |
||||||
|
|
||||||
|
- 点击项目管理-项目名称-工作流定义,点击“创建工作流”按钮,进入 DAG 编辑页面; |
||||||
|
- 拖动工具栏的 <img src="../../../../img/tasks/icons/sagemaker.png" width="15"/> 任务节点到画板中。 |
||||||
|
|
||||||
|
|
||||||
|
## 任务样例 |
||||||
|
|
||||||
|
首先介绍一些DS通用参数 |
||||||
|
|
||||||
|
- **节点名称** :设置任务的名称。一个工作流定义中的节点名称是唯一的。 |
||||||
|
- **运行标志** :标识这个节点是否能正常调度,如果不需要执行,可以打开禁止执行开关。 |
||||||
|
- **描述** :描述该节点的功能。 |
||||||
|
- **任务优先级** :worker 线程数不足时,根据优先级从高到低依次执行,优先级一样时根据先进先出原则执行。 |
||||||
|
- **Worker 分组** :任务分配给 worker 组的机器执行,选择 Default,会随机选择一台 worker 机执行。 |
||||||
|
- **环境名称** :配置运行脚本的环境。 |
||||||
|
- **失败重试次数** :任务失败重新提交的次数。 |
||||||
|
- **失败重试间隔** :任务失败重新提交任务的时间间隔,以分钟为单位。 |
||||||
|
- **延迟执行时间** :任务延迟执行的时间,以分钟为单位。 |
||||||
|
- **超时告警** :勾选超时告警、超时失败,当任务超过"超时时长"后,会发送告警邮件并且任务执行失败。 |
||||||
|
- **前置任务** :选择当前任务的前置任务,会将被选择的前置任务设置为当前任务的上游。 |
||||||
|
|
||||||
|
以上参数如无特殊需求,可以默认即可 |
||||||
|
|
||||||
|
- **SagemakerRequestJson**: 启动SageMakerPipeline的需要的请求参数,可见 [AWS API](https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_StartPipelineExecution.html) |
||||||
|
|
||||||
|
|
||||||
|
组件图示如下: |
||||||
|
|
||||||
|
![sagemaker_pipeline](../../../../img/tasks/demo/sagemaker_pipeline.png) |
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
## 环境配置 |
||||||
|
|
||||||
|
需要进行AWS的一些配置,修改`common.properties`中的`xxxxx`为你的配置信息 |
||||||
|
```yaml |
||||||
|
# The AWS access key. if resource.storage.type=S3 or use EMR-Task, This configuration is required |
||||||
|
resource.aws.access.key.id=<YOUR AWS ACCESS KEY> |
||||||
|
# The AWS secret access key. if resource.storage.type=S3 or use EMR-Task, This configuration is required |
||||||
|
resource.aws.secret.access.key=<YOUR AWS SECRET KEY> |
||||||
|
# The AWS Region to use. if resource.storage.type=S3 or use EMR-Task, This configuration is required |
||||||
|
resource.aws.region=<AWS REGION> |
||||||
|
``` |
After Width: | Height: | Size: 124 KiB |
After Width: | Height: | Size: 25 KiB |
@ -0,0 +1,201 @@ |
|||||||
|
Apache License |
||||||
|
Version 2.0, January 2004 |
||||||
|
http://www.apache.org/licenses/ |
||||||
|
|
||||||
|
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION |
||||||
|
|
||||||
|
1. Definitions. |
||||||
|
|
||||||
|
"License" shall mean the terms and conditions for use, reproduction, |
||||||
|
and distribution as defined by Sections 1 through 9 of this document. |
||||||
|
|
||||||
|
"Licensor" shall mean the copyright owner or entity authorized by |
||||||
|
the copyright owner that is granting the License. |
||||||
|
|
||||||
|
"Legal Entity" shall mean the union of the acting entity and all |
||||||
|
other entities that control, are controlled by, or are under common |
||||||
|
control with that entity. For the purposes of this definition, |
||||||
|
"control" means (i) the power, direct or indirect, to cause the |
||||||
|
direction or management of such entity, whether by contract or |
||||||
|
otherwise, or (ii) ownership of fifty percent (50%) or more of the |
||||||
|
outstanding shares, or (iii) beneficial ownership of such entity. |
||||||
|
|
||||||
|
"You" (or "Your") shall mean an individual or Legal Entity |
||||||
|
exercising permissions granted by this License. |
||||||
|
|
||||||
|
"Source" form shall mean the preferred form for making modifications, |
||||||
|
including but not limited to software source code, documentation |
||||||
|
source, and configuration files. |
||||||
|
|
||||||
|
"Object" form shall mean any form resulting from mechanical |
||||||
|
transformation or translation of a Source form, including but |
||||||
|
not limited to compiled object code, generated documentation, |
||||||
|
and conversions to other media types. |
||||||
|
|
||||||
|
"Work" shall mean the work of authorship, whether in Source or |
||||||
|
Object form, made available under the License, as indicated by a |
||||||
|
copyright notice that is included in or attached to the work |
||||||
|
(an example is provided in the Appendix below). |
||||||
|
|
||||||
|
"Derivative Works" shall mean any work, whether in Source or Object |
||||||
|
form, that is based on (or derived from) the Work and for which the |
||||||
|
editorial revisions, annotations, elaborations, or other modifications |
||||||
|
represent, as a whole, an original work of authorship. For the purposes |
||||||
|
of this License, Derivative Works shall not include works that remain |
||||||
|
separable from, or merely link (or bind by name) to the interfaces of, |
||||||
|
the Work and Derivative Works thereof. |
||||||
|
|
||||||
|
"Contribution" shall mean any work of authorship, including |
||||||
|
the original version of the Work and any modifications or additions |
||||||
|
to that Work or Derivative Works thereof, that is intentionally |
||||||
|
submitted to Licensor for inclusion in the Work by the copyright owner |
||||||
|
or by an individual or Legal Entity authorized to submit on behalf of |
||||||
|
the copyright owner. For the purposes of this definition, "submitted" |
||||||
|
means any form of electronic, verbal, or written communication sent |
||||||
|
to the Licensor or its representatives, including but not limited to |
||||||
|
communication on electronic mailing lists, source code control systems, |
||||||
|
and issue tracking systems that are managed by, or on behalf of, the |
||||||
|
Licensor for the purpose of discussing and improving the Work, but |
||||||
|
excluding communication that is conspicuously marked or otherwise |
||||||
|
designated in writing by the copyright owner as "Not a Contribution." |
||||||
|
|
||||||
|
"Contributor" shall mean Licensor and any individual or Legal Entity |
||||||
|
on behalf of whom a Contribution has been received by Licensor and |
||||||
|
subsequently incorporated within the Work. |
||||||
|
|
||||||
|
2. Grant of Copyright License. Subject to the terms and conditions of |
||||||
|
this License, each Contributor hereby grants to You a perpetual, |
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable |
||||||
|
copyright license to reproduce, prepare Derivative Works of, |
||||||
|
publicly display, publicly perform, sublicense, and distribute the |
||||||
|
Work and such Derivative Works in Source or Object form. |
||||||
|
|
||||||
|
3. Grant of Patent License. Subject to the terms and conditions of |
||||||
|
this License, each Contributor hereby grants to You a perpetual, |
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable |
||||||
|
(except as stated in this section) patent license to make, have made, |
||||||
|
use, offer to sell, sell, import, and otherwise transfer the Work, |
||||||
|
where such license applies only to those patent claims licensable |
||||||
|
by such Contributor that are necessarily infringed by their |
||||||
|
Contribution(s) alone or by combination of their Contribution(s) |
||||||
|
with the Work to which such Contribution(s) was submitted. If You |
||||||
|
institute patent litigation against any entity (including a |
||||||
|
cross-claim or counterclaim in a lawsuit) alleging that the Work |
||||||
|
or a Contribution incorporated within the Work constitutes direct |
||||||
|
or contributory patent infringement, then any patent licenses |
||||||
|
granted to You under this License for that Work shall terminate |
||||||
|
as of the date such litigation is filed. |
||||||
|
|
||||||
|
4. Redistribution. You may reproduce and distribute copies of the |
||||||
|
Work or Derivative Works thereof in any medium, with or without |
||||||
|
modifications, and in Source or Object form, provided that You |
||||||
|
meet the following conditions: |
||||||
|
|
||||||
|
(a) You must give any other recipients of the Work or |
||||||
|
Derivative Works a copy of this License; and |
||||||
|
|
||||||
|
(b) You must cause any modified files to carry prominent notices |
||||||
|
stating that You changed the files; and |
||||||
|
|
||||||
|
(c) You must retain, in the Source form of any Derivative Works |
||||||
|
that You distribute, all copyright, patent, trademark, and |
||||||
|
attribution notices from the Source form of the Work, |
||||||
|
excluding those notices that do not pertain to any part of |
||||||
|
the Derivative Works; and |
||||||
|
|
||||||
|
(d) If the Work includes a "NOTICE" text file as part of its |
||||||
|
distribution, then any Derivative Works that You distribute must |
||||||
|
include a readable copy of the attribution notices contained |
||||||
|
within such NOTICE file, excluding those notices that do not |
||||||
|
pertain to any part of the Derivative Works, in at least one |
||||||
|
of the following places: within a NOTICE text file distributed |
||||||
|
as part of the Derivative Works; within the Source form or |
||||||
|
documentation, if provided along with the Derivative Works; or, |
||||||
|
within a display generated by the Derivative Works, if and |
||||||
|
wherever such third-party notices normally appear. The contents |
||||||
|
of the NOTICE file are for informational purposes only and |
||||||
|
do not modify the License. You may add Your own attribution |
||||||
|
notices within Derivative Works that You distribute, alongside |
||||||
|
or as an addendum to the NOTICE text from the Work, provided |
||||||
|
that such additional attribution notices cannot be construed |
||||||
|
as modifying the License. |
||||||
|
|
||||||
|
You may add Your own copyright statement to Your modifications and |
||||||
|
may provide additional or different license terms and conditions |
||||||
|
for use, reproduction, or distribution of Your modifications, or |
||||||
|
for any such Derivative Works as a whole, provided Your use, |
||||||
|
reproduction, and distribution of the Work otherwise complies with |
||||||
|
the conditions stated in this License. |
||||||
|
|
||||||
|
5. Submission of Contributions. Unless You explicitly state otherwise, |
||||||
|
any Contribution intentionally submitted for inclusion in the Work |
||||||
|
by You to the Licensor shall be under the terms and conditions of |
||||||
|
this License, without any additional terms or conditions. |
||||||
|
Notwithstanding the above, nothing herein shall supersede or modify |
||||||
|
the terms of any separate license agreement you may have executed |
||||||
|
with Licensor regarding such Contributions. |
||||||
|
|
||||||
|
6. Trademarks. This License does not grant permission to use the trade |
||||||
|
names, trademarks, service marks, or product names of the Licensor, |
||||||
|
except as required for reasonable and customary use in describing the |
||||||
|
origin of the Work and reproducing the content of the NOTICE file. |
||||||
|
|
||||||
|
7. Disclaimer of Warranty. Unless required by applicable law or |
||||||
|
agreed to in writing, Licensor provides the Work (and each |
||||||
|
Contributor provides its Contributions) on an "AS IS" BASIS, |
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or |
||||||
|
implied, including, without limitation, any warranties or conditions |
||||||
|
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A |
||||||
|
PARTICULAR PURPOSE. You are solely responsible for determining the |
||||||
|
appropriateness of using or redistributing the Work and assume any |
||||||
|
risks associated with Your exercise of permissions under this License. |
||||||
|
|
||||||
|
8. Limitation of Liability. In no event and under no legal theory, |
||||||
|
whether in tort (including negligence), contract, or otherwise, |
||||||
|
unless required by applicable law (such as deliberate and grossly |
||||||
|
negligent acts) or agreed to in writing, shall any Contributor be |
||||||
|
liable to You for damages, including any direct, indirect, special, |
||||||
|
incidental, or consequential damages of any character arising as a |
||||||
|
result of this License or out of the use or inability to use the |
||||||
|
Work (including but not limited to damages for loss of goodwill, |
||||||
|
work stoppage, computer failure or malfunction, or any and all |
||||||
|
other commercial damages or losses), even if such Contributor |
||||||
|
has been advised of the possibility of such damages. |
||||||
|
|
||||||
|
9. Accepting Warranty or Additional Liability. While redistributing |
||||||
|
the Work or Derivative Works thereof, You may choose to offer, |
||||||
|
and charge a fee for, acceptance of support, warranty, indemnity, |
||||||
|
or other liability obligations and/or rights consistent with this |
||||||
|
License. However, in accepting such obligations, You may act only |
||||||
|
on Your own behalf and on Your sole responsibility, not on behalf |
||||||
|
of any other Contributor, and only if You agree to indemnify, |
||||||
|
defend, and hold each Contributor harmless for any liability |
||||||
|
incurred by, or claims asserted against, such Contributor by reason |
||||||
|
of your accepting any such warranty or additional liability. |
||||||
|
|
||||||
|
END OF TERMS AND CONDITIONS |
||||||
|
|
||||||
|
APPENDIX: How to apply the Apache License to your work. |
||||||
|
|
||||||
|
To apply the Apache License to your work, attach the following |
||||||
|
boilerplate notice, with the fields enclosed by brackets "[]" |
||||||
|
replaced with your own identifying information. (Don't include |
||||||
|
the brackets!) The text should be enclosed in the appropriate |
||||||
|
comment syntax for the file format. We also recommend that a |
||||||
|
file or class name and description of purpose be included on the |
||||||
|
same "printed page" as the copyright notice for easier |
||||||
|
identification within third-party archives. |
||||||
|
|
||||||
|
Copyright [yyyy] [name of copyright owner] |
||||||
|
|
||||||
|
Licensed under the Apache License, Version 2.0 (the "License"); |
||||||
|
you may not use this file except in compliance 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. |
@ -0,0 +1,54 @@ |
|||||||
|
<?xml version="1.0" encoding="UTF-8"?> |
||||||
|
<!-- |
||||||
|
~ Licensed to the Apache Software Foundation (ASF) under one or more |
||||||
|
~ contributor license agreements. See the NOTICE file distributed with |
||||||
|
~ this work for additional information regarding copyright ownership. |
||||||
|
~ The ASF licenses this file to You under the Apache License, Version 2.0 |
||||||
|
~ (the "License"); you may not use this file except in compliance 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. |
||||||
|
--> |
||||||
|
<project xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" |
||||||
|
xmlns="http://maven.apache.org/POM/4.0.0" |
||||||
|
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd"> |
||||||
|
<parent> |
||||||
|
<artifactId>dolphinscheduler-task-plugin</artifactId> |
||||||
|
<groupId>org.apache.dolphinscheduler</groupId> |
||||||
|
<version>dev-SNAPSHOT</version> |
||||||
|
</parent> |
||||||
|
<modelVersion>4.0.0</modelVersion> |
||||||
|
|
||||||
|
<artifactId>dolphinscheduler-task-sagemaker</artifactId> |
||||||
|
<packaging>jar</packaging> |
||||||
|
|
||||||
|
<dependencies> |
||||||
|
<dependency> |
||||||
|
<groupId>org.apache.dolphinscheduler</groupId> |
||||||
|
<artifactId>dolphinscheduler-spi</artifactId> |
||||||
|
<scope>provided</scope> |
||||||
|
</dependency> |
||||||
|
<dependency> |
||||||
|
<groupId>org.apache.dolphinscheduler</groupId> |
||||||
|
<artifactId>dolphinscheduler-task-api</artifactId> |
||||||
|
<scope>provided</scope> |
||||||
|
</dependency> |
||||||
|
<dependency> |
||||||
|
<groupId>org.apache.dolphinscheduler</groupId> |
||||||
|
<artifactId>dolphinscheduler-common</artifactId> |
||||||
|
<scope>provided</scope> |
||||||
|
</dependency> |
||||||
|
<dependency> |
||||||
|
<groupId>com.amazonaws</groupId> |
||||||
|
<artifactId>aws-java-sdk-sagemaker</artifactId> |
||||||
|
<version>1.12.160</version> |
||||||
|
</dependency> |
||||||
|
|
||||||
|
</dependencies> |
||||||
|
</project> |
@ -0,0 +1,125 @@ |
|||||||
|
/* |
||||||
|
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||||
|
* contributor license agreements. See the NOTICE file distributed with |
||||||
|
* this work for additional information regarding copyright ownership. |
||||||
|
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||||
|
* (the "License"); you may not use this file except in compliance 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. |
||||||
|
*/ |
||||||
|
|
||||||
|
package org.apache.dolphinscheduler.plugin.task.sagemaker; |
||||||
|
|
||||||
|
import org.apache.dolphinscheduler.common.thread.ThreadUtils; |
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.TaskConstants; |
||||||
|
|
||||||
|
import java.util.Collections; |
||||||
|
import java.util.List; |
||||||
|
|
||||||
|
import org.slf4j.Logger; |
||||||
|
import org.slf4j.LoggerFactory; |
||||||
|
|
||||||
|
import com.amazonaws.services.sagemaker.AmazonSageMaker; |
||||||
|
import com.amazonaws.services.sagemaker.model.DescribePipelineExecutionRequest; |
||||||
|
import com.amazonaws.services.sagemaker.model.DescribePipelineExecutionResult; |
||||||
|
import com.amazonaws.services.sagemaker.model.ListPipelineExecutionStepsRequest; |
||||||
|
import com.amazonaws.services.sagemaker.model.ListPipelineExecutionStepsResult; |
||||||
|
import com.amazonaws.services.sagemaker.model.PipelineExecutionStep; |
||||||
|
import com.amazonaws.services.sagemaker.model.StartPipelineExecutionRequest; |
||||||
|
import com.amazonaws.services.sagemaker.model.StartPipelineExecutionResult; |
||||||
|
import com.amazonaws.services.sagemaker.model.StopPipelineExecutionRequest; |
||||||
|
import com.amazonaws.services.sagemaker.model.StopPipelineExecutionResult; |
||||||
|
|
||||||
|
public class PipelineUtils { |
||||||
|
|
||||||
|
|
||||||
|
protected final Logger logger = LoggerFactory.getLogger(String.format(TaskConstants.TASK_LOG_LOGGER_NAME_FORMAT, getClass())); |
||||||
|
private final AmazonSageMaker client; |
||||||
|
private String pipelineExecutionArn; |
||||||
|
private String clientRequestToken; |
||||||
|
private String pipelineStatus; |
||||||
|
|
||||||
|
public PipelineUtils(AmazonSageMaker client) { |
||||||
|
this.client = client; |
||||||
|
} |
||||||
|
|
||||||
|
public int startPipelineExecution(StartPipelineExecutionRequest request) { |
||||||
|
int exitStatusCode = TaskConstants.EXIT_CODE_FAILURE; |
||||||
|
try { |
||||||
|
StartPipelineExecutionResult result = client.startPipelineExecution(request); |
||||||
|
pipelineExecutionArn = result.getPipelineExecutionArn(); |
||||||
|
clientRequestToken = request.getClientRequestToken(); |
||||||
|
exitStatusCode = TaskConstants.EXIT_CODE_SUCCESS; |
||||||
|
logger.info("Start pipeline: {} success", pipelineExecutionArn); |
||||||
|
} catch (Exception e) { |
||||||
|
logger.error("Start pipeline error: {}", e.getMessage()); |
||||||
|
} |
||||||
|
|
||||||
|
return exitStatusCode; |
||||||
|
} |
||||||
|
|
||||||
|
public void stopPipelineExecution() { |
||||||
|
StopPipelineExecutionRequest request = new StopPipelineExecutionRequest(); |
||||||
|
request.setPipelineExecutionArn(pipelineExecutionArn); |
||||||
|
request.setClientRequestToken(clientRequestToken); |
||||||
|
|
||||||
|
try { |
||||||
|
StopPipelineExecutionResult result = client.stopPipelineExecution(request); |
||||||
|
logger.info("Stop pipeline: {} success", result.getPipelineExecutionArn()); |
||||||
|
} catch (Exception e) { |
||||||
|
logger.error("Stop pipeline error: {}", e.getMessage()); |
||||||
|
} |
||||||
|
|
||||||
|
} |
||||||
|
|
||||||
|
public int checkPipelineExecutionStatus() { |
||||||
|
describePipelineExecution(); |
||||||
|
while (pipelineStatus.equals("Executing")) { |
||||||
|
logger.info("check Pipeline Steps running"); |
||||||
|
listPipelineExecutionSteps(); |
||||||
|
ThreadUtils.sleep(SagemakerConstants.CHECK_PIPELINE_EXECUTION_STATUS_INTERVAL); |
||||||
|
describePipelineExecution(); |
||||||
|
} |
||||||
|
|
||||||
|
int exitStatusCode = TaskConstants.EXIT_CODE_FAILURE; |
||||||
|
if (pipelineStatus.equals("Succeeded")) { |
||||||
|
exitStatusCode = TaskConstants.EXIT_CODE_SUCCESS; |
||||||
|
} |
||||||
|
logger.info("exit : {}", exitStatusCode); |
||||||
|
logger.info("PipelineExecutionStatus : {}", pipelineStatus); |
||||||
|
return exitStatusCode; |
||||||
|
} |
||||||
|
|
||||||
|
private void describePipelineExecution() { |
||||||
|
DescribePipelineExecutionRequest request = new DescribePipelineExecutionRequest(); |
||||||
|
request.setPipelineExecutionArn(pipelineExecutionArn); |
||||||
|
DescribePipelineExecutionResult result = client.describePipelineExecution(request); |
||||||
|
pipelineStatus = result.getPipelineExecutionStatus(); |
||||||
|
logger.info("PipelineExecutionStatus: {}", pipelineStatus); |
||||||
|
} |
||||||
|
|
||||||
|
private void listPipelineExecutionSteps() { |
||||||
|
ListPipelineExecutionStepsRequest request = new ListPipelineExecutionStepsRequest(); |
||||||
|
request.setPipelineExecutionArn(pipelineExecutionArn); |
||||||
|
request.setMaxResults(SagemakerConstants.PIPELINE_MAX_RESULTS); |
||||||
|
ListPipelineExecutionStepsResult result = client.listPipelineExecutionSteps(request); |
||||||
|
List<PipelineExecutionStep> steps = result.getPipelineExecutionSteps(); |
||||||
|
Collections.reverse(steps); |
||||||
|
logger.info("pipelineStepsStatus: "); |
||||||
|
for (PipelineExecutionStep step : steps) { |
||||||
|
String stepMessage = step.toString(); |
||||||
|
logger.info(stepMessage); |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
public String getPipelineExecutionArn() { |
||||||
|
return pipelineExecutionArn; |
||||||
|
} |
||||||
|
} |
@ -0,0 +1,27 @@ |
|||||||
|
/* |
||||||
|
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||||
|
* contributor license agreements. See the NOTICE file distributed with |
||||||
|
* this work for additional information regarding copyright ownership. |
||||||
|
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||||
|
* (the "License"); you may not use this file except in compliance 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. |
||||||
|
*/ |
||||||
|
|
||||||
|
package org.apache.dolphinscheduler.plugin.task.sagemaker; |
||||||
|
|
||||||
|
public class SagemakerConstants { |
||||||
|
public static final int CHECK_PIPELINE_EXECUTION_STATUS_INTERVAL = 5000; |
||||||
|
public static final int PIPELINE_MAX_RESULTS = 100; |
||||||
|
|
||||||
|
private SagemakerConstants() { |
||||||
|
throw new IllegalStateException("Utility class"); |
||||||
|
} |
||||||
|
} |
@ -0,0 +1,43 @@ |
|||||||
|
/* |
||||||
|
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||||
|
* contributor license agreements. See the NOTICE file distributed with |
||||||
|
* this work for additional information regarding copyright ownership. |
||||||
|
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||||
|
* (the "License"); you may not use this file except in compliance 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. |
||||||
|
*/ |
||||||
|
|
||||||
|
package org.apache.dolphinscheduler.plugin.task.sagemaker; |
||||||
|
|
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.parameters.AbstractParameters; |
||||||
|
|
||||||
|
import org.apache.commons.lang3.StringUtils; |
||||||
|
|
||||||
|
import lombok.Getter; |
||||||
|
import lombok.Setter; |
||||||
|
import lombok.ToString; |
||||||
|
|
||||||
|
@Getter |
||||||
|
@Setter |
||||||
|
@ToString |
||||||
|
public class SagemakerParameters extends AbstractParameters { |
||||||
|
|
||||||
|
/** |
||||||
|
* request script |
||||||
|
*/ |
||||||
|
private String sagemakerRequestJson; |
||||||
|
|
||||||
|
@Override |
||||||
|
public boolean checkParameters() { |
||||||
|
return StringUtils.isNotEmpty(sagemakerRequestJson); |
||||||
|
} |
||||||
|
|
||||||
|
} |
@ -0,0 +1,161 @@ |
|||||||
|
/* |
||||||
|
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||||
|
* contributor license agreements. See the NOTICE file distributed with |
||||||
|
* this work for additional information regarding copyright ownership. |
||||||
|
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||||
|
* (the "License"); you may not use this file except in compliance 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. |
||||||
|
*/ |
||||||
|
|
||||||
|
package org.apache.dolphinscheduler.plugin.task.sagemaker; |
||||||
|
|
||||||
|
import static com.fasterxml.jackson.databind.DeserializationFeature.ACCEPT_EMPTY_ARRAY_AS_NULL_OBJECT; |
||||||
|
import static com.fasterxml.jackson.databind.DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES; |
||||||
|
import static com.fasterxml.jackson.databind.DeserializationFeature.READ_UNKNOWN_ENUM_VALUES_AS_NULL; |
||||||
|
import static com.fasterxml.jackson.databind.MapperFeature.REQUIRE_SETTERS_FOR_GETTERS; |
||||||
|
|
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.AbstractTaskExecutor; |
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.TaskConstants; |
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.TaskExecutionContext; |
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.model.Property; |
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.parser.ParamUtils; |
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.parser.ParameterUtils; |
||||||
|
import org.apache.dolphinscheduler.spi.utils.JSONUtils; |
||||||
|
import org.apache.dolphinscheduler.spi.utils.PropertyUtils; |
||||||
|
|
||||||
|
import java.util.Map; |
||||||
|
|
||||||
|
import com.amazonaws.auth.AWSCredentialsProvider; |
||||||
|
import com.amazonaws.auth.AWSStaticCredentialsProvider; |
||||||
|
import com.amazonaws.auth.BasicAWSCredentials; |
||||||
|
import com.amazonaws.services.sagemaker.AmazonSageMaker; |
||||||
|
import com.amazonaws.services.sagemaker.AmazonSageMakerClientBuilder; |
||||||
|
import com.amazonaws.services.sagemaker.model.StartPipelineExecutionRequest; |
||||||
|
import com.fasterxml.jackson.databind.ObjectMapper; |
||||||
|
import com.fasterxml.jackson.databind.PropertyNamingStrategy; |
||||||
|
|
||||||
|
/** |
||||||
|
* SagemakerTask task, Used to start Sagemaker pipeline |
||||||
|
*/ |
||||||
|
public class SagemakerTask extends AbstractTaskExecutor { |
||||||
|
|
||||||
|
private static final ObjectMapper objectMapper = |
||||||
|
new ObjectMapper().configure(FAIL_ON_UNKNOWN_PROPERTIES, false).configure(ACCEPT_EMPTY_ARRAY_AS_NULL_OBJECT, true).configure(READ_UNKNOWN_ENUM_VALUES_AS_NULL, true) |
||||||
|
.configure(REQUIRE_SETTERS_FOR_GETTERS, true).setPropertyNamingStrategy(new PropertyNamingStrategy.UpperCamelCaseStrategy()); |
||||||
|
/** |
||||||
|
* taskExecutionContext |
||||||
|
*/ |
||||||
|
private final TaskExecutionContext taskExecutionContext; |
||||||
|
/** |
||||||
|
* SageMaker parameters |
||||||
|
*/ |
||||||
|
private SagemakerParameters parameters; |
||||||
|
private PipelineUtils utils; |
||||||
|
|
||||||
|
public SagemakerTask(TaskExecutionContext taskExecutionContext) { |
||||||
|
super(taskExecutionContext); |
||||||
|
|
||||||
|
this.taskExecutionContext = taskExecutionContext; |
||||||
|
|
||||||
|
} |
||||||
|
|
||||||
|
@Override |
||||||
|
public void init() { |
||||||
|
logger.info("Sagemaker task params {}", taskExecutionContext.getTaskParams()); |
||||||
|
|
||||||
|
parameters = JSONUtils.parseObject(taskExecutionContext.getTaskParams(), SagemakerParameters.class); |
||||||
|
|
||||||
|
if (!parameters.checkParameters()) { |
||||||
|
throw new SagemakerTaskException("Sagemaker task params is not valid"); |
||||||
|
} |
||||||
|
|
||||||
|
} |
||||||
|
|
||||||
|
@Override |
||||||
|
public void handle() throws SagemakerTaskException { |
||||||
|
try { |
||||||
|
int exitStatusCode = handleStartPipeline(); |
||||||
|
setExitStatusCode(exitStatusCode); |
||||||
|
} catch (Exception e) { |
||||||
|
setExitStatusCode(TaskConstants.EXIT_CODE_FAILURE); |
||||||
|
throw new SagemakerTaskException("SageMaker task error", e); |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
@Override |
||||||
|
public void cancelApplication(boolean cancelApplication) { |
||||||
|
// stop pipeline
|
||||||
|
utils.stopPipelineExecution(); |
||||||
|
} |
||||||
|
|
||||||
|
public int handleStartPipeline() { |
||||||
|
int exitStatusCode; |
||||||
|
StartPipelineExecutionRequest request = createStartPipelineRequest(); |
||||||
|
|
||||||
|
try { |
||||||
|
AmazonSageMaker client = createClient(); |
||||||
|
utils = new PipelineUtils(client); |
||||||
|
setAppIds(utils.getPipelineExecutionArn()); |
||||||
|
} catch (Exception e) { |
||||||
|
throw new SagemakerTaskException("can not connect aws ", e); |
||||||
|
} |
||||||
|
|
||||||
|
// Start pipeline
|
||||||
|
exitStatusCode = utils.startPipelineExecution(request); |
||||||
|
if (exitStatusCode == TaskConstants.EXIT_CODE_SUCCESS) { |
||||||
|
// Keep checking the health status
|
||||||
|
exitStatusCode = utils.checkPipelineExecutionStatus(); |
||||||
|
} |
||||||
|
return exitStatusCode; |
||||||
|
} |
||||||
|
|
||||||
|
public StartPipelineExecutionRequest createStartPipelineRequest() throws SagemakerTaskException { |
||||||
|
|
||||||
|
String requestJson = parameters.getSagemakerRequestJson(); |
||||||
|
requestJson = parseRequstJson(requestJson); |
||||||
|
|
||||||
|
StartPipelineExecutionRequest startPipelineRequest; |
||||||
|
try { |
||||||
|
startPipelineRequest = objectMapper.readValue(requestJson, StartPipelineExecutionRequest.class); |
||||||
|
} catch (Exception e) { |
||||||
|
logger.error("can not parse SagemakerRequestJson from json: {}", requestJson); |
||||||
|
throw new SagemakerTaskException("can not parse SagemakerRequestJson ", e); |
||||||
|
} |
||||||
|
|
||||||
|
logger.info("Sagemaker task create StartPipelineRequest: {}", startPipelineRequest); |
||||||
|
return startPipelineRequest; |
||||||
|
} |
||||||
|
|
||||||
|
@Override |
||||||
|
public SagemakerParameters getParameters() { |
||||||
|
return parameters; |
||||||
|
} |
||||||
|
|
||||||
|
private String parseRequstJson(String requestJson) { |
||||||
|
// combining local and global parameters
|
||||||
|
Map<String, Property> paramsMap = taskExecutionContext.getPrepareParamsMap(); |
||||||
|
return ParameterUtils.convertParameterPlaceholders(requestJson, ParamUtils.convert(paramsMap)); |
||||||
|
} |
||||||
|
|
||||||
|
private AmazonSageMaker createClient() { |
||||||
|
final String awsAccessKeyId = PropertyUtils.getString(TaskConstants.AWS_ACCESS_KEY_ID); |
||||||
|
final String awsSecretAccessKey = PropertyUtils.getString(TaskConstants.AWS_SECRET_ACCESS_KEY); |
||||||
|
final String awsRegion = PropertyUtils.getString(TaskConstants.AWS_REGION); |
||||||
|
final BasicAWSCredentials basicAWSCredentials = new BasicAWSCredentials(awsAccessKeyId, awsSecretAccessKey); |
||||||
|
final AWSCredentialsProvider awsCredentialsProvider = new AWSStaticCredentialsProvider(basicAWSCredentials); |
||||||
|
// create a SageMaker client
|
||||||
|
return AmazonSageMakerClientBuilder.standard() |
||||||
|
.withCredentials(awsCredentialsProvider) |
||||||
|
.withRegion(awsRegion) |
||||||
|
.build(); |
||||||
|
} |
||||||
|
|
||||||
|
} |
@ -0,0 +1,49 @@ |
|||||||
|
/* |
||||||
|
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||||
|
* contributor license agreements. See the NOTICE file distributed with |
||||||
|
* this work for additional information regarding copyright ownership. |
||||||
|
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||||
|
* (the "License"); you may not use this file except in compliance 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. |
||||||
|
*/ |
||||||
|
|
||||||
|
package org.apache.dolphinscheduler.plugin.task.sagemaker; |
||||||
|
|
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.TaskChannel; |
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.TaskExecutionContext; |
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.parameters.AbstractParameters; |
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.parameters.ParametersNode; |
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.parameters.resource.ResourceParametersHelper; |
||||||
|
import org.apache.dolphinscheduler.spi.utils.JSONUtils; |
||||||
|
|
||||||
|
public class SagemakerTaskChannel implements TaskChannel { |
||||||
|
|
||||||
|
@Override |
||||||
|
public void cancelApplication(boolean status) { |
||||||
|
|
||||||
|
} |
||||||
|
|
||||||
|
@Override |
||||||
|
public SagemakerTask createTask(TaskExecutionContext taskRequest) { |
||||||
|
return new SagemakerTask(taskRequest); |
||||||
|
} |
||||||
|
|
||||||
|
@Override |
||||||
|
public AbstractParameters parseParameters(ParametersNode parametersNode) { |
||||||
|
return JSONUtils.parseObject(parametersNode.getTaskParams(), SagemakerParameters.class); |
||||||
|
} |
||||||
|
|
||||||
|
@Override |
||||||
|
public ResourceParametersHelper getResources(String parameters) { |
||||||
|
return null; |
||||||
|
} |
||||||
|
|
||||||
|
} |
@ -0,0 +1,45 @@ |
|||||||
|
/* |
||||||
|
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||||
|
* contributor license agreements. See the NOTICE file distributed with |
||||||
|
* this work for additional information regarding copyright ownership. |
||||||
|
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||||
|
* (the "License"); you may not use this file except in compliance 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. |
||||||
|
*/ |
||||||
|
|
||||||
|
package org.apache.dolphinscheduler.plugin.task.sagemaker; |
||||||
|
|
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.TaskChannel; |
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.TaskChannelFactory; |
||||||
|
import org.apache.dolphinscheduler.spi.params.base.PluginParams; |
||||||
|
|
||||||
|
import java.util.Collections; |
||||||
|
import java.util.List; |
||||||
|
|
||||||
|
import com.google.auto.service.AutoService; |
||||||
|
|
||||||
|
@AutoService(TaskChannelFactory.class) |
||||||
|
public class SagemakerTaskChannelFactory implements TaskChannelFactory { |
||||||
|
@Override |
||||||
|
public TaskChannel create() { |
||||||
|
return new SagemakerTaskChannel(); |
||||||
|
} |
||||||
|
|
||||||
|
@Override |
||||||
|
public String getName() { |
||||||
|
return "SAGEMAKER"; |
||||||
|
} |
||||||
|
|
||||||
|
@Override |
||||||
|
public List<PluginParams> getParams() { |
||||||
|
return Collections.emptyList(); |
||||||
|
} |
||||||
|
} |
@ -0,0 +1,36 @@ |
|||||||
|
/* |
||||||
|
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||||
|
* contributor license agreements. See the NOTICE file distributed with |
||||||
|
* this work for additional information regarding copyright ownership. |
||||||
|
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||||
|
* (the "License"); you may not use this file except in compliance 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. |
||||||
|
*/ |
||||||
|
|
||||||
|
package org.apache.dolphinscheduler.plugin.task.sagemaker; |
||||||
|
|
||||||
|
/** |
||||||
|
* Custom SagemakerTaskException |
||||||
|
*/ |
||||||
|
public class SagemakerTaskException extends RuntimeException { |
||||||
|
|
||||||
|
public SagemakerTaskException() { |
||||||
|
super(); |
||||||
|
} |
||||||
|
|
||||||
|
public SagemakerTaskException(String message) { |
||||||
|
super(message); |
||||||
|
} |
||||||
|
|
||||||
|
public SagemakerTaskException(String message, Throwable cause) { |
||||||
|
super(message, cause); |
||||||
|
} |
||||||
|
} |
@ -0,0 +1,132 @@ |
|||||||
|
/* |
||||||
|
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||||
|
* contributor license agreements. See the NOTICE file distributed with |
||||||
|
* this work for additional information regarding copyright ownership. |
||||||
|
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||||
|
* (the "License"); you may not use this file except in compliance 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. |
||||||
|
*/ |
||||||
|
|
||||||
|
package org.apache.dolphinscheduler.plugin.task.sagemaker; |
||||||
|
|
||||||
|
import static org.mockito.ArgumentMatchers.any; |
||||||
|
import static org.powermock.api.mockito.PowerMockito.mock; |
||||||
|
import static org.powermock.api.mockito.PowerMockito.when; |
||||||
|
|
||||||
|
import org.apache.dolphinscheduler.plugin.task.api.TaskExecutionContext; |
||||||
|
import org.apache.dolphinscheduler.spi.utils.JSONUtils; |
||||||
|
import org.apache.dolphinscheduler.spi.utils.PropertyUtils; |
||||||
|
|
||||||
|
import org.apache.commons.io.IOUtils; |
||||||
|
|
||||||
|
import java.io.InputStream; |
||||||
|
import java.nio.charset.StandardCharsets; |
||||||
|
import java.util.ArrayList; |
||||||
|
import java.util.List; |
||||||
|
|
||||||
|
import org.junit.Assert; |
||||||
|
import org.junit.Before; |
||||||
|
import org.junit.Test; |
||||||
|
import org.junit.runner.RunWith; |
||||||
|
import org.mockito.Mockito; |
||||||
|
import org.powermock.api.mockito.PowerMockito; |
||||||
|
import org.powermock.core.classloader.annotations.PowerMockIgnore; |
||||||
|
import org.powermock.core.classloader.annotations.PrepareForTest; |
||||||
|
import org.powermock.core.classloader.annotations.SuppressStaticInitializationFor; |
||||||
|
import org.powermock.modules.junit4.PowerMockRunner; |
||||||
|
|
||||||
|
import com.amazonaws.services.sagemaker.AmazonSageMaker; |
||||||
|
import com.amazonaws.services.sagemaker.model.DescribePipelineExecutionResult; |
||||||
|
import com.amazonaws.services.sagemaker.model.ListPipelineExecutionStepsResult; |
||||||
|
import com.amazonaws.services.sagemaker.model.PipelineExecutionStep; |
||||||
|
import com.amazonaws.services.sagemaker.model.StartPipelineExecutionRequest; |
||||||
|
import com.amazonaws.services.sagemaker.model.StartPipelineExecutionResult; |
||||||
|
import com.amazonaws.services.sagemaker.model.StopPipelineExecutionResult; |
||||||
|
|
||||||
|
@RunWith(PowerMockRunner.class) |
||||||
|
@PrepareForTest({JSONUtils.class, PropertyUtils.class,}) |
||||||
|
@PowerMockIgnore({"javax.*"}) |
||||||
|
@SuppressStaticInitializationFor("org.apache.dolphinscheduler.spi.utils.PropertyUtils") |
||||||
|
public class SagemakerTaskTest { |
||||||
|
|
||||||
|
private final String pipelineExecutionArn = "test-pipeline-arn"; |
||||||
|
private SagemakerTask sagemakerTask; |
||||||
|
private AmazonSageMaker client; |
||||||
|
private PipelineUtils pipelineUtils; |
||||||
|
|
||||||
|
@Before |
||||||
|
public void before() { |
||||||
|
PowerMockito.mockStatic(PropertyUtils.class); |
||||||
|
String parameters = buildParameters(); |
||||||
|
TaskExecutionContext taskExecutionContext = Mockito.mock(TaskExecutionContext.class); |
||||||
|
Mockito.when(taskExecutionContext.getTaskParams()).thenReturn(parameters); |
||||||
|
sagemakerTask = new SagemakerTask(taskExecutionContext); |
||||||
|
sagemakerTask.init(); |
||||||
|
client = mock(AmazonSageMaker.class); |
||||||
|
pipelineUtils = new PipelineUtils(client); |
||||||
|
|
||||||
|
StartPipelineExecutionResult startPipelineExecutionResult = mock(StartPipelineExecutionResult.class); |
||||||
|
when(startPipelineExecutionResult.getPipelineExecutionArn()).thenReturn(pipelineExecutionArn); |
||||||
|
|
||||||
|
StopPipelineExecutionResult stopPipelineExecutionResult = mock(StopPipelineExecutionResult.class); |
||||||
|
when(stopPipelineExecutionResult.getPipelineExecutionArn()).thenReturn(pipelineExecutionArn); |
||||||
|
|
||||||
|
DescribePipelineExecutionResult describePipelineExecutionResult = mock(DescribePipelineExecutionResult.class); |
||||||
|
when(describePipelineExecutionResult.getPipelineExecutionStatus()).thenReturn("Executing", "Succeeded"); |
||||||
|
|
||||||
|
ListPipelineExecutionStepsResult listPipelineExecutionStepsResult = mock(ListPipelineExecutionStepsResult.class); |
||||||
|
PipelineExecutionStep pipelineExecutionStep = mock(PipelineExecutionStep.class); |
||||||
|
List<PipelineExecutionStep> pipelineExecutionSteps = new ArrayList<>(); |
||||||
|
pipelineExecutionSteps.add(pipelineExecutionStep); |
||||||
|
pipelineExecutionSteps.add(pipelineExecutionStep); |
||||||
|
|
||||||
|
when(pipelineExecutionStep.toString()).thenReturn("Test Step1", "Test Step2"); |
||||||
|
when(listPipelineExecutionStepsResult.getPipelineExecutionSteps()).thenReturn(pipelineExecutionSteps); |
||||||
|
|
||||||
|
when(client.startPipelineExecution(any())).thenReturn(startPipelineExecutionResult); |
||||||
|
when(client.stopPipelineExecution(any())).thenReturn(stopPipelineExecutionResult); |
||||||
|
when(client.describePipelineExecution(any())).thenReturn(describePipelineExecutionResult); |
||||||
|
when(client.listPipelineExecutionSteps(any())).thenReturn(listPipelineExecutionStepsResult); |
||||||
|
|
||||||
|
} |
||||||
|
|
||||||
|
@Test |
||||||
|
public void testStartPipelineRequest() throws Exception { |
||||||
|
StartPipelineExecutionRequest request = sagemakerTask.createStartPipelineRequest(); |
||||||
|
Assert.assertEquals("AbalonePipeline", request.getPipelineName()); |
||||||
|
Assert.assertEquals("test Pipeline", request.getPipelineExecutionDescription()); |
||||||
|
Assert.assertEquals("AbalonePipeline", request.getPipelineExecutionDisplayName()); |
||||||
|
Assert.assertEquals("AbalonePipeline", request.getPipelineName()); |
||||||
|
Assert.assertEquals(new Integer(1), request.getParallelismConfiguration().getMaxParallelExecutionSteps()); |
||||||
|
} |
||||||
|
|
||||||
|
@Test |
||||||
|
public void testPipelineExecution() throws Exception { |
||||||
|
pipelineUtils.startPipelineExecution(sagemakerTask.createStartPipelineRequest()); |
||||||
|
Assert.assertEquals(pipelineExecutionArn, pipelineUtils.getPipelineExecutionArn()); |
||||||
|
Assert.assertEquals(0, pipelineUtils.checkPipelineExecutionStatus()); |
||||||
|
pipelineUtils.stopPipelineExecution(); |
||||||
|
} |
||||||
|
|
||||||
|
private String buildParameters() { |
||||||
|
SagemakerParameters parameters = new SagemakerParameters(); |
||||||
|
String sagemakerRequestJson; |
||||||
|
try (InputStream i = this.getClass().getResourceAsStream("SagemakerRequestJson.json")) { |
||||||
|
assert i != null; |
||||||
|
sagemakerRequestJson = IOUtils.toString(i, StandardCharsets.UTF_8); |
||||||
|
} catch (Exception e) { |
||||||
|
throw new RuntimeException(e); |
||||||
|
} |
||||||
|
parameters.setSagemakerRequestJson(sagemakerRequestJson); |
||||||
|
|
||||||
|
return JSONUtils.toJsonString(parameters); |
||||||
|
} |
||||||
|
} |
@ -0,0 +1,9 @@ |
|||||||
|
{ |
||||||
|
"ParallelismConfiguration": { |
||||||
|
"MaxParallelExecutionSteps": 1 |
||||||
|
}, |
||||||
|
"PipelineExecutionDescription": "test Pipeline", |
||||||
|
"PipelineExecutionDisplayName": "AbalonePipeline", |
||||||
|
"PipelineName": "AbalonePipeline", |
||||||
|
"PipelineParameters": [ ] |
||||||
|
} |
After Width: | Height: | Size: 25 KiB |
After Width: | Height: | Size: 139 KiB |
@ -0,0 +1,38 @@ |
|||||||
|
/* |
||||||
|
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||||
|
* contributor license agreements. See the NOTICE file distributed with |
||||||
|
* this work for additional information regarding copyright ownership. |
||||||
|
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||||
|
* (the "License"); you may not use this file except in compliance 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. |
||||||
|
*/ |
||||||
|
import type { IJsonItem } from '../types' |
||||||
|
import { useCustomParams } from '.' |
||||||
|
|
||||||
|
export function useSagemaker(model: { [field: string]: any }): IJsonItem[] { |
||||||
|
|
||||||
|
return [ |
||||||
|
{ |
||||||
|
type: 'editor', |
||||||
|
field: 'sagemakerRequestJson', |
||||||
|
name: "SagemakerRequestJson", |
||||||
|
props: { |
||||||
|
language: 'json' |
||||||
|
}, |
||||||
|
validate: { |
||||||
|
trigger: ['input', 'trigger'], |
||||||
|
required: true, |
||||||
|
message: 'requestJson'
|
||||||
|
} |
||||||
|
}, |
||||||
|
...useCustomParams({ model, field: 'localParams', isSimple: false }) |
||||||
|
] |
||||||
|
} |
@ -0,0 +1,81 @@ |
|||||||
|
/* |
||||||
|
* Licensed to the Apache Software Foundation (ASF) under one or more |
||||||
|
* contributor license agreements. See the NOTICE file distributed with |
||||||
|
* this work for additional information regarding copyright ownership. |
||||||
|
* The ASF licenses this file to You under the Apache License, Version 2.0 |
||||||
|
* (the "License"); you may not use this file except in compliance 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. |
||||||
|
*/ |
||||||
|
|
||||||
|
import { reactive } from 'vue' |
||||||
|
import * as Fields from '../fields/index' |
||||||
|
import type { IJsonItem, INodeData, ITaskData } from '../types' |
||||||
|
|
||||||
|
export function userSagemaker({ |
||||||
|
projectCode, |
||||||
|
from = 0, |
||||||
|
readonly, |
||||||
|
data |
||||||
|
}: { |
||||||
|
projectCode: number |
||||||
|
from?: number |
||||||
|
readonly?: boolean |
||||||
|
data?: ITaskData |
||||||
|
}) { |
||||||
|
const model = reactive({ |
||||||
|
name: '', |
||||||
|
taskType: 'MLFLOW', |
||||||
|
flag: 'YES', |
||||||
|
description: '', |
||||||
|
timeoutFlag: false, |
||||||
|
localParams: [], |
||||||
|
environmentCode: null, |
||||||
|
failRetryInterval: 1, |
||||||
|
failRetryTimes: 0, |
||||||
|
workerGroup: 'default', |
||||||
|
delayTime: 0, |
||||||
|
timeout: 30, |
||||||
|
timeoutNotifyStrategy: ['WARN'], |
||||||
|
} as INodeData) |
||||||
|
|
||||||
|
let extra: IJsonItem[] = [] |
||||||
|
if (from === 1) { |
||||||
|
extra = [ |
||||||
|
Fields.useTaskType(model, readonly), |
||||||
|
Fields.useProcessName({ |
||||||
|
model, |
||||||
|
projectCode, |
||||||
|
isCreate: !data?.id, |
||||||
|
from, |
||||||
|
processName: data?.processName |
||||||
|
}) |
||||||
|
] |
||||||
|
} |
||||||
|
|
||||||
|
return { |
||||||
|
json: [ |
||||||
|
Fields.useName(from), |
||||||
|
...extra, |
||||||
|
Fields.useRunFlag(), |
||||||
|
Fields.useDescription(), |
||||||
|
Fields.useTaskPriority(), |
||||||
|
Fields.useWorkerGroup(), |
||||||
|
Fields.useEnvironmentName(model, !model.id), |
||||||
|
...Fields.useTaskGroup(model, projectCode), |
||||||
|
...Fields.useFailed(), |
||||||
|
Fields.useDelayTime(model), |
||||||
|
...Fields.useTimeoutAlarm(model), |
||||||
|
...Fields.useSagemaker(model), |
||||||
|
Fields.usePreTasks() |
||||||
|
] as IJsonItem[], |
||||||
|
model |
||||||
|
} |
||||||
|
} |
Loading…
Reference in new issue