Skip to content

Commit eb46f1e

Browse files
authored
add jupyterClient implement (#53)
* remove TaskResponse/TaskExecutionContext to protocol submodule && add jupyter start/stop * mode BaseRequest/BaseResponse to wedpr-core-protocol * add jupyterClient implement
1 parent 05bed23 commit eb46f1e

File tree

53 files changed

+859
-106
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

53 files changed

+859
-106
lines changed

db/wedpr_dml.sql

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,16 @@
11
-- the Wizard algorithm template
22
insert into `wedpr_config_table`(`config_key`, `config_value`) values("wedpr_algorithm_templates", '{"version":"1.0","templates":[{"name":"PSI","title":"数据对齐","detail":"","version":"1.0"},{"name":"XGB_TRAINING","title":"SecureLGBM训练","detail":"","version":"1.0"},{"name":"XGB_PREDICTING","title":"SecureLGBM预测","detail":"","version":"1.0"}]}');
33

4+
5+
-- the jupyter related host settings
6+
-- insert into `wedpr_config_table`(`config_key`, `config_value`) values("jupyter_entrypoints", '{"hostSettings":[{"entryPoint":"192.168.0.238","maxJupyterCount":5,"jupyterStartPort":14000}]}');
7+
48
-- the jupyter related code template
59
-- command to create user
6-
insert into `wedpr_config_table`(`config_key`, `config_value`) values("wedpr_create_user", 'useradd -m ${user_name} -p ${user_pass}');
10+
insert into `wedpr_config_table`(`config_key`, `config_value`) values("wedpr_create_user", 'useradd -m ${user_name}');
711
-- command to start jupyter
8-
insert into `wedpr_config_table`(`config_key`, `config_value`) values("wedpr_start_jupyter", 'su ${user_name} && mkdir -p ~/project && jupyter-lab --config=/home/share/.jupyter/jupyter_lab_config.py --ip ${listen_ip} --port ${listen_port} --notebook-dir ~/project');
9-
insert into `wedpr_config_table`(`config_key`, `config_value`) values("wedpr_stop_jupyter", 'ps aux | grep -i jupyter-lab | grep -i ${user_name} | grep -v grep | awk -F\' \' \'{print $2}\' | xargs kill');
12+
insert into `wedpr_config_table`(`config_key`, `config_value`) values("wedpr_start_jupyter", 'su ${user_name} && mkdir -p ${jupyter_project_path} && nohup ${jupyter_binary} --config=${jupyter_config_path} --ip ${listen_ip} --port ${listen_port} --notebook-dir ${jupyter_project_path} >> ${user_name}.jupyter.out 2>&1 &');
13+
insert into `wedpr_config_table`(`config_key`, `config_value`) values("wedpr_get_jupyter_pid", 'ps aux | grep -i ${jupyter_binary} | grep -i ${jupyter_project_path} | grep -v grep | awk -F\' \' \'{print $2}\'');
1014
insert into `wedpr_config_table`(`config_key`, `config_value`) values("wedpr_delete_user", 'userdel -m ${user_name}');
1115

1216

wedpr-adm/src/main/java/com/webank/wedpr/adm/controller/JupyterController.java

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -162,7 +162,11 @@ public WeDPRResponse destroy(@RequestParam String id, HttpServletRequest request
162162
try {
163163
UserToken userToken = TokenUtils.getLoginUser(request);
164164
boolean success =
165-
this.jupyterService.destroy(userToken.isAdmin(), userToken.getUsername(), id);
165+
this.jupyterService.destroy(
166+
userToken.isAdmin(),
167+
userToken.getUsername(),
168+
WeDPRCommonConfig.getAgency(),
169+
id);
166170
WeDPRResponse response =
167171
new WeDPRResponse(Constant.WEDPR_SUCCESS, Constant.WEDPR_SUCCESS_MSG);
168172
response.setData(success);

wedpr-components/env-integration/jupyter/build.gradle

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,5 +6,6 @@ dependencies{
66
compile project(":wedpr-components-mybatis")
77
compile project(":wedpr-components-uuid")
88
compile project(":wedpr-components-sys-config")
9+
compile project(":wedpr-components-http-client")
910
compile project(":wedpr-components-hook")
1011
}
Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
/*
2+
* Copyright 2017-2025 [webank-wedpr]
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
5+
* in compliance with the License. You may obtain a copy of the License at
6+
*
7+
* http://www.apache.org/licenses/LICENSE-2.0
8+
*
9+
* Unless required by applicable law or agreed to in writing, software distributed under the License
10+
* is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
11+
* or implied. See the License for the specific language governing permissions and limitations under
12+
* the License.
13+
*
14+
*/
15+
package com.webank.wedpr.components.integration.jupyter.client;
16+
17+
import com.webank.wedpr.components.integration.jupyter.core.JupyterConfig;
18+
import com.webank.wedpr.components.integration.jupyter.dao.JupyterInfoDO;
19+
import com.webank.wedpr.core.config.WeDPRCommonConfig;
20+
import com.webank.wedpr.core.protocol.task.TaskResponse;
21+
import java.util.HashMap;
22+
import java.util.Map;
23+
24+
public interface JupyterClient {
25+
26+
TaskResponse create(JupyterInfoDO jupyterInfo) throws Exception;
27+
28+
TaskResponse start(JupyterInfoDO jupyterInfo) throws Exception;
29+
30+
TaskResponse stop(JupyterInfoDO jupyterInfo) throws Exception;
31+
32+
static Map<String, String> generateParamMap(JupyterInfoDO jupyterInfoDO) {
33+
Map<String, String> paramMap = new HashMap<>();
34+
// the username
35+
paramMap.put(WeDPRCommonConfig.getParamKeyUserName(), jupyterInfoDO.getOwner());
36+
// the listen ip
37+
paramMap.put(
38+
WeDPRCommonConfig.getParamKeyListenIp(),
39+
jupyterInfoDO.getJupyterSetting().getListenIp());
40+
// the listen port
41+
paramMap.put(
42+
WeDPRCommonConfig.getParamKeyListenPort(),
43+
String.valueOf(jupyterInfoDO.getJupyterSetting().getListenPort()));
44+
// the jupyter binary
45+
paramMap.put(
46+
JupyterConfig.getParamKeyJupyterBinary(),
47+
jupyterInfoDO.getJupyterSetting().getJupyterBinary());
48+
// the jupyter config path
49+
paramMap.put(
50+
JupyterConfig.getParamKeyJupyterConfigPath(),
51+
jupyterInfoDO.getJupyterSetting().getConfigPath());
52+
// the notebook project directory
53+
paramMap.put(
54+
JupyterConfig.getParamKeyJupyterProjectPath(),
55+
jupyterInfoDO.getJupyterSetting().getNoteBookPath());
56+
return paramMap;
57+
}
58+
}
Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
/*
2+
* Copyright 2017-2025 [webank-wedpr]
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
5+
* in compliance with the License. You may obtain a copy of the License at
6+
*
7+
* http://www.apache.org/licenses/LICENSE-2.0
8+
*
9+
* Unless required by applicable law or agreed to in writing, software distributed under the License
10+
* is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
11+
* or implied. See the License for the specific language governing permissions and limitations under
12+
* the License.
13+
*
14+
*/
15+
package com.webank.wedpr.components.integration.jupyter.client.impl;
16+
17+
import com.webank.wedpr.components.http.client.HttpClientImpl;
18+
import com.webank.wedpr.components.http.client.HttpClientPool;
19+
import com.webank.wedpr.components.integration.jupyter.client.JupyterClient;
20+
import com.webank.wedpr.components.integration.jupyter.core.JupyterConfig;
21+
import com.webank.wedpr.components.integration.jupyter.dao.JupyterInfoDO;
22+
import com.webank.wedpr.components.meta.sys.config.dao.SysConfigDO;
23+
import com.webank.wedpr.components.meta.sys.config.dao.SysConfigMapper;
24+
import com.webank.wedpr.components.uuid.generator.WeDPRUuidGenerator;
25+
import com.webank.wedpr.core.config.WeDPRCommonConfig;
26+
import com.webank.wedpr.core.protocol.task.ShellParameters;
27+
import com.webank.wedpr.core.protocol.task.TaskExecutionContext;
28+
import com.webank.wedpr.core.protocol.task.TaskResponse;
29+
import com.webank.wedpr.core.protocol.task.TaskType;
30+
import com.webank.wedpr.core.utils.BaseResponseFactory;
31+
import com.webank.wedpr.core.utils.ShellConstant;
32+
import com.webank.wedpr.core.utils.WeDPRException;
33+
import com.webank.wedpr.core.utils.WeDPRResponseFactory;
34+
import java.util.ArrayList;
35+
import java.util.List;
36+
import java.util.stream.Collectors;
37+
import org.slf4j.Logger;
38+
import org.slf4j.LoggerFactory;
39+
40+
public class JupyterClientImpl implements JupyterClient {
41+
private static final Logger logger = LoggerFactory.getLogger(JupyterClientImpl.class);
42+
private final SysConfigMapper sysConfigMapper;
43+
private final BaseResponseFactory responseFactory = new WeDPRResponseFactory();
44+
45+
public JupyterClientImpl(SysConfigMapper sysConfigMapper) {
46+
this.sysConfigMapper = sysConfigMapper;
47+
}
48+
49+
protected String getCodeTemplate(String key, boolean mustExist) throws WeDPRException {
50+
SysConfigDO result = this.sysConfigMapper.queryConfig(key);
51+
if (result != null) {
52+
return result.getConfigValue();
53+
}
54+
if (mustExist) {
55+
throw new WeDPRException(
56+
"The code template for " + key + " not configure! Please configure firstly");
57+
}
58+
return null;
59+
}
60+
61+
protected HttpClientImpl generateHttpClient(String url) {
62+
String apiUri =
63+
String.format(
64+
"%s/%s/%s",
65+
url,
66+
WeDPRCommonConfig.getWedprWorkerApiPath(),
67+
WeDPRCommonConfig.getWedprWorkerSubmitTaskMethod());
68+
return new HttpClientImpl(
69+
HttpClientPool.getUrl(apiUri),
70+
JupyterConfig.getMaxTotalConnection(),
71+
JupyterConfig.buildConfig(),
72+
responseFactory);
73+
}
74+
75+
protected TaskResponse submitTask(String method, JupyterInfoDO jupyterInfoDO, String code)
76+
throws Exception {
77+
logger.info("Submit code to: {}", jupyterInfoDO.getAccessEntry());
78+
HttpClientImpl httpClient = generateHttpClient(jupyterInfoDO.getAccessEntry());
79+
TaskExecutionContext taskRequest = new TaskExecutionContext();
80+
taskRequest.setTaskID(WeDPRUuidGenerator.generateID());
81+
taskRequest.setTaskType(TaskType.SHELL.getType());
82+
ShellParameters shellParameters = new ShellParameters(code);
83+
taskRequest.setTaskParameters(shellParameters.serialize());
84+
taskRequest.setParameterMap(JupyterClient.generateParamMap(jupyterInfoDO));
85+
TaskResponse response = (TaskResponse) httpClient.executePost(taskRequest);
86+
if (response != null && response.statusOk()) {
87+
logger.info(
88+
"submitTask for method {} success, taskID: {}, jupyterInfo: {}, code: {}, response: {}",
89+
method,
90+
taskRequest.getTaskID(),
91+
jupyterInfoDO.toString(),
92+
code,
93+
response);
94+
return response;
95+
}
96+
logger.error(
97+
"submitTask for method {} failed, taskID: {}, jupyterInfo: {}, code: {}, response: {}",
98+
method,
99+
taskRequest.getTaskID(),
100+
jupyterInfoDO.toString(),
101+
code,
102+
(response == null ? "null" : response.toString()));
103+
throw new WeDPRException(
104+
"submitTask for "
105+
+ method
106+
+ " failed, taskID: "
107+
+ taskRequest.getTaskID()
108+
+ ", jupyterInfo: "
109+
+ jupyterInfoDO.toString()
110+
+ ", response: "
111+
+ (response == null ? "null" : response.toString()));
112+
}
113+
114+
@Override
115+
public TaskResponse create(JupyterInfoDO jupyterInfo) throws Exception {
116+
// create user and start jupyter
117+
String code =
118+
getCodeTemplate(WeDPRCommonConfig.getCodeTemplateKeyCreateUser(), true)
119+
+ WeDPRCommonConfig.getShellCodeConnector()
120+
+ getCodeTemplate(JupyterConfig.getCodeTemplateKeyStartJupyter(), true);
121+
logger.info("create jupyter, info: {}, code: {}", jupyterInfo.toString(), code);
122+
return submitTask("createJupyter", jupyterInfo, code);
123+
}
124+
125+
@Override
126+
public TaskResponse start(JupyterInfoDO jupyterInfo) throws Exception {
127+
// Note: should check the existence firstly
128+
List<String> commands = new ArrayList<>();
129+
commands.add(ShellConstant.BASH_HEADER);
130+
commands.add(ShellConstant.BASE_DIR_CMD);
131+
// check the pid existence
132+
String getJupyterPidCode =
133+
getCodeTemplate(JupyterConfig.getCodeTemplateKeyGetJupyterPid(), true);
134+
commands.add(String.format("jupyter_pid=$(%s)", getJupyterPidCode));
135+
commands.add("if [ ! -z ${jupyter_pid} ];then ");
136+
commands.add("echo \"${jupyter_pid}\"");
137+
commands.add("else");
138+
commands.add(getCodeTemplate(JupyterConfig.getCodeTemplateKeyStartJupyter(), true));
139+
commands.add("sleep 1.5");
140+
commands.add(String.format("echo $(%s)", getJupyterPidCode));
141+
commands.add("fi");
142+
String code = commands.stream().collect(Collectors.joining(System.lineSeparator()));
143+
logger.info("start jupyter, info: {}, code: {}", jupyterInfo.toString(), code);
144+
return submitTask("startJupyter", jupyterInfo, code);
145+
}
146+
147+
@Override
148+
public TaskResponse stop(JupyterInfoDO jupyterInfo) throws Exception {
149+
List<String> commands = new ArrayList<>();
150+
commands.add(ShellConstant.BASH_HEADER);
151+
commands.add(ShellConstant.BASE_DIR_CMD);
152+
// check the pid existence
153+
String getJupyterPidCode =
154+
getCodeTemplate(JupyterConfig.getCodeTemplateKeyGetJupyterPid(), true);
155+
commands.add(String.format("jupyter_pid=$(%s)", getJupyterPidCode));
156+
commands.add("[ ! -z ${jupyter_pid} ] && kill ${jupyter_pid} > /dev/null");
157+
commands.add("[ -z ${jupyter_pid} ] && exit 0");
158+
commands.add("sleep 2");
159+
commands.add(String.format("jupyter_pid=$(%s)", getJupyterPidCode));
160+
commands.add("[ ! -z ${jupyter_pid} ] && exit 1");
161+
commands.add("[ -z ${jupyter_pid} ] && exit 0");
162+
String code = commands.stream().collect(Collectors.joining(System.lineSeparator()));
163+
logger.info("stop jupyter, info: {}, code: {}", jupyterInfo.toString(), code);
164+
return submitTask("stopJupyter", jupyterInfo, code);
165+
}
166+
}

wedpr-components/env-integration/jupyter/src/main/java/com/webank/wedpr/components/integration/jupyter/core/JupyterConfig.java

Lines changed: 104 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,18 +16,62 @@
1616
package com.webank.wedpr.components.integration.jupyter.core;
1717

1818
import com.webank.wedpr.core.config.WeDPRConfig;
19+
import com.webank.wedpr.core.utils.Common;
20+
import org.apache.http.client.config.RequestConfig;
1921

2022
public class JupyterConfig {
2123

2224
private static Integer MAX_JUPYTER_PER_HOST =
2325
WeDPRConfig.apply("wedpr.jupyter.max_count_per_host", 3);
24-
private static String JUPYTER_HOST_CONFIGUATINON_KEY =
26+
private static String JUPYTER_HOST_CONFIGURATION_KEY =
2527
WeDPRConfig.apply("wedpr.jupyter.host_configuration_key", "jupyter_entrypoints");
2628
private static String JUPYTER_ENTRYPOINT_SPLITTER = ";";
2729
private static String JUPYTER_MODULE = "jupyter-integration";
30+
private static String DEFAULT_JUPYTER_BINARY =
31+
WeDPRConfig.apply("wedpr.jupyter.binary", "jupyter-lab");
32+
private static String DEFAULT_JUPYTER_CONFIG_PATH =
33+
WeDPRConfig.apply(
34+
"wedpr.juypter.config.path", "/home/share/.jupyter/jupyter_lab_config.py");
35+
private static String DEFAULT_JUPYTER_LISTEN_IP =
36+
WeDPRConfig.apply("wedpr.jupyter.listen.ip", "0.0.0.0");
37+
private static String DEFAULT_JUPYTER_PROJECT_DIR =
38+
WeDPRConfig.apply("wedpr.jupyter.project.dir", "project");
39+
private static Integer DEFAULT_JUPYTER_START_PORT =
40+
WeDPRConfig.apply("wedpr.jupyter.jupyter.default.startPort", 14001);
2841

29-
public static String getJupyterHostConfiguatinonKey() {
30-
return JUPYTER_HOST_CONFIGUATINON_KEY;
42+
///// the parameter map key configuration for jupyter
43+
private static String PARAM_KEY_JUPYTER_CONFIG_PATH =
44+
WeDPRConfig.apply("wedpr.jupyter.param.key.jupyter_config_path", "jupyter_config_path");
45+
private static String PARAM_KEY_JUPYTER_PROJECT_PATH =
46+
WeDPRConfig.apply(
47+
"wedpr.jupyter.param.key.jupyter_project_path", "jupyter_project_path");
48+
private static String PARAM_KEY_JUPYTER_BINARY =
49+
WeDPRConfig.apply("wedpr.jupyter.param.key.jupyter_binary", "jupyter_binary");
50+
///// the parameter map key configuration for jupyter
51+
52+
//// the key to store the jupyter code template
53+
private static String CODE_TEMPLATE_KEY_START_JUPYTER =
54+
WeDPRConfig.apply("wedpr.code.template.key.start_jupyter", "wedpr_start_jupyter");
55+
private static String CODE_TEMPLATE_KEY_GET_JUPYTER_PID =
56+
WeDPRConfig.apply(
57+
"wedpr.code.template.key.wedpr_get_jupyter_pid", "wedpr_get_jupyter_pid");
58+
//// the key to store the jupyter code template
59+
60+
/// the jupyter worker client config ///
61+
private static final Integer CONNECTION_REQUEST_TIME_OUT =
62+
WeDPRConfig.apply("wedpr.jupyter.connect.request.timeout.ms", 10000);
63+
private static final Integer CONNECTION_TIME_OUT =
64+
WeDPRConfig.apply("wedpr.jupyter.connect.timeout.ms", 5000);
65+
private static final Integer REQUEST_TIMEOUT =
66+
WeDPRConfig.apply("wedpr.jupyter.request.timeout.ms", 60000);
67+
private static final Integer MAX_TOTAL_CONNECTION =
68+
WeDPRConfig.apply("wedpr.jupyter.max.total.connection", 5);
69+
/// the jupyter worker client config ///
70+
71+
private static String DEFAULT_HOME_DIR = "/home";
72+
73+
public static String getJupyterHostConfigurationKey() {
74+
return JUPYTER_HOST_CONFIGURATION_KEY;
3175
}
3276

3377
public static Integer getMaxJupyterPerHost() {
@@ -41,4 +85,61 @@ public static String getJupyterEntrypointSplitter() {
4185
public static String getJupyterModule() {
4286
return JUPYTER_MODULE;
4387
}
88+
89+
public static String getDefaultJupyterBinary() {
90+
return DEFAULT_JUPYTER_BINARY;
91+
}
92+
93+
public static String getDefaultJupyterConfigPath() {
94+
return DEFAULT_JUPYTER_CONFIG_PATH;
95+
}
96+
97+
public static String getDefaultJupyterListenIp() {
98+
return DEFAULT_JUPYTER_LISTEN_IP;
99+
}
100+
101+
public static String getDefaultJupyterProjectDir() {
102+
return DEFAULT_JUPYTER_PROJECT_DIR;
103+
}
104+
105+
public static String getUserJupyterProjectPath(String userName) {
106+
return Common.joinPath(
107+
Common.joinPath(DEFAULT_HOME_DIR, userName), getDefaultJupyterProjectDir());
108+
}
109+
110+
public static Integer getDefaultJupyterStartPort() {
111+
return DEFAULT_JUPYTER_START_PORT;
112+
}
113+
114+
public static String getParamKeyJupyterConfigPath() {
115+
return PARAM_KEY_JUPYTER_CONFIG_PATH;
116+
}
117+
118+
public static String getParamKeyJupyterProjectPath() {
119+
return PARAM_KEY_JUPYTER_PROJECT_PATH;
120+
}
121+
122+
public static String getParamKeyJupyterBinary() {
123+
return PARAM_KEY_JUPYTER_BINARY;
124+
}
125+
126+
public static String getCodeTemplateKeyStartJupyter() {
127+
return CODE_TEMPLATE_KEY_START_JUPYTER;
128+
}
129+
130+
public static String getCodeTemplateKeyGetJupyterPid() {
131+
return CODE_TEMPLATE_KEY_GET_JUPYTER_PID;
132+
}
133+
134+
public static RequestConfig buildConfig() {
135+
return RequestConfig.custom()
136+
.setConnectionRequestTimeout(CONNECTION_REQUEST_TIME_OUT)
137+
.setConnectTimeout(CONNECTION_TIME_OUT)
138+
.setSocketTimeout(REQUEST_TIMEOUT)
139+
.build();
140+
}
141+
142+
public static Integer getMaxTotalConnection() {
143+
return MAX_TOTAL_CONNECTION;
144+
}
44145
}

0 commit comments

Comments
 (0)