[ML] Fix model size stats. (#121958)

The API for model size stats changed attributes and structure. This picks up the changes to fix the display of model sizes on the ML overview page and model management page.
elastic · Dec 23, 2021 · 058caf8 · 058caf8
1 parent ee13980
commit 058caf8
Show file tree

Hide file tree

Showing 6 changed files with 420 additions and 356 deletions.
diff --git a/x-pack/plugins/ml/common/types/trained_models.ts b/x-pack/plugins/ml/common/types/trained_models.ts
@@ -17,6 +17,11 @@ export interface IngestStats {
   failed: number;
 }
 
+export interface TrainedModelModelSizeStats {
+  model_size_bytes: number;
+  required_native_memory_bytes: number;
+}
+
 export interface TrainedModelStat {
   model_id?: string;
   pipeline_count?: number;
@@ -46,6 +51,7 @@ export interface TrainedModelStat {
     >;
   };
   deployment_stats?: Omit<TrainedModelDeploymentStatsResponse, 'model_id'>;
+  model_size_stats?: TrainedModelModelSizeStats;
 }
 
 type TreeNode = object;
@@ -126,7 +132,6 @@ export interface InferenceConfigResponse {
 
 export interface TrainedModelDeploymentStatsResponse {
   model_id: string;
-  model_size_bytes: number;
   inference_threads: number;
   model_threads: number;
   state: DeploymentState;
@@ -170,6 +175,7 @@ export interface AllocatedModel {
   state: string;
   model_threads: number;
   model_size_bytes: number;
+  required_native_memory_bytes: number;
   node: {
     /**
      * Not required for rendering in the Nodes overview

diff --git a/x-pack/plugins/ml/public/application/trained_models/models_management/expanded_row.tsx b/x-pack/plugins/ml/public/application/trained_models/models_management/expanded_row.tsx
@@ -122,13 +122,15 @@ export const ExpandedRow: FC<ExpandedRowProps> = ({ item }) => {
     function updateModelItems() {
       (async function () {
         const deploymentStats = stats.deployment_stats;
+        const modelSizeStats = stats.model_size_stats;
 
-        if (!deploymentStats) return;
+        if (!deploymentStats || !modelSizeStats) return;
 
         const items: AllocatedModel[] = deploymentStats.nodes.map((n) => {
           const nodeName = Object.values(n.node)[0].name;
           return {
             ...deploymentStats,
+            ...modelSizeStats,
             node: {
               ...pick(n, [
                 'average_inference_time_ms',

diff --git a/x-pack/plugins/ml/public/application/trained_models/nodes_overview/allocated_models.tsx b/x-pack/plugins/ml/public/application/trained_models/nodes_overview/allocated_models.tsx
@@ -59,7 +59,7 @@ export const AllocatedModels: FC<AllocatedModelsProps> = ({
       truncateText: true,
       'data-test-subj': 'mlAllocatedModelsTableSize',
       render: (v: AllocatedModel) => {
-        return bytesFormatter(v.model_size_bytes);
+        return bytesFormatter(v.required_native_memory_bytes);
       },
     },
     {