refactor origin_job_exists check

elastic · Nov 15, 2023 · 50ffcfd · 50ffcfd
1 parent 12f1a01
commit 50ffcfd
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 21 deletions.
diff --git a/x-pack/plugins/ml/common/types/trained_models.ts b/x-pack/plugins/ml/common/types/trained_models.ts
@@ -12,6 +12,7 @@ import type {
   TotalFeatureImportance,
 } from '@kbn/ml-data-frame-analytics-utils';
 import { IndexName, IndicesIndexState } from '@elastic/elasticsearch/lib/api/types';
+import { isPopulatedObject } from '@kbn/ml-is-populated-object';
 import type { XOR } from './common';
 import type { MlSavedObjectType } from './saved_objects';
 
@@ -291,3 +292,10 @@ export interface TrainedModelStatsResponse extends estypes.MlTrainedModelStats {
   deployment_stats?: Omit<TrainedModelDeploymentStatsResponse, 'model_id'>;
   model_size_stats?: TrainedModelModelSizeStats;
 }
+
+export function isTrainedModelConfigResponse(arg: unknown): arg is TrainedModelConfigResponse {
+  return (
+    isPopulatedObject(arg, ['metadata']) &&
+    isPopulatedObject(arg.metadata ?? {}, ['analytics_config'])
+  );
+}
diff --git a/x-pack/plugins/ml/server/routes/trained_models.ts b/x-pack/plugins/ml/server/routes/trained_models.ts
@@ -29,9 +29,10 @@ import {
   createIngestPipelineSchema,
   modelDownloadsQuery,
 } from './schemas/inference_schema';
-import type {
+import {
+  isTrainedModelConfigResponse,
   PipelineDefinition,
-  TrainedModelConfigResponse,
+  type TrainedModelConfigResponse,
 } from '../../common/types/trained_models';
 import { mlLog } from '../lib/log';
 import { forceQuerySchema } from './schemas/anomaly_detectors_schema';
@@ -190,39 +191,49 @@ export function trainedModelsRoutes(
             // we don't need to fill kibana's log with these messages.
             mlLog.debug(e);
           }
-          const enabledFeatures = getEnabledFeatures();
+
+          const filteredModels = filterForEnabledFeatureModels(result, getEnabledFeatures());
+
           try {
-            if (enabledFeatures.dfa) {
-              const jobIds = result.map((model) => {
-                let id = model.metadata?.analytics_config?.id;
-                if (id) {
-                  id = `${id}*`;
+            // @ts-ignore
+            const jobIdsString = filteredModels.reduce(
+              (
+                jobIdsStr: string,
+                currentModel: TrainedModelConfigResponse | estypes.MlTrainedModelConfig,
+                idx: number
+              ) => {
+                if (isTrainedModelConfigResponse(currentModel)) {
+                  let id = currentModel.metadata?.analytics_config?.id ?? '';
+                  if (id !== '') {
+                    id = `${idx > 0 ? ',' : ''}${id}*`;
+                  }
+                  return `${jobIdsStr}${id}`;
                 }
-                return id;
-              });
-              const filteredJobIds = jobIds.filter((id) => id !== undefined);
+                return jobIdsStr;
+              },
+              ''
+            );
 
+            if (jobIdsString !== '') {
               const { data_frame_analytics: jobs } = await mlClient.getDataFrameAnalytics({
-                id: filteredJobIds.join(','),
+                id: jobIdsString,
                 allow_no_match: true,
               });
 
-              jobs.forEach(({ id }) => {
-                const model = result.find(
-                  (modelWithJob) => id === modelWithJob.metadata?.analytics_config?.id
-                );
-
-                if (model) {
-                  model.origin_job_exists = true;
+              filteredModels.forEach((model) => {
+                if (isTrainedModelConfigResponse(model)) {
+                  const dfaId = model?.metadata?.analytics_config?.id;
+                  if (dfaId !== undefined) {
+                    // if this is a dfa model, set origin_job_exists
+                    model.origin_job_exists = jobs.find((job) => job.id === dfaId) !== undefined;
+                  }
                 }
               });
             }
           } catch (e) {
             // Swallow error to prevent blocking trained models result
           }
 
-          const filteredModels = filterForEnabledFeatureModels(result, enabledFeatures);
-
           return response.ok({
             body: filteredModels,
           });