<commit_before>/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License
 * 2.0; you may not use this file except in compliance with the Elastic License
 * 2.0.
 */

package org.elasticsearch.xpack.ml.action;

import org.elasticsearch.Version;
import org.elasticsearch.cluster.node.DiscoveryNode;
import org.elasticsearch.cluster.node.DiscoveryNodes;
import org.elasticsearch.common.settings.ClusterSettings;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.transport.TransportAddress;
import org.elasticsearch.common.unit.ByteSizeValue;
import org.elasticsearch.common.util.set.Sets;
import org.elasticsearch.test.ESTestCase;
import org.elasticsearch.xpack.core.ml.dataframe.DataFrameAnalyticsConfig;
import org.elasticsearch.xpack.core.ml.job.config.Job;
import org.elasticsearch.xpack.ml.MachineLearning;

import java.net.InetAddress;
import java.util.Collections;

import static org.elasticsearch.xpack.ml.MachineLearning.MAX_LAZY_ML_NODES;
import static org.elasticsearch.xpack.ml.MachineLearning.MAX_MACHINE_MEMORY_PERCENT;
import static org.elasticsearch.xpack.ml.MachineLearning.MAX_ML_NODE_SIZE;
import static org.elasticsearch.xpack.ml.MachineLearning.USE_AUTO_MACHINE_MEMORY_PERCENT;
import static org.hamcrest.Matchers.greaterThan;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.lessThanOrEqualTo;
import static org.hamcrest.Matchers.notNullValue;
import static org.hamcrest.Matchers.nullValue;

public class TransportMlInfoActionTests extends ESTestCase {

    public void testCalculateEffectiveMaxModelMemoryLimitWithoutMaxMlNodeSize() {

        int mlMemoryPercent = randomIntBetween(5, 90);
        long mlMachineMemory = randomLongBetween(2000000000L, 100000000000L);
        int numMlNodes = randomIntBetween(0, 10);
        int numNonMlNodes = randomIntBetween(0, 10);
        ClusterSettings clusterSettings = new ClusterSettings(
            Settings.builder().put(MAX_MACHINE_MEMORY_PERCENT.getKey(), mlMemoryPercent).build(),
            Sets.newHashSet(MAX_LAZY_ML_NODES, MAX_MACHINE_MEMORY_PERCENT, MAX_ML_NODE_SIZE, USE_AUTO_MACHINE_MEMORY_PERCENT));
        long totalMlMemoryBytes = numMlNodes * mlMachineMemory * mlMemoryPercent / 100;

        DiscoveryNodes nodes = randomNodes(numMlNodes, numNonMlNodes, mlMachineMemory);

        ByteSizeValue effectiveMaxModelMemoryLimit = TransportMlInfoAction.calculateEffectiveMaxModelMemoryLimit(clusterSettings, nodes);

        if (numMlNodes == 0) {
            // "Don't know"
            assertThat(effectiveMaxModelMemoryLimit, nullValue());
        } else {
            // Expect configured percentage of current node size (allowing for small rounding errors)
            assertThat(effectiveMaxModelMemoryLimit, notNullValue());
            assertThat(effectiveMaxModelMemoryLimit.getBytes()
                    + Math.max(Job.PROCESS_MEMORY_OVERHEAD.getBytes(), DataFrameAnalyticsConfig.PROCESS_MEMORY_OVERHEAD.getBytes())
                    + MachineLearning.NATIVE_EXECUTABLE_CODE_OVERHEAD.getBytes(),
                lessThanOrEqualTo(mlMachineMemory * mlMemoryPercent / 100));
        }

        ByteSizeValue totalMlMemory = TransportMlInfoAction.calculateTotalMlMemory(clusterSettings, nodes);

        assertThat(totalMlMemory, notNullValue());
        assertThat(totalMlMemory, is(ByteSizeValue.ofMb(totalMlMemoryBytes / (1024 * 1024))));
    }

