From 050184f51da03cfd925fc16d84cf7e5f990d37f2 Mon Sep 17 00:00:00 2001 From: Benjamin Trent <4357155+benwtrent@users.noreply.github.com> Date: Wed, 27 Oct 2021 08:52:14 -0400 Subject: [PATCH 1/4] [ML] stop using isAllowedByLicense for model license checks --- .../xpack/core/ml/MachineLearningField.java | 23 +++++++++++++++++++ .../core/ml/inference/TrainedModelConfig.java | 1 + .../TransportInternalInferModelAction.java | 10 ++++++-- .../TransportPutTrainedModelAliasAction.java | 4 +++- .../InferencePipelineAggregationBuilder.java | 3 ++- 5 files changed, 37 insertions(+), 4 deletions(-) diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/MachineLearningField.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/MachineLearningField.java index 3438c03956806..d15a7980176fb 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/MachineLearningField.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/MachineLearningField.java @@ -34,6 +34,17 @@ public final class MachineLearningField { License.OperationMode.PLATINUM ); + public static final LicensedFeature.Momentary ML_MODEL_INFERENCE_PLATINUM_FEATURE = LicensedFeature.momentary( + MachineLearningField.ML_FEATURE_FAMILY, + "model-inference-platinum-check", + License.OperationMode.PLATINUM + ); + public static final LicensedFeature.Momentary ML_MODEL_INFERENCE_BASIC_FEATURE = LicensedFeature.momentary( + MachineLearningField.ML_FEATURE_FAMILY, + "model-inference-basic-check", + License.OperationMode.BASIC + ); + private MachineLearningField() {} public static String valuesToId(String... values) { @@ -45,4 +56,16 @@ public static String valuesToId(String... values) { System.arraycopy(Numbers.longToBytes(hash.h2), 0, hashedBytes, 8, 8); return new BigInteger(hashedBytes) + "_" + combined.length(); } + + public static LicensedFeature.Momentary featureFromLicenseLevel(License.OperationMode mode) { + switch (mode) { + case BASIC: + return ML_MODEL_INFERENCE_BASIC_FEATURE; + case PLATINUM: + return ML_MODEL_INFERENCE_PLATINUM_FEATURE; + default: + assert false: "Unrecognized licensing mode for inference models [" + mode + "]"; + return ML_MODEL_INFERENCE_PLATINUM_FEATURE; + } + } } diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfig.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfig.java index 5f2ddae51436a..4f8955053a4ff 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfig.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfig.java @@ -341,6 +341,7 @@ public long getEstimatedOperations() { } //TODO if we ever support anything other than "basic" and platinum, we need to adjust our feature tracking logic + // Additionally, see `MachineLearningField.featureFromLicenseLevel` for handling modes public License.OperationMode getLicenseLevel() { return licenseLevel; } diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java index 91f54e4f2e407..cd137b6cff25f 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java @@ -16,6 +16,7 @@ import org.elasticsearch.common.inject.Inject; import org.elasticsearch.core.TimeValue; import org.elasticsearch.license.LicenseUtils; +import org.elasticsearch.license.LicensedFeature; import org.elasticsearch.license.XPackLicenseState; import org.elasticsearch.rest.RestStatus; import org.elasticsearch.tasks.Task; @@ -42,6 +43,7 @@ import static org.elasticsearch.xpack.core.ClientHelper.ML_ORIGIN; import static org.elasticsearch.xpack.core.ClientHelper.executeAsyncWithOrigin; +import static org.elasticsearch.xpack.core.ml.MachineLearningField.featureFromLicenseLevel; public class TransportInternalInferModelAction extends HandledTransportAction { @@ -79,8 +81,10 @@ protected void doExecute(Task task, Request request, ActionListener li } else { trainedModelProvider.getTrainedModel(request.getModelId(), GetTrainedModelsAction.Includes.empty(), ActionListener.wrap( trainedModelConfig -> { - responseBuilder.setLicensed(licenseState.isAllowedByLicense(trainedModelConfig.getLicenseLevel())); - if (licenseState.isAllowedByLicense(trainedModelConfig.getLicenseLevel()) || request.isPreviouslyLicensed()) { + LicensedFeature.Momentary check = featureFromLicenseLevel(trainedModelConfig.getLicenseLevel()); + final boolean allowed = check.check(licenseState); + responseBuilder.setLicensed(allowed); + if (allowed || request.isPreviouslyLicensed()) { doInfer(request, responseBuilder, listener); } else { listener.onFailure(LicenseUtils.newComplianceException(XPackField.MACHINE_LEARNING)); @@ -183,4 +187,6 @@ private void inferSingleDocAgainstAllocatedModel( ) ); } + + } diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportPutTrainedModelAliasAction.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportPutTrainedModelAliasAction.java index ff92d882d535e..d24d793fdf251 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportPutTrainedModelAliasAction.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportPutTrainedModelAliasAction.java @@ -47,6 +47,7 @@ import java.util.Set; import java.util.function.Predicate; +import static org.elasticsearch.xpack.core.ml.MachineLearningField.featureFromLicenseLevel; import static org.elasticsearch.xpack.core.ml.job.messages.Messages.TRAINED_MODEL_INPUTS_DIFFER_SIGNIFICANTLY; public class TransportPutTrainedModelAliasAction extends AcknowledgedTransportMasterNodeAction { @@ -90,7 +91,8 @@ protected void masterOperation( ActionListener listener ) throws Exception { final boolean mlSupported = MachineLearningField.ML_API_FEATURE.check(licenseState); - final Predicate isLicensed = (model) -> mlSupported || licenseState.isAllowedByLicense(model.getLicenseLevel()); + final Predicate isLicensed = (model) -> mlSupported + || featureFromLicenseLevel(model.getLicenseLevel()).check(licenseState); final String oldModelId = ModelAliasMetadata.fromState(state).getModelId(request.getModelAlias()); if (oldModelId != null && (request.isReassign() == false)) { diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/inference/InferencePipelineAggregationBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/inference/InferencePipelineAggregationBuilder.java index 9e6b0734595ee..fc6f60c8c1e83 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/inference/InferencePipelineAggregationBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/inference/InferencePipelineAggregationBuilder.java @@ -51,6 +51,7 @@ import java.util.function.Supplier; import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg; +import static org.elasticsearch.xpack.core.ml.MachineLearningField.featureFromLicenseLevel; import static org.elasticsearch.xpack.ml.utils.SecondaryAuthorizationUtils.useSecondaryAuthIfAvailable; public class InferencePipelineAggregationBuilder extends AbstractPipelineAggregationBuilder { @@ -231,7 +232,7 @@ public InferencePipelineAggregationBuilder rewrite(QueryRewriteContext context) loadedModel.set(model); boolean isLicensed = MachineLearningField.ML_API_FEATURE.check(licenseState) || - licenseState.isAllowedByLicense(model.getLicenseLevel()); + featureFromLicenseLevel(model.getLicenseLevel()).check(licenseState); if (isLicensed) { delegate.onResponse(null); } else { From bba0132ed2460119799c69f1ea4361dfa276bbff Mon Sep 17 00:00:00 2001 From: Benjamin Trent <4357155+benwtrent@users.noreply.github.com> Date: Thu, 28 Oct 2021 08:10:50 -0400 Subject: [PATCH 2/4] formatting --- .../org/elasticsearch/xpack/core/ml/MachineLearningField.java | 2 +- .../xpack/ml/action/TransportInternalInferModelAction.java | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/MachineLearningField.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/MachineLearningField.java index 4afa2500eddec..3278ad4d295b2 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/MachineLearningField.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/MachineLearningField.java @@ -70,7 +70,7 @@ public static LicensedFeature.Momentary featureFromLicenseLevel(License.Operatio case PLATINUM: return ML_MODEL_INFERENCE_PLATINUM_FEATURE; default: - assert false: "Unrecognized licensing mode for inference models [" + mode + "]"; + assert false : "Unrecognized licensing mode for inference models [" + mode + "]"; return ML_MODEL_INFERENCE_PLATINUM_FEATURE; } } diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java index 61774b436a4f3..51b4506f165a6 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java @@ -200,5 +200,4 @@ private void inferSingleDocAgainstAllocatedModel( ); } - } From 02fed700a2ac58a893fc5625d76bdb5e84861daa Mon Sep 17 00:00:00 2001 From: Benjamin Trent <4357155+benwtrent@users.noreply.github.com> Date: Thu, 28 Oct 2021 08:39:42 -0400 Subject: [PATCH 3/4] addressing PR comments --- .../xpack/core/ml/MachineLearningField.java | 19 +++++-------------- .../TransportInternalInferModelAction.java | 6 ++---- .../TransportPutTrainedModelAliasAction.java | 4 ++-- .../InferencePipelineAggregationBuilder.java | 4 ++-- 4 files changed, 11 insertions(+), 22 deletions(-) diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/MachineLearningField.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/MachineLearningField.java index 3278ad4d295b2..8be3d4e938893 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/MachineLearningField.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/MachineLearningField.java @@ -13,6 +13,7 @@ import org.elasticsearch.core.TimeValue; import org.elasticsearch.license.License; import org.elasticsearch.license.LicensedFeature; +import org.elasticsearch.license.XPackLicenseState; import java.math.BigInteger; import java.nio.charset.StandardCharsets; @@ -45,11 +46,6 @@ public final class MachineLearningField { "model-inference-platinum-check", License.OperationMode.PLATINUM ); - public static final LicensedFeature.Momentary ML_MODEL_INFERENCE_BASIC_FEATURE = LicensedFeature.momentary( - MachineLearningField.ML_FEATURE_FAMILY, - "model-inference-basic-check", - License.OperationMode.BASIC - ); private MachineLearningField() {} @@ -63,15 +59,10 @@ public static String valuesToId(String... values) { return new BigInteger(hashedBytes) + "_" + combined.length(); } - public static LicensedFeature.Momentary featureFromLicenseLevel(License.OperationMode mode) { - switch (mode) { - case BASIC: - return ML_MODEL_INFERENCE_BASIC_FEATURE; - case PLATINUM: - return ML_MODEL_INFERENCE_PLATINUM_FEATURE; - default: - assert false : "Unrecognized licensing mode for inference models [" + mode + "]"; - return ML_MODEL_INFERENCE_PLATINUM_FEATURE; + public static boolean featureCheckForMode(License.OperationMode mode, XPackLicenseState licenseState) { + if (mode.equals(License.OperationMode.PLATINUM)) { + return ML_MODEL_INFERENCE_PLATINUM_FEATURE.check(licenseState); } + return true; } } diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java index 51b4506f165a6..668cc54ae81e8 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java @@ -16,7 +16,6 @@ import org.elasticsearch.common.inject.Inject; import org.elasticsearch.core.TimeValue; import org.elasticsearch.license.LicenseUtils; -import org.elasticsearch.license.LicensedFeature; import org.elasticsearch.license.XPackLicenseState; import org.elasticsearch.rest.RestStatus; import org.elasticsearch.tasks.Task; @@ -43,7 +42,7 @@ import static org.elasticsearch.xpack.core.ClientHelper.ML_ORIGIN; import static org.elasticsearch.xpack.core.ClientHelper.executeAsyncWithOrigin; -import static org.elasticsearch.xpack.core.ml.MachineLearningField.featureFromLicenseLevel; +import static org.elasticsearch.xpack.core.ml.MachineLearningField.featureCheckForMode; public class TransportInternalInferModelAction extends HandledTransportAction { @@ -84,8 +83,7 @@ protected void doExecute(Task task, Request request, ActionListener li request.getModelId(), GetTrainedModelsAction.Includes.empty(), ActionListener.wrap(trainedModelConfig -> { - LicensedFeature.Momentary check = featureFromLicenseLevel(trainedModelConfig.getLicenseLevel()); - final boolean allowed = check.check(licenseState); + final boolean allowed = featureCheckForMode(trainedModelConfig.getLicenseLevel(), licenseState); responseBuilder.setLicensed(allowed); if (allowed || request.isPreviouslyLicensed()) { doInfer(request, responseBuilder, listener); diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportPutTrainedModelAliasAction.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportPutTrainedModelAliasAction.java index d6fd44d78e2df..37e7a5cefbcd0 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportPutTrainedModelAliasAction.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportPutTrainedModelAliasAction.java @@ -47,7 +47,7 @@ import java.util.Set; import java.util.function.Predicate; -import static org.elasticsearch.xpack.core.ml.MachineLearningField.featureFromLicenseLevel; +import static org.elasticsearch.xpack.core.ml.MachineLearningField.featureCheckForMode; import static org.elasticsearch.xpack.core.ml.job.messages.Messages.TRAINED_MODEL_INPUTS_DIFFER_SIGNIFICANTLY; public class TransportPutTrainedModelAliasAction extends AcknowledgedTransportMasterNodeAction { @@ -93,7 +93,7 @@ protected void masterOperation( ) throws Exception { final boolean mlSupported = MachineLearningField.ML_API_FEATURE.check(licenseState); final Predicate isLicensed = (model) -> mlSupported - || featureFromLicenseLevel(model.getLicenseLevel()).check(licenseState); + || featureCheckForMode(model.getLicenseLevel(), licenseState); final String oldModelId = ModelAliasMetadata.fromState(state).getModelId(request.getModelAlias()); if (oldModelId != null && (request.isReassign() == false)) { diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/inference/InferencePipelineAggregationBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/inference/InferencePipelineAggregationBuilder.java index e012fddd4e4f2..e2e13810cd28c 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/inference/InferencePipelineAggregationBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/inference/InferencePipelineAggregationBuilder.java @@ -51,7 +51,7 @@ import java.util.function.Supplier; import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg; -import static org.elasticsearch.xpack.core.ml.MachineLearningField.featureFromLicenseLevel; +import static org.elasticsearch.xpack.core.ml.MachineLearningField.featureCheckForMode; import static org.elasticsearch.xpack.ml.utils.SecondaryAuthorizationUtils.useSecondaryAuthIfAvailable; public class InferencePipelineAggregationBuilder extends AbstractPipelineAggregationBuilder { @@ -268,7 +268,7 @@ public InferencePipelineAggregationBuilder rewrite(QueryRewriteContext context) loadedModel.set(model); boolean isLicensed = MachineLearningField.ML_API_FEATURE.check(licenseState) - || featureFromLicenseLevel(model.getLicenseLevel()).check(licenseState); + || featureCheckForMode(model.getLicenseLevel(), licenseState); if (isLicensed) { delegate.onResponse(null); } else { From dee500e5892837d34ab7c6b948e22b299b567615 Mon Sep 17 00:00:00 2001 From: Benjamin Trent Date: Thu, 28 Oct 2021 08:55:27 -0400 Subject: [PATCH 4/4] Apply suggestions from code review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Przemysław Witek --- .../xpack/core/ml/inference/TrainedModelConfig.java | 2 +- .../xpack/ml/action/TransportInternalInferModelAction.java | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfig.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfig.java index 74ce2b8799fd6..9ff0893de7207 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfig.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfig.java @@ -351,7 +351,7 @@ public long getEstimatedOperations() { } // TODO if we ever support anything other than "basic" and platinum, we need to adjust our feature tracking logic - // Additionally, see `MachineLearningField.featureFromLicenseLevel` for handling modes + // Additionally, see `MachineLearningField. featureCheckForMode` for handling modes public License.OperationMode getLicenseLevel() { return licenseLevel; } diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java index 668cc54ae81e8..1a3782225a037 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java @@ -197,5 +197,4 @@ private void inferSingleDocAgainstAllocatedModel( }) ); } - }