[ML] track feature usage for data frame analytics, inference, and ano…

…maly jobs (#76789) This adds feature tracking for machine learning features. Model Snapshot upgrader Anomaly jobs Data frame analytics jobs Can all take advantage of the license state tracking built for persistent tasks. The ModelLoadingService needed special handling to ensure that models cached and referenced by pipelines are tracked. License tracking is done per-node and allows for a simple view into when a feature was last used on a given node.
elastic · Sep 1, 2021 · 2d70114 · 2d70114
1 parent 100f222
commit 2d70114
Show file tree

Hide file tree

Showing 20 changed files with 454 additions and 85 deletions.
diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/license/GetFeatureUsageResponse.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/license/GetFeatureUsageResponse.java
@@ -27,11 +27,11 @@
 public class GetFeatureUsageResponse extends ActionResponse implements ToXContentObject {
 
     public static class FeatureUsageInfo implements Writeable {
-        final String family;
-        final String name;
-        final ZonedDateTime lastUsedTime;
-        final String context;
-        final String licenseLevel;
+        private final String family;
+        private final String name;
+        private final ZonedDateTime lastUsedTime;
+        private final String context;
+        private final String licenseLevel;
 
         public FeatureUsageInfo(@Nullable String family, String name, ZonedDateTime lastUsedTime,
                                 @Nullable String context, String licenseLevel) {
@@ -70,6 +70,26 @@ public void writeTo(StreamOutput out) throws IOException {
             }
             out.writeString(licenseLevel);
         }
+
+        public String getFamily() {
+            return family;
+        }
+
+        public String getName() {
+            return name;
+        }
+
+        public ZonedDateTime getLastUsedTime() {
+            return lastUsedTime;
+        }
+
+        public String getContext() {
+            return context;
+        }
+
+        public String getLicenseLevel() {
+            return licenseLevel;
+        }
     }
 
     private List<FeatureUsageInfo> features;

diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/license/LicensedFeature.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/license/LicensedFeature.java
@@ -76,10 +76,10 @@ public void stopTracking(XPackLicenseState state, String contextName) {
         }
     }
 
-    final String family;
-    final String name;
-    final License.OperationMode minimumOperationMode;
-    final boolean needsActive;
+    private final String family;
+    private final String name;
+    private final License.OperationMode minimumOperationMode;
+    private final boolean needsActive;
 
     protected LicensedFeature(String family, String name, License.OperationMode minimumOperationMode, boolean needsActive) {
         this.family = family;
@@ -88,6 +88,22 @@ protected LicensedFeature(String family, String name, License.OperationMode mini
         this.needsActive = needsActive;
     }
 
+    public String getFamily() {
+        return family;
+    }
+
+    public String getName() {
+        return name;
+    }
+
+    public License.OperationMode getMinimumOperationMode() {
+        return minimumOperationMode;
+    }
+
+    public boolean isNeedsActive() {
+        return needsActive;
+    }
+
     /** Create a momentary feature for hte given license level */
     public static Momentary momentary(String family, String name, License.OperationMode licenseLevel) {
         return new Momentary(family, name, licenseLevel, true);

diff --git a/...k/plugin/core/src/main/java/org/elasticsearch/license/TransportGetFeatureUsageAction.java b/...k/plugin/core/src/main/java/org/elasticsearch/license/TransportGetFeatureUsageAction.java
@@ -45,11 +45,11 @@ protected void doExecute(Task task, GetFeatureUsageRequest request, ActionListen
             ZonedDateTime lastUsedTime = Instant.ofEpochMilli(lastUsed).atZone(ZoneOffset.UTC);
             usageInfos.add(
                 new GetFeatureUsageResponse.FeatureUsageInfo(
-                    usage.feature().family,
-                    usage.feature().name,
+                    usage.feature().getFamily(),
+                    usage.feature().getName(),
                     lastUsedTime,
                     usage.contextName(),
-                    usage.feature().minimumOperationMode.description()
+                    usage.feature().getMinimumOperationMode().description()
                 )
             );
         });

diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/license/XPackLicenseState.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/license/XPackLicenseState.java
@@ -511,7 +511,7 @@ public boolean isAllowed(Feature feature) {
 
     // Package protected: Only allowed to be called by LicensedFeature
     boolean isAllowed(LicensedFeature feature) {
-        if (isAllowedByLicense(feature.minimumOperationMode, feature.needsActive)) {
+        if (isAllowedByLicense(feature.getMinimumOperationMode(), feature.isNeedsActive())) {
             return true;
         }
         return false;
@@ -521,7 +521,7 @@ private void checkForExpiry(LicensedFeature feature) {
         final long licenseExpiryDate = getLicenseExpiryDate();
         // TODO: this should use epochMillisProvider to avoid a system call + testability
         final long diff = licenseExpiryDate - System.currentTimeMillis();
-        if (feature.minimumOperationMode.compareTo(OperationMode.BASIC) > 0 &&
+        if (feature.getMinimumOperationMode().compareTo(OperationMode.BASIC) > 0 &&
             LICENSE_EXPIRATION_WARNING_PERIOD.getMillis() > diff) {
             final long days = TimeUnit.MILLISECONDS.toDays(diff);
             final String expiryMessage = (days == 0 && diff > 0)? "expires today":
@@ -640,7 +640,7 @@ private FeatureUsage(LicensedFeature feature, String context) {
 
         @Override
         public String toString() {
-            return context == null ? feature.name : feature.name + ":" + context;
+            return context == null ? feature.getName() : feature.getName() + ":" + context;
         }
 
         @Override

diff --git a/...ugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfig.java b/...ugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfig.java
@@ -192,6 +192,8 @@ public static TrainedModelConfig.Builder fromXContent(XContentParser parser, boo
         }
         this.estimatedOperations = estimatedOperations;
         this.licenseLevel = License.OperationMode.parse(ExceptionsHelper.requireNonNull(licenseLevel, LICENSE_LEVEL));
+        assert this.licenseLevel.equals(License.OperationMode.PLATINUM) || this.licenseLevel.equals(License.OperationMode.BASIC) :
+                "[" + LICENSE_LEVEL.getPreferredName() + "] only [platinum] or [basic] is supported";
         this.defaultFieldMap = defaultFieldMap == null ? null : Collections.unmodifiableMap(defaultFieldMap);
         this.inferenceConfig = inferenceConfig;
         this.location = location;
@@ -330,6 +332,7 @@ public long getEstimatedOperations() {
         return estimatedOperations;
     }
 
+    //TODO if we ever support anything other than "basic" and platinum, we need to adjust our feature tracking logic
     public License.OperationMode getLicenseLevel() {
         return licenseLevel;
     }

diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/license/GetFeatureUsageResponseTests.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/license/GetFeatureUsageResponseTests.java
@@ -37,12 +37,12 @@ public void assertStreamInputOutput(Version version, String family, String conte
         GetFeatureUsageResponse finalResponse = new GetFeatureUsageResponse(input);
         assertThat(finalResponse.getFeatures(), hasSize(1));
         FeatureUsageInfo fui2 = finalResponse.getFeatures().get(0);
-        assertThat(fui2.family, equalTo(family));
-        assertThat(fui2.name, equalTo("feature"));
+        assertThat(fui2.getFamily(), equalTo(family));
+        assertThat(fui2.getName(), equalTo("feature"));
         // time is truncated to nearest second
-        assertThat(fui2.lastUsedTime, equalTo(zdt.withZoneSameInstant(ZoneOffset.UTC).withNano(0)));
-        assertThat(fui2.context, equalTo(context));
-        assertThat(fui2.licenseLevel, equalTo("gold"));
+        assertThat(fui2.getLastUsedTime(), equalTo(zdt.withZoneSameInstant(ZoneOffset.UTC).withNano(0)));
+        assertThat(fui2.getContext(), equalTo(context));
+        assertThat(fui2.getLicenseLevel(), equalTo("gold"));
     }
 
     public void testPre715StreamFormat() throws IOException {

diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/license/XPackLicenseStateTests.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/license/XPackLicenseStateTests.java
@@ -459,7 +459,7 @@ public void testLastUsedMomentaryFeature() {
         assertThat("feature.check tracks usage", lastUsed, aMapWithSize(1));
 
         XPackLicenseState.FeatureUsage usage = Iterables.get(lastUsed.keySet(), 0);
-        assertThat(usage.feature().name, equalTo("goldFeature"));
+        assertThat(usage.feature().getName(), equalTo("goldFeature"));
         assertThat(usage.contextName(), nullValue());
         assertThat(lastUsed.get(usage), equalTo(100L));
 
@@ -487,7 +487,7 @@ public void testLastUsedPersistentFeature() {
         assertThat(lastUsed, aMapWithSize(1));
 
         XPackLicenseState.FeatureUsage usage = Iterables.get(lastUsed.keySet(), 0);
-        assertThat(usage.feature().name, equalTo("goldFeature"));
+        assertThat(usage.feature().getName(), equalTo("goldFeature"));
         assertThat(usage.contextName(), equalTo("somecontext"));
         assertThat(lastUsed.get(usage), equalTo(200L));
 

diff --git a/...core/src/test/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfigTests.java b/...core/src/test/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelConfigTests.java
@@ -69,8 +69,6 @@ public static TrainedModelConfig.Builder createTestInstance(String modelId) {
             .setEstimatedHeapMemory(randomNonNegativeLong())
             .setEstimatedOperations(randomNonNegativeLong())
             .setLicenseLevel(randomFrom(License.OperationMode.PLATINUM.description(),
-                License.OperationMode.ENTERPRISE.description(),
-                License.OperationMode.GOLD.description(),
                 License.OperationMode.BASIC.description()))
             .setInferenceConfig(randomFrom(ClassificationConfigTests.randomClassificationConfig(),
                 RegressionConfigTests.randomRegressionConfig()))

diff --git a/...ternalClusterTest/java/org/elasticsearch/xpack/ml/integration/ModelInferenceActionIT.java b/...ternalClusterTest/java/org/elasticsearch/xpack/ml/integration/ModelInferenceActionIT.java
@@ -355,7 +355,7 @@ public void testInferMissingFields() throws Exception {
         }
     }
 
-    private static TrainedModelConfig.Builder buildTrainedModelConfigBuilder(String modelId) {
+    static TrainedModelConfig.Builder buildTrainedModelConfigBuilder(String modelId) {
         return TrainedModelConfig.builder()
             .setCreatedBy("ml_test")
             .setParsedDefinition(TrainedModelDefinitionTests.createRandomBuilder())