From ded18eeaac82787ea57d4991bad28b983d43ad0c Mon Sep 17 00:00:00 2001
From: Dima Arnautov <dmitrii.arnautov@elastic.co>
Date: Tue, 3 Dec 2024 15:29:28 +0100
Subject: [PATCH] [ML] Trained Models: Fixes spaces sync to retrieve 10000
 models (#202712)

## Summary

The default page size for the /trained_models API is 100. As a result,
the spaces sync task only fetched the first 100 models, leaving the rest
unassigned to spaces and therefore invisible in the ML UI.

This PR increases the page size to 10,000 to ensure all models are
properly assigned to Kibana spaces.

### Checklist

- [ ] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios
---
 x-pack/plugins/ml/common/constants/trained_models.ts | 12 ++++++++++++
 .../models/data_frame_analytics/analytics_manager.ts |  2 +-
 x-pack/plugins/ml/server/routes/trained_models.ts    |  3 +--
 x-pack/plugins/ml/server/saved_objects/util.ts       |  5 ++++-
 4 files changed, 18 insertions(+), 4 deletions(-)
 create mode 100644 x-pack/plugins/ml/common/constants/trained_models.ts

diff --git a/x-pack/plugins/ml/common/constants/trained_models.ts b/x-pack/plugins/ml/common/constants/trained_models.ts
new file mode 100644
index 0000000000000..c6092037cbfb4
--- /dev/null
+++ b/x-pack/plugins/ml/common/constants/trained_models.ts
@@ -0,0 +1,12 @@
+/*
+ * 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.
+ */
+
+/**
+ * Default page for the trained_models endpoint is 100,
+ * which is too small for the most cases, so we set it to 10000.
+ */
+export const DEFAULT_TRAINED_MODELS_PAGE_SIZE = 10000;
diff --git a/x-pack/plugins/ml/server/models/data_frame_analytics/analytics_manager.ts b/x-pack/plugins/ml/server/models/data_frame_analytics/analytics_manager.ts
index 4e112a2ff313b..e720f12fa4dd5 100644
--- a/x-pack/plugins/ml/server/models/data_frame_analytics/analytics_manager.ts
+++ b/x-pack/plugins/ml/server/models/data_frame_analytics/analytics_manager.ts
@@ -20,6 +20,7 @@ import {
 } from '@kbn/ml-data-frame-analytics-utils';
 import { isPopulatedObject } from '@kbn/ml-is-populated-object';
 import type { CloudSetup } from '@kbn/cloud-plugin/server';
+import { DEFAULT_TRAINED_MODELS_PAGE_SIZE } from '../../../common/constants/trained_models';
 import type { MlFeatures } from '../../../common/constants/app';
 import type { ModelService } from '../model_management/models_provider';
 import { modelsProvider } from '../model_management';
@@ -38,7 +39,6 @@ import {
   isTransformLinkReturnType,
 } from './types';
 import type { MlClient } from '../../lib/ml_client';
-import { DEFAULT_TRAINED_MODELS_PAGE_SIZE } from '../../routes/trained_models';
 
 export class AnalyticsManager {
   private _trainedModels: estypes.MlTrainedModelConfig[] = [];
diff --git a/x-pack/plugins/ml/server/routes/trained_models.ts b/x-pack/plugins/ml/server/routes/trained_models.ts
index c0010777ecf18..2782c4be18207 100644
--- a/x-pack/plugins/ml/server/routes/trained_models.ts
+++ b/x-pack/plugins/ml/server/routes/trained_models.ts
@@ -17,6 +17,7 @@ import type {
 } from '@kbn/ml-trained-models-utils';
 import { isDefined } from '@kbn/ml-is-defined';
 import type { IScopedClusterClient } from '@kbn/core-elasticsearch-server';
+import { DEFAULT_TRAINED_MODELS_PAGE_SIZE } from '../../common/constants/trained_models';
 import { type MlFeatures, ML_INTERNAL_BASE_PATH } from '../../common/constants/app';
 import type { RouteInitialization } from '../types';
 import { wrapError } from '../client/error_wrapper';
@@ -44,8 +45,6 @@ import { mlLog } from '../lib/log';
 import { forceQuerySchema } from './schemas/anomaly_detectors_schema';
 import { modelsProvider } from '../models/model_management';
 
-export const DEFAULT_TRAINED_MODELS_PAGE_SIZE = 10000;
-
 export function filterForEnabledFeatureModels<
   T extends TrainedModelConfigResponse | estypes.MlTrainedModelConfig
 >(models: T[], enabledFeatures: MlFeatures) {
diff --git a/x-pack/plugins/ml/server/saved_objects/util.ts b/x-pack/plugins/ml/server/saved_objects/util.ts
index 27562e919b79f..7f3f087ea61be 100644
--- a/x-pack/plugins/ml/server/saved_objects/util.ts
+++ b/x-pack/plugins/ml/server/saved_objects/util.ts
@@ -12,6 +12,7 @@ import {
   type IScopedClusterClient,
   SavedObjectsClient,
 } from '@kbn/core/server';
+import { DEFAULT_TRAINED_MODELS_PAGE_SIZE } from '../../common/constants/trained_models';
 import type { TrainedModelJob, MLSavedObjectService } from './service';
 import { ML_JOB_SAVED_OBJECT_TYPE } from '../../common/types/saved_objects';
 
@@ -86,7 +87,9 @@ export function mlFunctionsFactory(client: IScopedClusterClient) {
     },
     async getTrainedModels() {
       try {
-        return await client.asInternalUser.ml.getTrainedModels();
+        return await client.asInternalUser.ml.getTrainedModels({
+          size: DEFAULT_TRAINED_MODELS_PAGE_SIZE,
+        });
       } catch (error) {
         return null;
       }