    public void testCalculateEffectiveMaxModelMemoryLimitNoMlNodesButMaxMlNodeSizeAndLazyNodesAllowed() {

        int mlMemoryPercent = randomIntBetween(5, 90);
        long mlMaxNodeSize = randomLongBetween(2000000000L, 100000000000L);
        int numNonMlNodes = randomIntBetween(0, 10);
        ClusterSettings clusterSettings = new ClusterSettings(
            Settings.builder().put(MAX_ML_NODE_SIZE.getKey(), mlMaxNodeSize + "b")
                .put(MAX_LAZY_ML_NODES.getKey(), randomIntBetween(1, 100))
                .put(MAX_MACHINE_MEMORY_PERCENT.getKey(), mlMemoryPercent).build(),
            Sets.newHashSet(MAX_LAZY_ML_NODES, MAX_MACHINE_MEMORY_PERCENT, MAX_ML_NODE_SIZE, USE_AUTO_MACHINE_MEMORY_PERCENT));

        DiscoveryNodes nodes = randomNodes(0, numNonMlNodes, 0);

        ByteSizeValue effectiveMaxModelMemoryLimit = TransportMlInfoAction.calculateEffectiveMaxModelMemoryLimit(clusterSettings, nodes);

        // Expect configured percentage of maximum declared node size (allowing for small rounding errors)
        assertThat(effectiveMaxModelMemoryLimit, notNullValue());
        assertThat(effectiveMaxModelMemoryLimit.getBytes()
                + Math.max(Job.PROCESS_MEMORY_OVERHEAD.getBytes(), DataFrameAnalyticsConfig.PROCESS_MEMORY_OVERHEAD.getBytes())
                + MachineLearning.NATIVE_EXECUTABLE_CODE_OVERHEAD.getBytes(),
            lessThanOrEqualTo(mlMaxNodeSize * mlMemoryPercent / 100));

        ByteSizeValue totalMlMemory = TransportMlInfoAction.calculateTotalMlMemory(clusterSettings, nodes);

        assertThat(totalMlMemory, notNullValue());
        assertThat(totalMlMemory, is(ByteSizeValue.ofMb(0)));
    }

    public void testCalculateEffectiveMaxModelMemoryLimitSmallMlNodesButMaxMlNodeSizeBiggerAndLazyNodesAllowed() {

        int mlMemoryPercent = randomIntBetween(5, 90);
        long mlMaxNodeSize = randomLongBetween(2000000000L, 100000000000L);
        long mlMachineMemory = mlMaxNodeSize / randomLongBetween(3, 5);
        int numMlNodes = randomIntBetween(1, 10);
        int numNonMlNodes = randomIntBetween(0, 10);
        ClusterSettings clusterSettings = new ClusterSettings(
            Settings.builder().put(MAX_ML_NODE_SIZE.getKey(), mlMaxNodeSize + "b")
                .put(MAX_LAZY_ML_NODES.getKey(), randomIntBetween(numMlNodes + 1, 100))
                .put(MAX_MACHINE_MEMORY_PERCENT.getKey(), mlMemoryPercent).build(),
            Sets.newHashSet(MAX_LAZY_ML_NODES, MAX_MACHINE_MEMORY_PERCENT, MAX_ML_NODE_SIZE, USE_AUTO_MACHINE_MEMORY_PERCENT));
        long totalMlMemoryBytes = numMlNodes * mlMachineMemory * mlMemoryPercent / 100;

        DiscoveryNodes nodes = randomNodes(numMlNodes, numNonMlNodes, mlMachineMemory);

        ByteSizeValue effectiveMaxModelMemoryLimit = TransportMlInfoAction.calculateEffectiveMaxModelMemoryLimit(clusterSettings, nodes);

        // Expect configured percentage of maximum declared node size (allowing for small rounding errors) - bigger than current node size
        assertThat(effectiveMaxModelMemoryLimit, notNullValue());
        assertThat(effectiveMaxModelMemoryLimit.getBytes()
                + Math.max(Job.PROCESS_MEMORY_OVERHEAD.getBytes(), DataFrameAnalyticsConfig.PROCESS_MEMORY_OVERHEAD.getBytes())
                + MachineLearning.NATIVE_EXECUTABLE_CODE_OVERHEAD.getBytes(),
            lessThanOrEqualTo(mlMaxNodeSize * mlMemoryPercent / 100));
        assertThat(effectiveMaxModelMemoryLimit.getBytes()
                + Math.max(Job.PROCESS_MEMORY_OVERHEAD.getBytes(), DataFrameAnalyticsConfig.PROCESS_MEMORY_OVERHEAD.getBytes())
                + MachineLearning.NATIVE_EXECUTABLE_CODE_OVERHEAD.getBytes(),
            greaterThan(2 * mlMachineMemory * mlMemoryPercent / 100));

        ByteSizeValue totalMlMemory = TransportMlInfoAction.calculateTotalMlMemory(clusterSettings, nodes);

        assertThat(totalMlMemory, notNullValue());
        assertThat(totalMlMemory, is(ByteSizeValue.ofMb(totalMlMemoryBytes / (1024 * 1024))));
    }

