Fetch numNodes from ClusterTrainingRuntime (opendatahub-io#5508)

* Fetch numNodes from ClusterTrainingRuntime

* resolving type check

Author: dpanshug

frontend/src/api/models/kubeflow.ts

@@ -13,3 +13,10 @@ export const TrainJobModel: K8sModelCommon = {
   kind: 'TrainJob',
   plural: 'trainjobs',
 };
+
+export const ClusterTrainingRuntimeModel: K8sModelCommon = {
+  apiVersion: 'v1alpha1',
+  apiGroup: 'trainer.kubeflow.org',
+  kind: 'ClusterTrainingRuntime',
+  plural: 'clustertrainingruntimes',
+};

packages/model-training/src/__mocks__/mockClusterTrainingRuntimeK8sResource.ts

@@ -0,0 +1,69 @@
+import { ClusterTrainingRuntimeKind } from '@odh-dashboard/model-training/k8sTypes';
+
+type MockClusterTrainingRuntimeConfigType = {
+  name?: string;
+  numNodes?: number;
+  numProcPerNode?: string | number;
+  framework?: string;
+};
+
+export const mockClusterTrainingRuntimeK8sResource = ({
+  name = 'training-cuda128-torch28-py312',
+  numNodes = 1,
+  numProcPerNode = 'auto',
+  framework = 'torch',
+}: MockClusterTrainingRuntimeConfigType = {}): ClusterTrainingRuntimeKind => ({
+  apiVersion: 'trainer.kubeflow.org/v1alpha1',
+  kind: 'ClusterTrainingRuntime',
+  metadata: {
+    name,
+    creationTimestamp: '2025-11-05T14:36:39Z',
+    labels: {
+      'app.kubernetes.io/component': 'controller',
+      'app.kubernetes.io/name': 'trainer',
+      [`trainer.kubeflow.org/framework`]: framework,
+    },
+    resourceVersion: '277649620',
+    generation: 2,
+  },
+  spec: {
+    mlPolicy: {
+      numNodes,
+      torch: {
+        numProcPerNode,
+      },
+    },
+    template: {
+      spec: {
+        replicatedJobs: [
+          {
+            groupName: 'default',
+            name: 'node',
+            replicas: 1,
+            template: {
+              metadata: {
+                labels: {
+                  'trainer.kubeflow.org/trainjob-ancestor-step': 'trainer',
+                },
+              },
+              spec: {
+                template: {
+                  metadata: {},
+                  spec: {
+                    containers: [
+                      {
+                        image: 'quay.io/rhoai/odh-training-cuda128-torch28-py312-rhel9:rhoai-3.0',
+                        name: 'node',
+                        resources: {},
+                      },
+                    ],
+                  },
+                },
+              },
+            },
+          },
+        ],
+      },
+    },
+  },
+});

packages/model-training/src/api.ts

@@ -2,6 +2,7 @@ import {
   k8sDeleteResource,
   K8sStatus,
   k8sPatchResource,
+  k8sGetResource,
 } from '@openshift/dynamic-plugin-sdk-utils';
 import { applyK8sAPIOptions } from '@odh-dashboard/internal/api/apiMergeUtils';
 import { K8sAPIOptions, WorkloadKind } from '@odh-dashboard/internal/k8sTypes';
@@ -10,8 +11,11 @@ import { WorkloadModel } from '@odh-dashboard/internal/api/models/kueue';
 import { groupVersionKind } from '@odh-dashboard/internal/api/k8sUtils';
 import { CustomWatchK8sResult } from '@odh-dashboard/internal/types';
 import useK8sWatchResourceList from '@odh-dashboard/internal/utilities/useK8sWatchResourceList';
-import { TrainJobModel } from '@odh-dashboard/internal/api/models/kubeflow';
-import { TrainJobKind } from './k8sTypes';
+import {
+  TrainJobModel,
+  ClusterTrainingRuntimeModel,
+} from '@odh-dashboard/internal/api/models/kubeflow';
+import { TrainJobKind, ClusterTrainingRuntimeKind } from './k8sTypes';
 
 export const useTrainJobs = (namespace: string): CustomWatchK8sResult<TrainJobKind[]> =>
   useK8sWatchResourceList(
@@ -166,3 +170,17 @@ export const toggleTrainJobHibernation = async (
     };
   }
 };
+
+export const getClusterTrainingRuntime = (
+  name: string,
+  opts?: K8sAPIOptions,
+): Promise<ClusterTrainingRuntimeKind> =>
+  k8sGetResource<ClusterTrainingRuntimeKind>(
+    applyK8sAPIOptions(
+      {
+        model: ClusterTrainingRuntimeModel,
+        queryOptions: { name },
+      },
+      opts,
+    ),
+  );

packages/model-training/src/global/trainingJobDetailsDrawer/TrainingJobResourcesTab.tsx

