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 128f6315994a9..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; @@ -40,6 +41,12 @@ 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 + ); + private MachineLearningField() {} public static String valuesToId(String... values) { @@ -51,4 +58,11 @@ public static String valuesToId(String... values) { System.arraycopy(Numbers.longToBytes(hash.h2), 0, hashedBytes, 8, 8); return new BigInteger(hashedBytes) + "_" + combined.length(); } + + 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/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 cad6a3668c223..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,6 +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. 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 c3742b670bd61..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 @@ -42,6 +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.featureCheckForMode; public class TransportInternalInferModelAction extends HandledTransportAction { @@ -82,8 +83,9 @@ protected void doExecute(Task task, Request request, ActionListener li request.getModelId(), GetTrainedModelsAction.Includes.empty(), ActionListener.wrap(trainedModelConfig -> { - responseBuilder.setLicensed(licenseState.isAllowedByLicense(trainedModelConfig.getLicenseLevel())); - if (licenseState.isAllowedByLicense(trainedModelConfig.getLicenseLevel()) || request.isPreviouslyLicensed()) { + final boolean allowed = featureCheckForMode(trainedModelConfig.getLicenseLevel(), licenseState); + responseBuilder.setLicensed(allowed); + if (allowed || request.isPreviouslyLicensed()) { doInfer(request, responseBuilder, listener); } else { listener.onFailure(LicenseUtils.newComplianceException(XPackField.MACHINE_LEARNING)); 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 d0c6611296b1c..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,6 +47,7 @@ import java.util.Set; import java.util.function.Predicate; +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 { @@ -91,7 +92,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 + || 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 e321d8ccf28e4..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,6 +51,7 @@ import java.util.function.Supplier; import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg; +import static org.elasticsearch.xpack.core.ml.MachineLearningField.featureCheckForMode; import static org.elasticsearch.xpack.ml.utils.SecondaryAuthorizationUtils.useSecondaryAuthIfAvailable; public class InferencePipelineAggregationBuilder extends AbstractPipelineAggregationBuilder { @@ -267,7 +268,7 @@ public InferencePipelineAggregationBuilder rewrite(QueryRewriteContext context) loadedModel.set(model); boolean isLicensed = MachineLearningField.ML_API_FEATURE.check(licenseState) - || licenseState.isAllowedByLicense(model.getLicenseLevel()); + || featureCheckForMode(model.getLicenseLevel(), licenseState); if (isLicensed) { delegate.onResponse(null); } else {