[ML] Set threading params for _start trained model deployment API (#135134)

* set threading params * model header * useMemo for threadsPerAllocationsOptions * remove magic number * id variable * rename modal wrapper function * remove placeholder * update legend text * min value for numOfAllocations * add validation * add extra text for cloud * change the layout * fix i18n message id * docs url * set header sizes * change messages * info callout * fix the scrollbar issue * change the message * move the docs link * rename the docs key * fix typo * change doc link alignment
2025-04-23 17:28:26 -04:00 · 2022-06-29 17:48:00 +02:00 · 2022-06-29 17:48:00 +02:00 · 6ff66841a7
commit 6ff66841a7
parent 0853a0b005
10 changed files with 415 additions and 4 deletions
--- a/packages/kbn-doc-links/src/get_doc_links.ts
+++ b/packages/kbn-doc-links/src/get_doc_links.ts
@ -395,6 +395,7 @@ export const getDocLinks = ({ kibanaBranch }: GetDocLinkOptions): DocLinks => {
      classificationAucRoc: `${ELASTIC_WEBSITE_URL}guide/en/machine-learning/${DOC_LINK_VERSION}/ml-dfa-classification.html#ml-dfanalytics-class-aucroc`,
      setUpgradeMode: `${ELASTICSEARCH_DOCS}ml-set-upgrade-mode.html`,
      trainedModels: `${ELASTIC_WEBSITE_URL}guide/en/machine-learning/${DOC_LINK_VERSION}/ml-trained-models.html`,
+      startTrainedModelsDeploymentQueryParams: `${ELASTIC_WEBSITE_URL}guide/en/elasticsearch/reference/${DOC_LINK_VERSION}/start-trained-model-deployment.html#start-trained-model-deployment-query-params`,
    },
    transforms: {
      guide: `${ELASTICSEARCH_DOCS}transforms.html`,
--- a/x-pack/plugins/ml/common/types/trained_models.ts
+++ b/x-pack/plugins/ml/common/types/trained_models.ts
@ -129,6 +129,8 @@ export interface TrainedModelDeploymentStatsResponse {
  inference_threads: number;
  model_threads: number;
  state: DeploymentState;
+  threads_per_allocation: number;
+  number_of_allocations: number;
  allocation_status: { target_allocation_count: number; state: string; allocation_count: number };
  nodes: Array<{
    node: Record<
@ -153,6 +155,8 @@ export interface TrainedModelDeploymentStatsResponse {
    number_of_pending_requests: number;
    start_time: number;
    throughput_last_minute: number;
+    threads_per_allocation: number;
+    number_of_allocations: number;
  }>;
 }

@ -163,6 +167,8 @@ export interface AllocatedModel {
    state: string;
    allocation_count: number;
  };
+  number_of_allocations: number;
+  threads_per_allocation: number;
  /**
   * Not required for rendering in the Model stats
   */
@ -186,6 +192,8 @@ export interface AllocatedModel {
    number_of_pending_requests: number;
    start_time: number;
    throughput_last_minute: number;
+    number_of_allocations: number;
+    threads_per_allocation: number;
  };
 }

--- a/x-pack/plugins/ml/common/util/validators.ts
+++ b/x-pack/plugins/ml/common/util/validators.ts
@ -104,9 +104,14 @@ export function timeIntervalInputValidator() {
 export interface NumberValidationResult {
  min: boolean;
  max: boolean;
+  integerOnly: boolean;
 }

-export function numberValidator(conditions?: { min?: number; max?: number }) {
+export function numberValidator(conditions?: {
+  min?: number;
+  max?: number;
+  integerOnly?: boolean;
+}) {
  if (
    conditions?.min !== undefined &&
    conditions.max !== undefined &&
@ -123,6 +128,9 @@ export function numberValidator(conditions?: { min?: number; max?: number }) {
    if (conditions?.max !== undefined && value > conditions.max) {
      result.max = true;
    }
+    if (!!conditions?.integerOnly && !Number.isInteger(value)) {
+      result.integerOnly = true;
+    }
    if (isPopulatedObject(result)) {
      return result;
    }
--- a/x-pack/plugins/ml/public/application/services/ml_api_service/trained_models.ts
+++ b/x-pack/plugins/ml/public/application/services/ml_api_service/trained_models.ts
@ -128,10 +128,14 @@ export function trainedModelsApiProvider(httpService: HttpService) {
      });
    },

-    startModelAllocation(modelId: string) {
+    startModelAllocation(
+      modelId: string,
+      queryParams?: { number_of_allocations: number; threads_per_allocation: number }
+    ) {
      return httpService.http<{ acknowledge: boolean }>({
        path: `${apiBasePath}/trained_models/${modelId}/deployment/_start`,
        method: 'POST',
+        query: queryParams,
      });
    },

--- 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
@ -156,6 +156,8 @@ export const ExpandedRow: FC<ExpandedRowProps> = ({ item }) => {
                'number_of_pending_requests',
                'start_time',
                'throughput_last_minute',
+                'number_of_allocations',
+                'threads_per_allocation',
              ]),
              name: nodeName,
            } as AllocatedModel['node'],
--- a/x-pack/plugins/ml/public/application/trained_models/models_management/models_list.tsx
+++ b/x-pack/plugins/ml/public/application/trained_models/models_management/models_list.tsx
@ -27,6 +27,7 @@ import { EuiTableSelectionType } from '@elastic/eui/src/components/basic_table/t
 import { Action } from '@elastic/eui/src/components/basic_table/action_types';
 import { FIELD_FORMAT_IDS } from '@kbn/field-formats-plugin/common';
 import { isPopulatedObject } from '@kbn/ml-is-populated-object';
+import { getUserInputThreadingParamsProvider } from './start_deployment_setup';
 import { getAnalysisType } from '../../data_frame_analytics/common';
 import { ModelsTableToConfigMapping } from '.';
 import { ModelsBarStats, StatsBar } from '../../components/stats_bar';
@ -94,10 +95,13 @@ export const ModelsList: FC<Props> = ({
      overlays,
      theme,
      spacesApi,
+      docLinks,
    },
  } = useMlKibana();
  const urlLocator = useMlLocator()!;

+  const startModelDeploymentDocUrl = docLinks.links.ml.startTrainedModelsDeploymentQueryParams;
+
  useTimefilter({ timeRangeSelector: false, autoRefreshSelector: true });

  const dateFormatter = useFieldFormatter(FIELD_FORMAT_IDS.DATE);
@ -139,6 +143,11 @@ export const ModelsList: FC<Props> = ({
  const [showTestFlyout, setShowTestFlyout] = useState<ModelItem | null>(null);
  const getUserConfirmation = useMemo(() => getUserConfirmationProvider(overlays, theme), []);

+  const getUserInputThreadingParams = useMemo(
+    () => getUserInputThreadingParamsProvider(overlays, theme.theme$, startModelDeploymentDocUrl),
+    [overlays, theme.theme$, startModelDeploymentDocUrl]
+  );
+
  const navigateToPath = useNavigateToPath();

  const isBuiltInModel = useCallback(
@ -369,9 +378,16 @@ export const ModelsList: FC<Props> = ({
          },
          available: (item) => item.model_type === TRAINED_MODEL_TYPE.PYTORCH,
          onClick: async (item) => {
+            const threadingParams = await getUserInputThreadingParams(item.model_id);
+
+            if (!threadingParams) return;
+
            try {
              setIsLoading(true);
-              await trainedModelsApiService.startModelAllocation(item.model_id);
+              await trainedModelsApiService.startModelAllocation(item.model_id, {
+                number_of_allocations: threadingParams.numOfAllocations,
+                threads_per_allocation: threadingParams.threadsPerAllocations,
+              });
              displaySuccessToast(
                i18n.translate('xpack.ml.trainedModels.modelsList.startSuccess', {
                  defaultMessage: 'Deployment for "{modelId}" has been started successfully.',
--- a/x-pack/plugins/ml/public/application/trained_models/models_management/start_deployment_setup.tsx
+++ b/x-pack/plugins/ml/public/application/trained_models/models_management/start_deployment_setup.tsx
@ -0,0 +1,340 @@
+/*
+ * 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.
+ */
+
+import React, { FC, useState, useMemo } from 'react';
+import { i18n } from '@kbn/i18n';
+import { FormattedMessage } from '@kbn/i18n-react';
+import {
+  EuiForm,
+  EuiButtonGroup,
+  EuiFormRow,
+  EuiFieldNumber,
+  EuiModal,
+  EuiModalHeader,
+  EuiModalHeaderTitle,
+  EuiModalBody,
+  EuiModalFooter,
+  EuiButtonEmpty,
+  EuiButton,
+  EuiCallOut,
+  EuiSpacer,
+  EuiDescribedFormGroup,
+  EuiFlexGroup,
+  EuiFlexItem,
+  EuiLink,
+  EuiTitle,
+} from '@elastic/eui';
+import { toMountPoint, wrapWithTheme } from '@kbn/kibana-react-plugin/public';
+import type { Observable } from 'rxjs';
+import type { CoreTheme, OverlayStart } from '@kbn/core/public';
+import { css } from '@emotion/react';
+import { isCloud } from '../../services/ml_server_info';
+import {
+  composeValidators,
+  numberValidator,
+  requiredValidator,
+} from '../../../../common/util/validators';
+
+interface StartDeploymentSetup {
+  config: ThreadingParams;
+  onConfigChange: (config: ThreadingParams) => void;
+}
+
+export interface ThreadingParams {
+  numOfAllocations: number;
+  threadsPerAllocations: number;
+}
+
+const THREADS_MAX_EXPONENT = 6;
+
+/**
+ * Form for setting threading params.
+ */
+export const StartDeploymentSetup: FC<StartDeploymentSetup> = ({ config, onConfigChange }) => {
+  const numOfAllocation = config.numOfAllocations;
+  const threadsPerAllocations = config.threadsPerAllocations;
+
+  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,
+        };
+      }),
+    []
+  );
+
+  const toggleIdSelected = threadsPerAllocationsOptions.find(
+    (v) => v.value === threadsPerAllocations
+  )!.id;
+
+  return (
+    <EuiForm component={'form'} id={'startDeploymentForm'}>
+      <EuiDescribedFormGroup
+        titleSize={'xxs'}
+        title={
+          <h3>
+            <FormattedMessage
+              id="xpack.ml.trainedModels.modelsList.startDeployment.numbersOfAllocationsLabel"
+              defaultMessage="Number of allocations"
+            />
+          </h3>
+        }
+        description={
+          <FormattedMessage
+            id="xpack.ml.trainedModels.modelsList.startDeployment.numbersOfAllocationsHelp"
+            defaultMessage="Increase to improve throughput of all requests."
+          />
+        }
+      >
+        <EuiFormRow
+          label={
+            <FormattedMessage
+              id="xpack.ml.trainedModels.modelsList.startDeployment.numbersOfAllocationsLabel"
+              defaultMessage="Number of allocations"
+            />
+          }
+          hasChildLabel={false}
+        >
+          <EuiFieldNumber
+            fullWidth
+            min={1}
+            step={1}
+            name={'numOfAllocations'}
+            value={numOfAllocation}
+            onChange={(event) => {
+              onConfigChange({ ...config, numOfAllocations: Number(event.target.value) });
+            }}
+          />
+        </EuiFormRow>
+      </EuiDescribedFormGroup>
+
+      <EuiDescribedFormGroup
+        titleSize={'xxs'}
+        title={
+          <h3>
+            <FormattedMessage
+              id="xpack.ml.trainedModels.modelsList.startDeployment.threadsPerAllocationLabel"
+              defaultMessage="Threads per allocation"
+            />
+          </h3>
+        }
+        description={
+          <FormattedMessage
+            id="xpack.ml.trainedModels.modelsList.startDeployment.threadsPerAllocationHelp"
+            defaultMessage="Increase to improve latency for each request."
+          />
+        }
+      >
+        <EuiFormRow
+          label={
+            <FormattedMessage
+              id="xpack.ml.trainedModels.modelsList.startDeployment.threadsPerAllocationLabel"
+              defaultMessage="Threads per allocation"
+            />
+          }
+          hasChildLabel={false}
+        >
+          <EuiButtonGroup
+            legend={i18n.translate(
+              'xpack.ml.trainedModels.modelsList.startDeployment.threadsPerAllocationLegend',
+              {
+                defaultMessage: 'Threads per allocation selector',
+              }
+            )}
+            name={'threadsPerAllocation'}
+            isFullWidth
+            idSelected={toggleIdSelected}
+            onChange={(optionId) => {
+              const value = threadsPerAllocationsOptions.find((v) => v.id === optionId)!.value;
+              onConfigChange({ ...config, threadsPerAllocations: value });
+            }}
+            options={threadsPerAllocationsOptions}
+          />
+        </EuiFormRow>
+      </EuiDescribedFormGroup>
+    </EuiForm>
+  );
+};
+
+interface StartDeploymentModalProps {
+  modelId: string;
+  startModelDeploymentDocUrl: string;
+  onConfigChange: (config: ThreadingParams) => void;
+  onClose: () => void;
+}
+
+/**
+ * Modal window wrapper for {@link StartDeploymentSetup}
+ *
+ * @param onConfigChange
+ * @param onClose
+ */
+export const StartDeploymentModal: FC<StartDeploymentModalProps> = ({
+  modelId,
+  onConfigChange,
+  onClose,
+  startModelDeploymentDocUrl,
+}) => {
+  const [config, setConfig] = useState<ThreadingParams>({
+    numOfAllocations: 1,
+    threadsPerAllocations: 1,
+  });
+
+  const numOfAllocationsValidator = composeValidators(
+    requiredValidator(),
+    numberValidator({ min: 1, integerOnly: true })
+  );
+
+  const errors = numOfAllocationsValidator(config.numOfAllocations);
+
+  return (
+    <EuiModal onClose={onClose} initialFocus="[name=numOfAllocations]" maxWidth={false}>
+      <EuiModalHeader>
+        <EuiModalHeaderTitle>
+          <EuiFlexGroup justifyContent={'spaceBetween'}>
+            <EuiFlexItem grow={false}>
+              <EuiTitle size={'s'}>
+                <h2>
+                  <FormattedMessage
+                    id="xpack.ml.trainedModels.modelsList.startDeployment.modalTitle"
+                    defaultMessage="Start {modelId} deployment"
+                    values={{ modelId }}
+                  />
+                </h2>
+              </EuiTitle>
+            </EuiFlexItem>
+
+            <EuiFlexItem grow={false} />
+          </EuiFlexGroup>
+        </EuiModalHeaderTitle>
+      </EuiModalHeader>
+
+      <EuiModalBody>
+        {isCloud() ? (
+          <>
+            <EuiCallOut
+              size={'s'}
+              title={
+                <FormattedMessage
+                  id="xpack.ml.trainedModels.modelsList.startDeployment.cloudWarningHeader"
+                  defaultMessage="In the future Cloud deployments will autoscale to have the required number of processors."
+                />
+              }
+              iconType="iInCircle"
+              color={'warning'}
+            >
+              <p>
+                <FormattedMessage
+                  id="xpack.ml.trainedModels.modelsList.startDeployment.cloudWarningText"
+                  defaultMessage="However, in this release you must increase the size of your ML nodes manually in the Cloud console to get more processors."
+                />
+              </p>
+            </EuiCallOut>
+            <EuiSpacer size={'m'} />
+          </>
+        ) : null}
+
+        <EuiCallOut
+          size={'s'}
+          title={
+            <FormattedMessage
+              id="xpack.ml.trainedModels.modelsList.startDeployment.maxNumOfProcessorsWarning"
+              defaultMessage="The product of the number of allocations and threads per allocation should be less than the total number of processors on your ML nodes."
+            />
+          }
+          iconType="iInCircle"
+          color={'primary'}
+        />
+        <EuiSpacer size={'m'} />
+
+        <StartDeploymentSetup config={config} onConfigChange={setConfig} />
+
+        <EuiSpacer size={'m'} />
+      </EuiModalBody>
+
+      <EuiModalFooter>
+        <EuiLink
+          href={startModelDeploymentDocUrl}
+          external
+          target={'_blank'}
+          css={css`
+            align-self: center;
+            margin-right: auto;
+          `}
+        >
+          <FormattedMessage
+            id="xpack.ml.trainedModels.modelsList.startDeployment.docLinkTitle"
+            defaultMessage="Learn more"
+          />
+        </EuiLink>
+
+        <EuiButtonEmpty onClick={onClose}>
+          <FormattedMessage
+            id="xpack.ml.trainedModels.modelsList.startDeployment.cancelButton"
+            defaultMessage="Cancel"
+          />
+        </EuiButtonEmpty>
+
+        <EuiButton
+          type="submit"
+          form={'startDeploymentForm'}
+          onClick={onConfigChange.bind(null, config)}
+          fill
+          disabled={!!errors}
+        >
+          <FormattedMessage
+            id="xpack.ml.trainedModels.modelsList.startDeployment.startButton"
+            defaultMessage="Start"
+          />
+        </EuiButton>
+      </EuiModalFooter>
+    </EuiModal>
+  );
+};
+
+/**
+ * Returns a callback for requesting user's input for threading params
+ * with a form rendered in a modal window.
+ *
+ * @param overlays
+ * @param theme$
+ */
+export const getUserInputThreadingParamsProvider =
+  (overlays: OverlayStart, theme$: Observable<CoreTheme>, startModelDeploymentDocUrl: string) =>
+  (modelId: string): Promise<ThreadingParams | void> => {
+    return new Promise(async (resolve, reject) => {
+      try {
+        const modalSession = overlays.openModal(
+          toMountPoint(
+            wrapWithTheme(
+              <StartDeploymentModal
+                startModelDeploymentDocUrl={startModelDeploymentDocUrl}
+                modelId={modelId}
+                onConfigChange={(config) => {
+                  modalSession.close();
+                  resolve(config);
+                }}
+                onClose={() => {
+                  modalSession.close();
+                  reject();
+                }}
+              />,
+              theme$
+            )
+          )
+        );
+      } catch (e) {
+        reject();
+      }
+    });
+  };
--- 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
@ -5,7 +5,7 @@
 * 2.0.
 */
 import React, { FC } from 'react';
-import { EuiBadge, EuiInMemoryTable, EuiToolTip } from '@elastic/eui';
+import { EuiBadge, EuiIcon, EuiInMemoryTable, EuiToolTip } from '@elastic/eui';
 import { i18n } from '@kbn/i18n';
 import { EuiBasicTableColumn } from '@elastic/eui/src/components/basic_table/basic_table';
 import { FIELD_FORMAT_IDS } from '@kbn/field-formats-plugin/common';
@ -62,6 +62,28 @@ export const AllocatedModels: FC<AllocatedModelsProps> = ({
        return bytesFormatter(v.required_native_memory_bytes);
      },
    },
+    {
+      name: (
+        <EuiToolTip
+          content={i18n.translate('xpack.ml.trainedModels.nodesList.modelsList.allocationTooltip', {
+            defaultMessage: 'number_of_allocations times threads_per_allocation',
+          })}
+        >
+          <span>
+            {i18n.translate('xpack.ml.trainedModels.nodesList.modelsList.allocationHeader', {
+              defaultMessage: 'Allocation',
+            })}
+            <EuiIcon size="s" color="subdued" type="questionInCircle" className="eui-alignTop" />
+          </span>
+        </EuiToolTip>
+      ),
+      width: '100px',
+      truncateText: false,
+      'data-test-subj': 'mlAllocatedModelsTableAllocation',
+      render: (v: AllocatedModel) => {
+        return `${v.node.number_of_allocations} * ${v.node.threads_per_allocation}`;
+      },
+    },
    {
      field: 'node.throughput_last_minute',
      name: i18n.translate(
--- a/x-pack/plugins/ml/server/routes/schemas/inference_schema.ts
+++ b/x-pack/plugins/ml/server/routes/schemas/inference_schema.ts
@ -14,6 +14,13 @@ export const modelIdSchema = schema.object({
  modelId: schema.string(),
 });

+export const threadingParamsSchema = schema.maybe(
+  schema.object({
+    number_of_allocations: schema.number(),
+    threads_per_allocation: schema.number(),
+  })
+);
+
 export const optionalModelIdSchema = schema.object({
  /**
   * Model ID
--- a/x-pack/plugins/ml/server/routes/trained_models.ts
+++ b/x-pack/plugins/ml/server/routes/trained_models.ts
@ -15,6 +15,7 @@ import {
  putTrainedModelQuerySchema,
  inferTrainedModelQuery,
  inferTrainedModelBody,
+  threadingParamsSchema,
 } from './schemas/inference_schema';
 import { modelsProvider } from '../models/data_frame_analytics';
 import { TrainedModelConfigResponse } from '../../common/types/trained_models';
@ -301,6 +302,7 @@ export function trainedModelsRoutes({ router, routeGuard }: RouteInitialization)
      path: '/api/ml/trained_models/{modelId}/deployment/_start',
      validate: {
        params: modelIdSchema,
+        query: threadingParamsSchema,
      },
      options: {
        tags: ['access:ml:canStartStopTrainedModels'],
@ -311,6 +313,7 @@ export function trainedModelsRoutes({ router, routeGuard }: RouteInitialization)
        const { modelId } = request.params;
        const body = await mlClient.startTrainedModelDeployment({
          model_id: modelId,
+          ...(request.query ? request.query : {}),
        });
        return response.ok({
          body,