@@ -52,7 +52,7 @@ const TrainingJobResourcesTab: React.FC<TrainingJobResourcesTabProps> = ({ job }
                 isDisabled // TODO: RHOAIENG-37576 Uncomment this when scaling is implemented
                 data-testid="nodes-edit-button"
               >
-                {job.spec.trainer.numNodes}
+                {job.spec.trainer?.numNodes}
               </Button>
             </DescriptionListDescription>
           </DescriptionListGroup>
@@ -61,7 +61,7 @@ const TrainingJobResourcesTab: React.FC<TrainingJobResourcesTabProps> = ({ job }
               Processes per node:
             </DescriptionListTerm>
             <DescriptionListDescription data-testid="processes-per-node-value">
-              {job.spec.trainer.numProcPerNode || '-'}
+              {job.spec.trainer?.numProcPerNode || '-'}
             </DescriptionListDescription>
           </DescriptionListGroup>
         </DescriptionList>
@@ -76,29 +76,29 @@ const TrainingJobResourcesTab: React.FC<TrainingJobResourcesTabProps> = ({ job }
               CPU requests:
             </DescriptionListTerm>
             <DescriptionListDescription data-testid="cpu-requests-value">
-              {job.spec.trainer.resourcesPerNode?.requests?.cpu || '-'}
+              {job.spec.trainer?.resourcesPerNode?.requests?.cpu || '-'}
             </DescriptionListDescription>
           </DescriptionListGroup>
           <DescriptionListGroup>
             <DescriptionListTerm style={{ fontWeight: 'normal' }}>CPU limits:</DescriptionListTerm>
             <DescriptionListDescription data-testid="cpu-limits-value">
-              {job.spec.trainer.resourcesPerNode?.limits?.cpu || '-'}
+              {job.spec.trainer?.resourcesPerNode?.limits?.cpu || '-'}
             </DescriptionListDescription>
           </DescriptionListGroup>
           <DescriptionListGroup>
             <DescriptionListTerm style={{ fontWeight: 'normal' }}>
               Memory requests:
             </DescriptionListTerm>
             <DescriptionListDescription data-testid="memory-requests-value">
-              {job.spec.trainer.resourcesPerNode?.requests?.memory || '-'}
+              {job.spec.trainer?.resourcesPerNode?.requests?.memory || '-'}
             </DescriptionListDescription>
           </DescriptionListGroup>
           <DescriptionListGroup>
             <DescriptionListTerm style={{ fontWeight: 'normal' }}>
               Memory limits:
             </DescriptionListTerm>
             <DescriptionListDescription data-testid="memory-limits-value">
-              {job.spec.trainer.resourcesPerNode?.limits?.memory || '-'}
+              {job.spec.trainer?.resourcesPerNode?.limits?.memory || '-'}
             </DescriptionListDescription>
           </DescriptionListGroup>
         </DescriptionList>

packages/model-training/src/global/trainingJobList/TrainingJobClusterQueue.tsx

@@ -16,6 +16,10 @@ const TrainingJobClusterQueue: React.FC<TrainingJobClusterQueueProps> = ({
     namespace,
   );
 
+  if (!localQueueName) {
+    return '-';
+  }
+
   if (!clusterQueueLoaded) {
     return <Skeleton width="100px" />;
   }

packages/model-training/src/global/trainingJobList/TrainingJobTableRow.tsx

@@ -14,6 +14,7 @@ import StateActionToggle from './StateActionToggle';
 import { TrainJobKind } from '../../k8sTypes';
 import { TrainingJobState } from '../../types';
 import { toggleTrainJobHibernation } from '../../api';
+import useClusterTrainingRuntime from '../../hooks/useClusterTrainingRuntime';
 
 type TrainingJobTableRowProps = {
   job: TrainJobKind;
@@ -35,7 +36,25 @@ const TrainingJobTableRow: React.FC<TrainingJobTableRowProps> = ({
   const [isToggling, setIsToggling] = React.useState(false);
 
   const displayName = getDisplayNameFromK8sResource(job);
-  const nodesCount = job.spec.trainer.numNodes || 0;
+
+  // Fetch ClusterTrainingRuntime if trainer spec is not available
+  const runtimeName =
+    job.spec.runtimeRef.kind === 'ClusterTrainingRuntime' ? job.spec.runtimeRef.name : null;
+  const { clusterTrainingRuntime, loaded: runtimeLoaded } = useClusterTrainingRuntime(
+    !job.spec.trainer ? runtimeName : null,
+  );
+
+  // Get numNodes from trainer spec or ClusterTrainingRuntime
+  const nodesCount = React.useMemo(() => {
+    if (job.spec.trainer?.numNodes) {
+      return job.spec.trainer.numNodes;
+    }
+    if (runtimeLoaded && clusterTrainingRuntime?.spec.mlPolicy?.numNodes) {
+      return clusterTrainingRuntime.spec.mlPolicy.numNodes;
+    }
+    return 0;
+  }, [job.spec.trainer?.numNodes, runtimeLoaded, clusterTrainingRuntime]);
+
   const localQueueName = job.metadata.labels?.['kueue.x-k8s.io/queue-name'];
 
   const status = jobStatus || getTrainingJobStatusSync(job);

packages/model-training/src/global/trainingJobList/const.ts

@@ -24,8 +24,8 @@ export const columns: SortableData<TrainJobKind>[] = [
     label: 'Nodes',
     width: 15,
     sortable: (a: TrainJobKind, b: TrainJobKind): number => {
-      const aNodes = a.spec.trainer.numNodes || 0;
-      const bNodes = b.spec.trainer.numNodes || 0;
+      const aNodes = a.spec.trainer?.numNodes || 0;
+      const bNodes = b.spec.trainer?.numNodes || 0;
 
       return aNodes - bNodes;
     },