[ML] Trained models: Update the start model allocation response type (e…

…lastic#210966) Updates the type correctly, as it is currently set to `{acknowledge: boolean}`, which is the wrong type.
JoseLuisGJ · Feb 13, 2025 · 26548ae · 26548ae
1 parent 18ef744
commit 26548ae
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 4 deletions.
diff --git a/x-pack/platform/plugins/shared/ml/common/types/trained_models.ts b/x-pack/platform/plugins/shared/ml/common/types/trained_models.ts
@@ -168,6 +168,17 @@ export type TrainedModelDeploymentStatsResponse = estypes.MlTrainedModelDeployme
   };
 };
 
+export interface StartTrainedModelDeploymentResponse {
+  // TODO update types in elasticsearch-specification
+  assignment: estypes.MlStartTrainedModelDeploymentResponse['assignment'] & {
+    adaptive_allocations?: {
+      enabled: boolean;
+      min_number_of_allocations?: number;
+      max_number_of_allocations?: number;
+    };
+  };
+}
+
 export interface AllocatedModel {
   key: string;
   deployment_id: string;

diff --git a/...form/plugins/shared/ml/public/application/model_management/trained_models_service.test.ts b/...form/plugins/shared/ml/public/application/model_management/trained_models_service.test.ts
@@ -12,7 +12,10 @@ import type {
   TrainedModelsApiService,
 } from '../services/ml_api_service/trained_models';
 import { TrainedModelsService } from './trained_models_service';
-import type { TrainedModelUIItem } from '../../../common/types/trained_models';
+import type {
+  StartTrainedModelDeploymentResponse,
+  TrainedModelUIItem,
+} from '../../../common/types/trained_models';
 import { MODEL_STATE } from '@kbn/ml-trained-models-utils';
 import { i18n } from '@kbn/i18n';
 import type { MlTrainedModelConfig } from '@elastic/elasticsearch/lib/api/types';
@@ -185,7 +188,39 @@ describe('TrainedModelsService', () => {
 
     mockTrainedModelsApiService.getTrainedModelsList.mockResolvedValueOnce([mockModel]);
 
-    mockTrainedModelsApiService.startModelAllocation.mockReturnValueOnce(of({ acknowledge: true }));
+    mockTrainedModelsApiService.startModelAllocation.mockReturnValueOnce(
+      of({
+        assignment: {
+          task_parameters: {
+            model_id: 'deploy-model',
+            model_bytes: 1000,
+            allocation_id: 'test-allocation',
+            priority: 'normal',
+            number_of_allocations: 1,
+            threads_per_allocation: 1,
+            queue_capacity: 1024,
+            deployment_id: 'my-deployment-id',
+            cache_size: '1mb',
+          },
+          node_count: 1,
+          routing_table: {
+            'node-1': {
+              routing_state: 'started',
+              reason: '',
+              current_allocations: 1,
+              target_allocations: 1,
+            },
+          },
+          assignment_state: 'started',
+          start_time: 1234567890,
+          adaptive_allocations: {
+            enabled: true,
+            min_number_of_allocations: 1,
+            max_number_of_allocations: 4,
+          },
+        },
+      })
+    );
 
     // Start deployment
     trainedModelsService.startModelDeployment('deploy-model', {
@@ -230,7 +265,9 @@ describe('TrainedModelsService', () => {
     const deploymentError = new Error('Deployment error');
 
     mockTrainedModelsApiService.startModelAllocation.mockReturnValueOnce(
-      throwError(() => deploymentError) as unknown as Observable<{ acknowledge: boolean }>
+      throwError(
+        () => deploymentError
+      ) as unknown as Observable<StartTrainedModelDeploymentResponse>
     );
 
     trainedModelsService.startModelDeployment('error-model', {

diff --git a/...k/platform/plugins/shared/ml/public/application/services/ml_api_service/trained_models.ts b/...k/platform/plugins/shared/ml/public/application/services/ml_api_service/trained_models.ts
@@ -27,6 +27,7 @@ import type {
   ModelDownloadState,
   TrainedModelUIItem,
   TrainedModelConfigResponse,
+  StartTrainedModelDeploymentResponse,
 } from '../../../../common/types/trained_models';
 
 export interface InferenceQueryParams {
@@ -238,7 +239,7 @@ export function trainedModelsApiProvider(httpService: HttpService) {
       deploymentParams,
       adaptiveAllocationsParams,
     }: StartAllocationParams) {
-      return httpService.http$<{ acknowledge: boolean }>({
+      return httpService.http$<StartTrainedModelDeploymentResponse>({
         path: `${ML_INTERNAL_BASE_PATH}/trained_models/${modelId}/deployment/_start`,
         method: 'POST',
         query: deploymentParams,