From ea06a3521f8ae6f788f02807569bc22b4bf45363 Mon Sep 17 00:00:00 2001
From: Dima Arnautov <dmitrii.arnautov@elastic.co>
Date: Fri, 24 Nov 2023 14:33:21 +0100
Subject: [PATCH] validate threading params

---
 .../plugins/ml/common/types/ml_server_info.ts |  2 +
 .../model_management/deployment_setup.tsx     | 55 ++++++++++++++-----
 2 files changed, 43 insertions(+), 14 deletions(-)

diff --git a/x-pack/plugins/ml/common/types/ml_server_info.ts b/x-pack/plugins/ml/common/types/ml_server_info.ts
index e5141d6f2e78f..215f46c26bef8 100644
--- a/x-pack/plugins/ml/common/types/ml_server_info.ts
+++ b/x-pack/plugins/ml/common/types/ml_server_info.ts
@@ -20,6 +20,8 @@ export interface MlServerDefaults {
 export interface MlServerLimits {
   max_model_memory_limit?: string;
   effective_max_model_memory_limit?: string;
+  max_single_ml_node_processors?: number;
+  total_ml_processors?: number;
 }
 
 export interface MlInfoResponse {
diff --git a/x-pack/plugins/ml/public/application/model_management/deployment_setup.tsx b/x-pack/plugins/ml/public/application/model_management/deployment_setup.tsx
index 3ac6960af55ca..f0c0898652a05 100644
--- a/x-pack/plugins/ml/public/application/model_management/deployment_setup.tsx
+++ b/x-pack/plugins/ml/public/application/model_management/deployment_setup.tsx
@@ -31,7 +31,7 @@ import type { I18nStart, OverlayStart, ThemeServiceStart } from '@kbn/core/publi
 import { css } from '@emotion/react';
 import { numberValidator } from '@kbn/ml-agg-utils';
 import { toMountPoint } from '@kbn/react-kibana-mount';
-import { isCloudTrial } from '../services/ml_server_info';
+import { getNewJobLimits, isCloudTrial } from '../services/ml_server_info';
 import {
   composeValidators,
   dictionaryValidator,
@@ -42,7 +42,7 @@ import { ModelItem } from './models_list';
 interface DeploymentSetupProps {
   config: ThreadingParams;
   onConfigChange: (config: ThreadingParams) => void;
-  errors: Partial<Record<keyof ThreadingParams, object>>;
+  errors: Partial<Record<keyof ThreadingParams, Record<string, unknown>>>;
   isUpdate?: boolean;
   deploymentsParams?: Record<string, ThreadingParams>;
 }
@@ -66,6 +66,11 @@ export const DeploymentSetup: FC<DeploymentSetupProps> = ({
   isUpdate,
   deploymentsParams,
 }) => {
+  const {
+    total_ml_processors: totalMlProcessors,
+    max_single_ml_node_processors: maxSingleMlNodeProcessors,
+  } = getNewJobLimits();
+
   const numOfAllocation = config.numOfAllocations;
   const threadsPerAllocations = config.threadsPerAllocations;
 
@@ -76,17 +81,20 @@ export const DeploymentSetup: FC<DeploymentSetupProps> = ({
 
   const threadsPerAllocationsOptions = useMemo(
     () =>
-      new Array(THREADS_MAX_EXPONENT).fill(null).map((v, i) => {
-        const value = Math.pow(2, i);
-        const id = value.toString();
-
-        return {
-          id,
-          label: id,
-          value,
-        };
-      }),
-    []
+      new Array(THREADS_MAX_EXPONENT)
+        .fill(null)
+        .map((v, i) => Math.pow(2, i))
+        .filter(maxSingleMlNodeProcessors ? (v) => v <= maxSingleMlNodeProcessors : (v) => v)
+        .map((value) => {
+          const id = value.toString();
+
+          return {
+            id,
+            label: id,
+            value,
+          };
+        }),
+    [maxSingleMlNodeProcessors]
   );
 
   const disableThreadingControls = config.priority === 'low';
@@ -251,11 +259,28 @@ export const DeploymentSetup: FC<DeploymentSetupProps> = ({
           }
           hasChildLabel={false}
           isDisabled={disableThreadingControls}
+          isInvalid={!!errors.numOfAllocations}
+          error={
+            errors?.numOfAllocations?.min ? (
+              <FormattedMessage
+                id="xpack.ml.trainedModels.modelsList.startDeployment.numbersOfAllocationsMinError"
+                defaultMessage="At least one allocation is required."
+              />
+            ) : errors?.numOfAllocations?.max ? (
+              <FormattedMessage
+                id="xpack.ml.trainedModels.modelsList.startDeployment.numbersOfAllocationsMaxError"
+                defaultMessage="Number of allocations cannot exceed {max}."
+                values={{ max: totalMlProcessors }}
+              />
+            ) : null
+          }
         >
           <EuiFieldNumber
             disabled={disableThreadingControls}
+            isInvalid={!!errors.numOfAllocations}
             fullWidth
             min={1}
+            max={totalMlProcessors}
             step={1}
             name={'numOfAllocations'}
             value={disableThreadingControls ? 1 : numOfAllocation}
@@ -346,6 +371,8 @@ export const StartUpdateDeploymentModal: FC<StartDeploymentModalProps> = ({
 }) => {
   const isUpdate = !!initialParams;
 
+  const { total_ml_processors: totalMlProcessors } = getNewJobLimits();
+
   const [config, setConfig] = useState<ThreadingParams>(
     initialParams ?? {
       numOfAllocations: 1,
@@ -373,7 +400,7 @@ export const StartUpdateDeploymentModal: FC<StartDeploymentModalProps> = ({
 
   const numOfAllocationsValidator = composeValidators(
     requiredValidator(),
-    numberValidator({ min: 1, integerOnly: true })
+    numberValidator({ min: 1, max: totalMlProcessors, integerOnly: true })
   );
 
   const numOfAllocationsErrors = numOfAllocationsValidator(config.numOfAllocations);