    public void testCalculateEffectiveMaxModelMemoryLimitSmallMlNodesButMaxMlNodeSizeBiggerAndLazyNodesExhausted() {

        int mlMemoryPercent = randomIntBetween(5, 90);
        long mlMaxNodeSize = randomLongBetween(2000000000L, 100000000000L);
        long mlMachineMemory = mlMaxNodeSize / randomLongBetween(3, 5);<commit_msg>[ML] Fixing edge cases in TransportMlInfoActionTests (#71930)<commit_after>/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License
 * 2.0; you may not use this file except in compliance with the Elastic License
 * 2.0.
 */

package org.elasticsearch.xpack.ml.action;

import org.elasticsearch.Version;
import org.elasticsearch.cluster.node.DiscoveryNode;
import org.elasticsearch.cluster.node.DiscoveryNodes;
import org.elasticsearch.common.settings.ClusterSettings;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.transport.TransportAddress;
import org.elasticsearch.common.unit.ByteSizeValue;
import org.elasticsearch.common.util.set.Sets;
import org.elasticsearch.test.ESTestCase;
import org.elasticsearch.xpack.core.ml.dataframe.DataFrameAnalyticsConfig;
import org.elasticsearch.xpack.core.ml.job.config.Job;
import org.elasticsearch.xpack.ml.MachineLearning;

import java.net.InetAddress;
import java.util.Collections;

import static org.elasticsearch.xpack.ml.MachineLearning.MAX_LAZY_ML_NODES;
import static org.elasticsearch.xpack.ml.MachineLearning.MAX_MACHINE_MEMORY_PERCENT;
import static org.elasticsearch.xpack.ml.MachineLearning.MAX_ML_NODE_SIZE;
import static org.elasticsearch.xpack.ml.MachineLearning.USE_AUTO_MACHINE_MEMORY_PERCENT;
import static org.hamcrest.Matchers.greaterThan;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.lessThanOrEqualTo;
import static org.hamcrest.Matchers.notNullValue;
import static org.hamcrest.Matchers.nullValue;

public class TransportMlInfoActionTests extends ESTestCase {

    public void testCalculateEffectiveMaxModelMemoryLimitWithoutMaxMlNodeSize() {

        int mlMemoryPercent = randomIntBetween(5, 90);
        long mlMachineMemory = randomLongBetween(2000000000L, 100000000000L);
        int numMlNodes = randomIntBetween(0, 10);
        int numNonMlNodes = randomIntBetween(0, 10);
        ClusterSettings clusterSettings = new ClusterSettings(
            Settings.builder().put(MAX_MACHINE_MEMORY_PERCENT.getKey(), mlMemoryPercent).build(),
            Sets.newHashSet(MAX_LAZY_ML_NODES, MAX_MACHINE_MEMORY_PERCENT, MAX_ML_NODE_SIZE, USE_AUTO_MACHINE_MEMORY_PERCENT));
        long totalMlMemoryBytes = numMlNodes * mlMachineMemory * mlMemoryPercent / 100;

        DiscoveryNodes nodes = randomNodes(numMlNodes, numNonMlNodes, mlMachineMemory);

        ByteSizeValue effectiveMaxModelMemoryLimit = TransportMlInfoAction.calculateEffectiveMaxModelMemoryLimit(clusterSettings, nodes);

        if (numMlNodes == 0) {
            // "Don't know"
            assertThat(effectiveMaxModelMemoryLimit, nullValue());
        } else {
            // Expect configured percentage of current node size (allowing for small rounding errors)
            assertThat(effectiveMaxModelMemoryLimit, notNullValue());
            assertThat(effectiveMaxModelMemoryLimit.getBytes()
                    + Math.max(Job.PROCESS_MEMORY_OVERHEAD.getBytes(), DataFrameAnalyticsConfig.PROCESS_MEMORY_OVERHEAD.getBytes())
                    + MachineLearning.NATIVE_EXECUTABLE_CODE_OVERHEAD.getBytes(),
                lessThanOrEqualTo(mlMachineMemory * mlMemoryPercent / 100));
        }

        ByteSizeValue totalMlMemory = TransportMlInfoAction.calculateTotalMlMemory(clusterSettings, nodes);

        assertThat(totalMlMemory, notNullValue());
        assertThat(totalMlMemory, is(ByteSizeValue.ofMb(totalMlMemoryBytes / (1024 * 1024))));
    }

    public void testCalculateEffectiveMaxModelMemoryLimitNoMlNodesButMaxMlNodeSizeAndLazyNodesAllowed() {

        int mlMemoryPercent = randomIntBetween(5, 90);
        long mlMaxNodeSize = randomLongBetween(2000000000L, 100000000000L);
        int numNonMlNodes = randomIntBetween(0, 10);
        ClusterSettings clusterSettings = new ClusterSettings(
            Settings.builder().put(MAX_ML_NODE_SIZE.getKey(), mlMaxNodeSize + "b")
                .put(MAX_LAZY_ML_NODES.getKey(), randomIntBetween(1, 100))
                .put(MAX_MACHINE_MEMORY_PERCENT.getKey(), mlMemoryPercent).build(),
            Sets.newHashSet(MAX_LAZY_ML_NODES, MAX_MACHINE_MEMORY_PERCENT, MAX_ML_NODE_SIZE, USE_AUTO_MACHINE_MEMORY_PERCENT));

        DiscoveryNodes nodes = randomNodes(0, numNonMlNodes, 0);

        ByteSizeValue effectiveMaxModelMemoryLimit = TransportMlInfoAction.calculateEffectiveMaxModelMemoryLimit(clusterSettings, nodes);

        // Expect configured percentage of maximum declared node size (allowing for small rounding errors)
        assertThat(effectiveMaxModelMemoryLimit, notNullValue());
        assertThat(effectiveMaxModelMemoryLimit.getBytes()
                + Math.max(Job.PROCESS_MEMORY_OVERHEAD.getBytes(), DataFrameAnalyticsConfig.PROCESS_MEMORY_OVERHEAD.getBytes())
                + MachineLearning.NATIVE_EXECUTABLE_CODE_OVERHEAD.getBytes(),
            lessThanOrEqualTo(mlMaxNodeSize * mlMemoryPercent / 100));

        ByteSizeValue totalMlMemory = TransportMlInfoAction.calculateTotalMlMemory(clusterSettings, nodes);

        assertThat(totalMlMemory, notNullValue());
        assertThat(totalMlMemory, is(ByteSizeValue.ofMb(0)));
    }

    public void testCalculateEffectiveMaxModelMemoryLimitSmallMlNodesButMaxMlNodeSizeBiggerAndLazyNodesAllowed() {

        int mlMemoryPercent = randomIntBetween(5, 90);
        long mlMaxNodeSize = randomLongBetween(2000000000L, 100000000000L);
        long mlMachineMemory = mlMaxNodeSize / randomLongBetween(3, 5);
        int numMlNodes = randomIntBetween(1, 10);
        int numNonMlNodes = randomIntBetween(0, 10);
        ClusterSettings clusterSettings = new ClusterSettings(
            Settings.builder().put(MAX_ML_NODE_SIZE.getKey(), mlMaxNodeSize + "b")
                .put(MAX_LAZY_ML_NODES.getKey(), randomIntBetween(numMlNodes + 1, 100))
                .put(MAX_MACHINE_MEMORY_PERCENT.getKey(), mlMemoryPercent).build(),
            Sets.newHashSet(MAX_LAZY_ML_NODES, MAX_MACHINE_MEMORY_PERCENT, MAX_ML_NODE_SIZE, USE_AUTO_MACHINE_MEMORY_PERCENT));
        long totalMlMemoryBytes = numMlNodes * mlMachineMemory * mlMemoryPercent / 100;

        DiscoveryNodes nodes = randomNodes(numMlNodes, numNonMlNodes, mlMachineMemory);

        ByteSizeValue effectiveMaxModelMemoryLimit = TransportMlInfoAction.calculateEffectiveMaxModelMemoryLimit(clusterSettings, nodes);

        // Expect configured percentage of maximum declared node size (allowing for small rounding errors) - bigger than current node size
        assertThat(effectiveMaxModelMemoryLimit, notNullValue());
        assertThat(effectiveMaxModelMemoryLimit.getBytes()
                + Math.max(Job.PROCESS_MEMORY_OVERHEAD.getBytes(), DataFrameAnalyticsConfig.PROCESS_MEMORY_OVERHEAD.getBytes())
                + MachineLearning.NATIVE_EXECUTABLE_CODE_OVERHEAD.getBytes(),
            lessThanOrEqualTo(mlMaxNodeSize * mlMemoryPercent / 100));
        assertThat(effectiveMaxModelMemoryLimit.getBytes()
                + Math.max(Job.PROCESS_MEMORY_OVERHEAD.getBytes(), DataFrameAnalyticsConfig.PROCESS_MEMORY_OVERHEAD.getBytes())
                + MachineLearning.NATIVE_EXECUTABLE_CODE_OVERHEAD.getBytes(),
            greaterThan(2 * mlMachineMemory * mlMemoryPercent / 100));

        ByteSizeValue totalMlMemory = TransportMlInfoAction.calculateTotalMlMemory(clusterSettings, nodes);

        assertThat(totalMlMemory, notNullValue());
        assertThat(totalMlMemory, is(ByteSizeValue.ofMb(totalMlMemoryBytes / (1024 * 1024))));
    }

    public void testCalculateEffectiveMaxModelMemoryLimitSmallMlNodesButMaxMlNodeSizeBiggerAndLazyNodesExhausted() {

        int mlMemoryPercent = randomIntBetween(10, 90);
        long mlMaxNodeSize = randomLongBetween(2000000000L, 100000000000L);
        long mlMachineMemory = mlMaxNodeSize / randomLongBetween(3, 4);
{'fraction_non_alphanumeric': 0.0804100082765646, 'fraction_numerical': 0.02202839498312854, 'mean_word_length': 4.425906735751295, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 16, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1798809', 'n_tokens_mistral': 4851, 'n_tokens_neox': 4481, 'n_words': 741}