[Rest Server] Use list pods with label selector in getting job detail (…

…#4270) * change get each pod to list pod with label selector * fix eslint * update getPods * Add accept header to list pods (#4274) * add accept header * fix mock gpu Co-authored-by: yiyione <yiyi_one@foxmail.com> Co-authored-by: yiyione <yiyi_one@foxmail.com>
microsoft · Mar 11, 2020 · 910fe78 · 910fe78
1 parent ebcf558
commit 910fe78
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 17 deletions.
diff --git a/src/rest-server/src/models/kubernetes/kubernetes.js b/src/rest-server/src/models/kubernetes/kubernetes.js
@@ -94,17 +94,15 @@ const getNodes = async () => {
   return res.data;
 };
 
-const getPods = async (options = {}) => {
-  const {namespace} = options;
+const getPods = async (options = {}, headers={}) => {
+  const {namespace, ...params} = options;
   const client = getClient();
 
-  const requestOptions = Object.assign({}, options);
-  delete requestOptions.namespace;
   let url = '/api/v1/pods';
   if (namespace) {
     url = `/api/v1/namespaces/${namespace}/pods`;
   }
-  const res = await client.get(url, {params: requestOptions});
+  const res = await client.get(url, {params, headers});
   return res.data;
 };
 

diff --git a/src/rest-server/src/models/v2/job/k8s.js b/src/rest-server/src/models/v2/job/k8s.js
@@ -178,7 +178,7 @@ const convertFrameworkSummary = (framework) => {
   };
 };
 
-const convertTaskDetail = async (taskStatus, ports, userName, jobName, taskRoleName) => {
+const convertTaskDetail = async (taskStatus, ports, userName, jobName, taskRoleName, pod, gpuNumber) => {
   // get container ports
   const containerPorts = {};
   if (ports) {
@@ -192,13 +192,6 @@ const convertTaskDetail = async (taskStatus, ports, userName, jobName, taskRoleN
   // get container gpus
   let containerGpus = null;
   try {
-    const response = await k8sModel.getClient().get(
-      launcherConfig.podPath(taskStatus.attemptStatus.podName),
-      {
-        headers: launcherConfig.requestHeaders,
-      }
-    );
-    const pod = response.data;
     if (launcherConfig.enabledHived) {
       const hivedSpec = yaml.load(pod.metadata.annotations['hivedscheduler.microsoft.com/pod-scheduling-spec']);
       if (hivedSpec && hivedSpec.affinityGroup && hivedSpec.affinityGroup.name) {
@@ -209,7 +202,6 @@ const convertTaskDetail = async (taskStatus, ports, userName, jobName, taskRoleN
       const isolation = pod.metadata.annotations['hivedscheduler.microsoft.com/pod-gpu-isolation'];
       containerGpus = isolation.split(',').reduce((attr, id) => attr + Math.pow(2, id), 0);
     } else {
-      const gpuNumber = k8s.atoi(pod.spec.containers[0].resources.limits['nvidia.com/gpu']);
       // mock GPU ids from 0 to (gpuNumber - 1)
       containerGpus = Math.pow(2, gpuNumber) - 1;
     }
@@ -295,21 +287,46 @@ const convertFrameworkDetail = async (framework) => {
     taskRoles: {},
   };
   const ports = {};
+  const gpuNumbers = {};
   for (let taskRoleSpec of framework.spec.taskRoles) {
     ports[taskRoleSpec.name] = taskRoleSpec.task.pod.metadata.annotations['rest-server/port-scheduling-spec'];
+    gpuNumbers[taskRoleSpec.name] = k8s.atoi(
+      taskRoleSpec.task.pod.spec.containers[0].resources.limits['nvidia.com/gpu']
+    );
   }
 
   const userName = framework.metadata.labels ? framework.metadata.labels.userName : 'unknown';
   const jobName = decodeName(framework.metadata.name, framework.metadata.annotations);
+  let pods = {};
+  try {
+    const podList = await k8sModel.getPods({
+      namespace: 'default',
+      labelSelector: `FC_FRAMEWORK_NAME=${framework.metadata.name}`,
+    }, {
+      accept: 'application/json;as=PartialObjectMetadataList;g=meta.k8s.io;v=v1beta1',
+    });
+    pods = _.keyBy(podList.items, (obj) => obj.metadata.name);
+  } catch (err) {
+    pods = {};
+  }
 
   for (let taskRoleStatus of framework.status.attemptStatus.taskRoleStatuses) {
+    const taskStatuses = await Promise.all(taskRoleStatus.taskStatuses.map(
+      async (status) => await convertTaskDetail(
+        status,
+        ports[taskRoleStatus.name],
+        userName,
+        jobName,
+        taskRoleStatus.name,
+        pods[status.attemptStatus.podName],
+        gpuNumbers[taskRoleStatus.name]
+      )
+    ));
     detail.taskRoles[taskRoleStatus.name] = {
       taskRoleStatus: {
         name: taskRoleStatus.name,
       },
-      taskStatuses: await Promise.all(taskRoleStatus.taskStatuses.map(
-        async (status) => await convertTaskDetail(status, ports[taskRoleStatus.name], userName, jobName, taskRoleStatus.name))
-      ),
+      taskStatuses: taskStatuses,
     };
   }