fix: xpk workload list should sum podSet counts (#1106)

Fix wildcard counting logic and priority for super-slicing

Author: jamOne-

src/xpk/core/workload.py

@@ -115,20 +115,23 @@ def _parse_workload_item(item: dict[str, Any]) -> _WorkloadListRow:
         or None
     )
 
-  tpu_vms_needed = _safe_int(pod_sets[0].get('count')) if pod_sets else None
-
-  pod_set_assignments = (
-      item.get('status', {}).get('admission', {}).get('podSetAssignments') or []
+  tpu_vms_needed = (
+      sum(_safe_int(ps.get('count')) for ps in pod_sets) if pod_sets else None
   )
+
+  admission_status = item.get('status', {}).get('admission', {})
+  pod_set_assignments = admission_status.get('podSetAssignments') or []
   tpu_vms_running_ran = (
-      _safe_int(pod_set_assignments[0].get('count'))
+      sum(_safe_int(psa.get('count')) for psa in pod_set_assignments)
       if pod_set_assignments
       else None
   )
 
   reclaimable_pods = item.get('status', {}).get('reclaimablePods') or []
   tpu_vms_done = (
-      _safe_int(reclaimable_pods[0].get('count')) if reclaimable_pods else None
+      sum(_safe_int(rp.get('count')) for rp in reclaimable_pods)
+      if reclaimable_pods
+      else None
   )
 
   conditions = item.get('status', {}).get('conditions') or [{}]

src/xpk/core/workload_test.py

@@ -46,9 +46,9 @@ class _MockWorkloadData:
   jobset_name: str
   created_time: str
   priority: str
-  needed: int | str
-  running: int | str
-  done: int | str
+  needed: list[int]
+  running: list[int]
+  done: list[int]
   status: str
   message: str
   status_time: str
@@ -61,14 +61,19 @@ def _create_mock_workload_json(data: _MockWorkloadData):
           'ownerReferences': [{'name': data.jobset_name}],
       },
       'spec': {
-          'podSets': [{
-              'count': data.needed,
-              'template': {'spec': {'priorityClassName': data.priority}},
-          }]
+          'podSets': [
+              {
+                  'count': v,
+                  'template': {'spec': {'priorityClassName': data.priority}},
+              }
+              for v in data.needed
+          ]
       },
       'status': {
-          'admission': {'podSetAssignments': [{'count': data.running}]},
-          'reclaimablePods': [{'count': data.done}],
+          'admission': {
+              'podSetAssignments': [{'count': v} for v in data.running]
+          },
+          'reclaimablePods': [{'count': v} for v in data.done],
           'conditions': [{
               'type': data.status,
               'message': data.message,
@@ -102,9 +107,9 @@ def test_get_workload_list(commands_tester: CommandsTester):
                   jobset_name='job-test',
                   created_time='2024-01-01T00:00:00Z',
                   priority='high',
-                  needed=32,
-                  running=32,
-                  done=0,
+                  needed=[32],
+                  running=[32],
+                  done=[0],
                   status='Running',
                   message='All good',
                   status_time='2024-01-01T00:01:00Z',
@@ -143,9 +148,9 @@ def test_get_workload_list_filter_by_job(commands_tester: CommandsTester):
                   jobset_name='job-test-1',
                   created_time='2024-01-01T00:00:00Z',
                   priority='high',
-                  needed=32,
-                  running=32,
-                  done=0,
+                  needed=[32],
+                  running=[32],
+                  done=[0],
                   status='Running',
                   message='All good',
                   status_time='2024-01-01T00:01:00Z',
@@ -156,9 +161,9 @@ def test_get_workload_list_filter_by_job(commands_tester: CommandsTester):
                   jobset_name='job-test-2',
                   created_time='2024-01-02T00:00:00Z',
                   priority='low',
-                  needed=4,
-                  running=4,
-                  done=0,
+                  needed=[4],
+                  running=[4],
+                  done=[0],
                   status='Running',
                   message='All good',
                   status_time='2024-01-02T00:01:00Z',
@@ -169,9 +174,9 @@ def test_get_workload_list_filter_by_job(commands_tester: CommandsTester):
                   jobset_name='other-job',
                   created_time='2024-01-03T00:00:00Z',
                   priority='high',
-                  needed=16,
-                  running='',
-                  done=0,
+                  needed=[16],
+                  running=[],
+                  done=[0],
                   status='Admitted',
                   message='Waiting',
                   status_time='2024-01-03T00:01:00Z',
@@ -195,6 +200,37 @@ def test_get_workload_list_filter_by_job(commands_tester: CommandsTester):
   assert parsed_table[1]['Jobset Name'] == 'job-test-2'
 
 
+def test_get_workload_list_multiple_pod_sets(commands_tester: CommandsTester):
+  mock_data = _MockWorkloadData(
+      jobset_name='multi-podset-job',
+      created_time='2024-01-01T00:00:00Z',
+      priority='high',
+      needed=[16, 32],
+      running=[16, 32],
+      done=[16, 32],
+      status='Running',
+      message='All good',
+      status_time='2024-01-01T00:01:00Z',
+  )
+  mock_output = json.dumps({'items': [_create_mock_workload_json(mock_data)]})
+  commands_tester.set_result_for_command(
+      (0, mock_output), 'kubectl', 'get', 'workloads'
+  )
+  args = MagicMock()
+  args.filter_by_status = 'EVERYTHING'
+  args.filter_by_job = None
+
+  return_code, return_value = get_workload_list(args)
+
+  assert return_code == 0
+  parsed_table = _parse_workload_table(return_value)
+  assert len(parsed_table) == 1
+  assert parsed_table[0]['Jobset Name'] == 'multi-podset-job'
+  assert parsed_table[0]['TPU VMs Needed'] == '48'
+  assert parsed_table[0]['TPU VMs Running/Ran'] == '48'
+  assert parsed_table[0]['TPU VMs Done'] == '48'
+
+
 @pytest.mark.parametrize(
     'filter_by_status, expected_job_names',
     [
@@ -226,9 +262,9 @@ def test_get_workload_list_filters(
                   jobset_name='queued-job',
                   created_time='2024-01-01T00:00:00Z',
                   priority='high',
-                  needed=4,
-                  running='',
-                  done=0,
+                  needed=[4],
+                  running=[],
+                  done=[0],
                   status='Admitted',
                   message='Waiting',
                   status_time='2024-01-01T00:01:00Z',
@@ -239,9 +275,9 @@ def test_get_workload_list_filters(
                   jobset_name='running-job',
                   created_time='2024-01-01T00:00:00Z',
                   priority='high',
-                  needed=4,
-                  running=4,
-                  done=0,
+                  needed=[4],
+                  running=[4],
+                  done=[0],
                   status='Admitted',
                   message='Running',
                   status_time='2024-01-01T00:01:00Z',
@@ -252,9 +288,9 @@ def test_get_workload_list_filters(
                   jobset_name='success-job',
                   created_time='2024-01-01T00:00:00Z',
                   priority='high',
-                  needed=4,
-                  running=4,
-                  done=4,
+                  needed=[4],
+                  running=[4],
+                  done=[4],
                   status='Finished',
                   message='Job finishedsuccessfully',
                   status_time='2024-01-01T00:01:00Z',
@@ -265,9 +301,9 @@ def test_get_workload_list_filters(
                   jobset_name='failed-job',
                   created_time='2024-01-01T00:00:00Z',
                   priority='high',
-                  needed=4,
-                  running=4,
-                  done=0,
+                  needed=[4],
+                  running=[4],
+                  done=[0],
                   status='Finished',
                   message='Job failed witherror',
                   status_time='2024-01-01T00:01:00Z',