|
|
|
@ -22,6 +22,7 @@ import java.util.UUID;
|
|
|
|
|
|
|
|
|
|
import org.apache.dolphinscheduler.plugin.task.api.TaskExecutionContext; |
|
|
|
|
import org.apache.dolphinscheduler.plugin.task.api.TaskExecutionContextCacheManager; |
|
|
|
|
import org.apache.dolphinscheduler.plugin.task.api.utils.OSUtils; |
|
|
|
|
import org.apache.dolphinscheduler.plugin.task.mlflow.MlflowConstants; |
|
|
|
|
import org.apache.dolphinscheduler.plugin.task.mlflow.MlflowParameters; |
|
|
|
|
import org.apache.dolphinscheduler.plugin.task.mlflow.MlflowTask; |
|
|
|
@ -41,7 +42,6 @@ import org.powermock.modules.junit4.PowerMockRunner;
|
|
|
|
|
import org.powermock.core.classloader.annotations.SuppressStaticInitializationFor; |
|
|
|
|
import org.apache.dolphinscheduler.spi.utils.JSONUtils; |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@RunWith(PowerMockRunner.class) |
|
|
|
|
@PrepareForTest({ |
|
|
|
|
JSONUtils.class, |
|
|
|
@ -129,9 +129,7 @@ public class MlflowTaskTest {
|
|
|
|
|
MlflowTask mlflowTask = initTask(createModelDeplyMlflowParameters()); |
|
|
|
|
Assert.assertEquals(mlflowTask.buildCommand(), |
|
|
|
|
"export MLFLOW_TRACKING_URI=http://127.0.0.1:5000\n" + |
|
|
|
|
"mlflow models serve -m runs:/a272ec279fc34a8995121ae04281585f/model " + |
|
|
|
|
"--port 7000 " + |
|
|
|
|
"-h 0.0.0.0"); |
|
|
|
|
"mlflow models serve -m models:/model/1 --port 7000 -h 0.0.0.0"); |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
@Test |
|
|
|
@ -139,12 +137,11 @@ public class MlflowTaskTest {
|
|
|
|
|
MlflowTask mlflowTask = initTask(createModelDeplyDockerParameters()); |
|
|
|
|
Assert.assertEquals(mlflowTask.buildCommand(), |
|
|
|
|
"export MLFLOW_TRACKING_URI=http://127.0.0.1:5000\n" + |
|
|
|
|
"mlflow models build-docker -m runs:/a272ec279fc34a8995121ae04281585f/model " + |
|
|
|
|
"-n mlflow/a272ec279fc34a8995121ae04281585f:model " + |
|
|
|
|
"--enable-mlserver\n" + |
|
|
|
|
"docker rm -f ds-mlflow-a272ec279fc34a8995121ae04281585f-model\n" + |
|
|
|
|
"docker run --name=ds-mlflow-a272ec279fc34a8995121ae04281585f-model " + |
|
|
|
|
"-p=7000:8080 mlflow/a272ec279fc34a8995121ae04281585f:model"); |
|
|
|
|
"mlflow models build-docker -m models:/model/1 -n mlflow/model:1 --enable-mlserver\n" + |
|
|
|
|
"docker rm -f ds-mlflow-model-1\n" + |
|
|
|
|
"docker run -d --name=ds-mlflow-model-1 -p=7000:8080 " + |
|
|
|
|
"--health-cmd \"curl --fail http://127.0.0.1:8080/ping || exit 1\" --health-interval 5s --health-retries 20 " + |
|
|
|
|
"mlflow/model:1"); |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
@Test |
|
|
|
@ -154,16 +151,14 @@ public class MlflowTaskTest {
|
|
|
|
|
"export MLFLOW_TRACKING_URI=http://127.0.0.1:5000\n" + |
|
|
|
|
"cp " + mlflowTask.getTemplatePath(MlflowConstants.TEMPLATE_DOCKER_COMPOSE) + |
|
|
|
|
" /tmp/dolphinscheduler_test\n" + |
|
|
|
|
"mlflow models build-docker -m models:/22222/1 -n mlflow/22222:1 --enable-mlserver\n" + |
|
|
|
|
"export DS_TASK_MLFLOW_IMAGE_NAME=mlflow/22222:1\n" + |
|
|
|
|
"export DS_TASK_MLFLOW_CONTAINER_NAME=ds-mlflow-22222-1\n" + |
|
|
|
|
"mlflow models build-docker -m models:/model/1 -n mlflow/model:1 --enable-mlserver\n" + |
|
|
|
|
"docker rm -f ds-mlflow-model-1\n" + |
|
|
|
|
"export DS_TASK_MLFLOW_IMAGE_NAME=mlflow/model:1\n" + |
|
|
|
|
"export DS_TASK_MLFLOW_CONTAINER_NAME=ds-mlflow-model-1\n" + |
|
|
|
|
"export DS_TASK_MLFLOW_DEPLOY_PORT=7000\n" + |
|
|
|
|
"export DS_TASK_MLFLOW_CPU_LIMIT=0.5\n" + |
|
|
|
|
"export DS_TASK_MLFLOW_MEMORY_LIMIT=200m\n" + |
|
|
|
|
"docker-compose up -d\n" + |
|
|
|
|
"for i in $(seq 1 300); do " + |
|
|
|
|
"[ $(docker inspect --format \"{{json .State.Health.Status }}\" ds-mlflow-22222-1) = '\"healthy\"' ] && exit 0 && break;sleep 1; " + |
|
|
|
|
"done; docker-compose down; exit 1"); |
|
|
|
|
"docker-compose up -d"); |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
private MlflowTask initTask(MlflowParameters mlflowParameters) { |
|
|
|
@ -172,7 +167,6 @@ public class MlflowTaskTest {
|
|
|
|
|
mlflowTask.init(); |
|
|
|
|
mlflowTask.getParameters().setVarPool(taskExecutionContext.getVarPool()); |
|
|
|
|
return mlflowTask; |
|
|
|
|
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
private MlflowParameters createBasicAlgorithmParameters() { |
|
|
|
@ -218,7 +212,7 @@ public class MlflowTaskTest {
|
|
|
|
|
mlflowParameters.setMlflowTaskType(MlflowConstants.MLFLOW_TASK_TYPE_MODELS); |
|
|
|
|
mlflowParameters.setDeployType(MlflowConstants.MLFLOW_MODELS_DEPLOY_TYPE_MLFLOW); |
|
|
|
|
mlflowParameters.setMlflowTrackingUris("http://127.0.0.1:5000"); |
|
|
|
|
mlflowParameters.setDeployModelKey("runs:/a272ec279fc34a8995121ae04281585f/model"); |
|
|
|
|
mlflowParameters.setDeployModelKey("models:/model/1"); |
|
|
|
|
mlflowParameters.setDeployPort("7000"); |
|
|
|
|
return mlflowParameters; |
|
|
|
|
} |
|
|
|
@ -228,7 +222,7 @@ public class MlflowTaskTest {
|
|
|
|
|
mlflowParameters.setMlflowTaskType(MlflowConstants.MLFLOW_TASK_TYPE_MODELS); |
|
|
|
|
mlflowParameters.setDeployType(MlflowConstants.MLFLOW_MODELS_DEPLOY_TYPE_DOCKER); |
|
|
|
|
mlflowParameters.setMlflowTrackingUris("http://127.0.0.1:5000"); |
|
|
|
|
mlflowParameters.setDeployModelKey("runs:/a272ec279fc34a8995121ae04281585f/model"); |
|
|
|
|
mlflowParameters.setDeployModelKey("models:/model/1"); |
|
|
|
|
mlflowParameters.setDeployPort("7000"); |
|
|
|
|
return mlflowParameters; |
|
|
|
|
} |
|
|
|
@ -238,7 +232,7 @@ public class MlflowTaskTest {
|
|
|
|
|
mlflowParameters.setMlflowTaskType(MlflowConstants.MLFLOW_TASK_TYPE_MODELS); |
|
|
|
|
mlflowParameters.setDeployType(MlflowConstants.MLFLOW_MODELS_DEPLOY_TYPE_DOCKER_COMPOSE); |
|
|
|
|
mlflowParameters.setMlflowTrackingUris("http://127.0.0.1:5000"); |
|
|
|
|
mlflowParameters.setDeployModelKey("models:/22222/1"); |
|
|
|
|
mlflowParameters.setDeployModelKey("models:/model/1"); |
|
|
|
|
mlflowParameters.setDeployPort("7000"); |
|
|
|
|
mlflowParameters.setCpuLimit("0.5"); |
|
|
|
|
mlflowParameters.setMemoryLimit("200m"); |
|
|
|
|