import type * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey' ; import numeral from '@elastic/numeral' ; import { IScopedClusterClient } from 'kibana/server' ; import { MLCATEGORY } from '../../../common/constants/field_types' ; import { AnalysisConfig , Datafeed } from '../../../common/types/anomaly_detection_jobs' ; import { fieldsServiceProvider } from '../fields_service' ; import { MlInfoResponse } from '../../../common/types/ml_server_info' ; import type { MlClient } from '../../lib/ml_client' ; export interface ModelMemoryEstimationResult { /**
     * Result model memory limit
     */ modelMemoryLimit : string ; /**
     * Estimated model memory by elasticsearch ml endpoint
     */ estimatedModelMemoryLimit : string ; /**
     * Maximum model memory limit
     */ maxModelMemoryLimit ? : string ; } /**
 * Response of the _estimate_model_memory endpoint.
 */ export interface ModelMemoryEstimateResponse { model_memory_estimate : string ; } /**
 * Retrieves overall and max bucket cardinalities.
 */ const cardinalityCheckProvider = ( client ) => { const fieldsService = fieldsServiceProvider ( client ) ; return async ( analysisConfig , indexPattern , query , timeFieldName , earliestMs , latestMs , datafeedConfig ? ) : Promise < { overallCardinality : { [ key ] : number ; } ; maxBucketCardinality : { [ key ] : number ; } ; } > => { /**
         * Fields not involved in cardinality check
         */ const excludedKeywords = new Set < string > ( MLCATEGORY ) ; const { detectors , influencers , bucket_span : bucketSpan } = analysisConfig ; let overallCardinality = { } ; let maxBucketCardinality = { } ; const overallCardinalityFields = detectors . reduce ( ( acc , { by_field_name : byFieldName , partition_field_name : partitionFieldName , over_field_name : overFieldName , } ) => { [ byFieldName , partitionFieldName , overFieldName ] . filter ( ( field ) => field !== undefined && field !== '' && ! excludedKeywords . has ( field ) ) . forEach ( ( key ) => { acc . add ( key as string ) ; } ) ; return acc ; } , new Set < string > ( ) ) ; const normalizedInfluencers = Array . isArray ( influencers ) ? influencers : [ influencers ] ; const maxBucketFieldCardinalities = normalizedInfluencers . filter ( ( influencerField ) => ! ! influencerField && ! excludedKeywords . has ( influencerField ) && ! overallCardinalityFields . has ( influencerField ) ) ; if ( overallCardinalityFields . size > 0 ) { overallCardinality = await fieldsService . getCardinalityOfFields ( indexPattern , [ ... overallCardinalityFields ] , query , timeFieldName , earliestMs , latestMs , datafeedConfig ) ; } if ( maxBucketFieldCardinalities . length > 0 ) { maxBucketCardinality = await fieldsService . getMaxBucketCardinalities ( indexPattern , maxBucketFieldCardinalities , query , timeFieldName , earliestMs , latestMs , bucketSpan as string , datafeedConfig ) ; } return { overallCardinality , maxBucketCardinality , } ; } ; } ; export function calculateModelMemoryLimitProvider ( client , mlClient ) { const getCardinalities = cardinalityCheckProvider ( client ) ; /**
     * Retrieves an estimated size of the model memory limit used in the job config
     * based on the cardinality of the fields being used to split the data
     * and influencers.
     */ return async function calculateModelMemoryLimit ( analysisConfig , indexPattern , query , timeFieldName , earliestMs , latestMs , allowMMLGreaterThanMax = false , datafeedConfig ? ) : Promise < ModelMemoryEstimationResult > { const { body : info } = await mlClient . info < MlInfoResponse > ( ) ; const maxModelMemoryLimit = info . limits . max_model_memory_limit ?. toUpperCase ( ) ; const effectiveMaxModelMemoryLimit = info . limits . effective_max_model_memory_limit ?. toUpperCase ( ) ; const { overallCardinality , maxBucketCardinality } = await getCardinalities ( analysisConfig , indexPattern , query , timeFieldName , earliestMs , latestMs , datafeedConfig ) ; const { body } = await mlClient . estimateModelMemory < ModelMemoryEstimateResponse > ( { body : { analysis_config : analysisConfig , overall_cardinality : overallCardinality , max_bucket_cardinality : maxBucketCardinality , } , } ) ; const estimatedModelMemoryLimit = body . model_memory_estimate . toUpperCase ( ) ; let modelMemoryLimit = estimatedModelMemoryLimit ; let mmlCappedAtMax = false ; if ( allowMMLGreaterThanMax === false ) { const mmlBytes = numeral ( estimatedModelMemoryLimit ) . value ( ) ; if ( maxModelMemoryLimit !== undefined ) { const maxBytes = numeral ( maxModelMemoryLimit ) . value ( ) ; if ( mmlBytes > maxBytes ) { modelMemoryLimit = `${ Math . floor ( maxBytes / numeral ( '1MB' ) . value ( ) ) }MB` ; mmlCappedAtMax = true ; } } if ( mmlCappedAtMax === false && effectiveMaxModelMemoryLimit !== undefined ) { const effectiveMaxMmlBytes = numeral ( effectiveMaxModelMemoryLimit ) . value ( ) ; if ( mmlBytes > effectiveMaxMmlBytes ) { modelMemoryLimit = `${ Math . floor ( effectiveMaxMmlBytes / numeral ( '1MB' ) . value ( ) ) }MB` ; } } } return { estimatedModelMemoryLimit , modelMemoryLimit , ... ( maxModelMemoryLimit ? { maxModelMemoryLimit } : { } ) , ... ( effectiveMaxModelMemoryLimit ? { effectiveMaxModelMemoryLimit } : { } ) , } ; } ; }