From a1af62be290d390ee59e3326cb3fe2e25abd5044 Mon Sep 17 00:00:00 2001
From: Daniel Goldstein <danielgold95@gmail.com>
Date: Fri, 14 Jun 2024 16:50:49 -0400
Subject: [PATCH] [batch] Expose job group structure in the UI

---
 batch/batch/front_end/front_end.py         | 24 ++++++++++++++--
 batch/batch/front_end/templates/batch.html | 32 ++++++++++++++++++++++
 2 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/batch/batch/front_end/front_end.py b/batch/batch/front_end/front_end.py
index d5348e5364a..bb2cfed5eaf 100644
--- a/batch/batch/front_end/front_end.py
+++ b/batch/batch/front_end/front_end.py
@@ -2171,20 +2171,35 @@ async def delete_batch(request: web.Request, _, batch_id: int) -> web.Response:
 @billing_project_users_only()
 @catch_ui_error_in_dev
 async def ui_batch(request, userdata, batch_id):
+    return await _ui_job_group(request, userdata, batch_id, ROOT_JOB_GROUP_ID)
+
+
+@routes.get('/batches/{batch_id}/job-groups/{job_group_id}')
+@billing_project_users_only()
+@catch_ui_error_in_dev
+async def ui_job_group(request, userdata, batch_id):
+    job_group_id = int(request.match_info['job_group_id'])
+    return await _ui_job_group(request, userdata, batch_id, job_group_id)
+
+
+async def _ui_job_group(request, userdata, batch_id, job_group_id):
     app = request.app
     batch = await _get_batch(app, batch_id)
 
     q = request.query.get('q', '')
     last_job_id = cast_query_param_to_int(request.query.get('last_job_id'))
+    last_job_group_id = cast_query_param_to_int(request.query.get('last_job_group_id'))
 
     try:
         jobs, last_job_id = await _query_job_group_jobs(
-            request, batch_id, ROOT_JOB_GROUP_ID, CURRENT_QUERY_VERSION, q, last_job_id, recursive=True
+            request, batch_id, job_group_id, CURRENT_QUERY_VERSION, q, last_job_id, recursive=False
         )
+        job_groups, last_job_group_id = await _query_job_groups(request, batch_id, job_group_id, last_job_group_id)
     except QueryError as e:
         session = await aiohttp_session.get_session(request)
         set_message(session, e.message, 'error')
         jobs = []
+        job_groups = []
         last_job_id = None
 
     for j in jobs:
@@ -2195,8 +2210,12 @@ async def ui_batch(request, userdata, batch_id):
             if j['always_run'] and j['state'] not in {'Success', 'Failed', 'Error'}
             else j['state']
         )
-    batch['jobs'] = jobs
+    for jg in job_groups:
+        jg['duration'] = humanize_timedelta_msecs(jg['duration'])
+        jg['cost'] = cost_str(jg['cost'])
 
+    batch['jobs'] = jobs
+    batch['job_groups'] = job_groups
     batch['cost'] = cost_str(batch['cost'])
 
     if batch['cost_breakdown'] is not None:
@@ -2208,6 +2227,7 @@ async def ui_batch(request, userdata, batch_id):
         'batch': batch,
         'q': q,
         'last_job_id': last_job_id,
+        'last_job_group_id': last_job_group_id,
     }
     return await render_template('batch', request, userdata, 'batch.html', page_context)
 
diff --git a/batch/batch/front_end/templates/batch.html b/batch/batch/front_end/templates/batch.html
index caf911f368c..6d0c0218be3 100644
--- a/batch/batch/front_end/templates/batch.html
+++ b/batch/batch/front_end/templates/batch.html
@@ -84,6 +84,38 @@
           <th class='h-12 bg-slate-200 font-light text-md text-left hidden md:table-cell rounded-tr'>Cost</th>
         </thead>
         <tbody class='border border-collapse border-slate-50'>
+          {% for jg in batch['job_groups'] %}
+          <tr class='border border-collapse hover:bg-slate-100'>
+            <td class='font-light pl-4 w-20'>
+              {{ link(base_path ~ '/batches/' ~ jg['batch_id'] ~ '/job-groups/' ~ jg['job_group_id'],
+              jg['job_group_id']) }}
+            </td>
+            <td class='py-1 block overflow-x-auto'>
+              <div class='flex flex-col space-x-0 md:space-x-2 md:flex-row md:flex-wrap'>
+                {% if 'attributes' in jg and 'name' in jg['attributes'] %}
+                <div class='text-wrap'>
+                  {{ link(base_path ~ '/batches/' ~ jg['batch_id'] ~ '/job-groups/' ~ jg['job_group_id'],
+                  jg['attributes']['name'])
+                  }}
+                </div>
+                {% else %}
+                <div class='text-wrap text-zinc-400 italic'>
+                  {{ link(base_path ~ '/batches/' ~ jg['batch_id'] ~ '/job-groups/' ~ jg['job_group_id'], 'no name') }}
+                </div>
+                {% endif %}
+                <div class='flex items-center'>
+                  {{ batch_state_indicator(jg) }}
+                </div>
+              </div>
+            </td>
+            <td class='hidden lg:table-cell font-light py-2 px-4'>
+              {{ jg.get('duration') or '' }}
+            </td>
+            <td class='hidden md:table-cell font-light py-2 px-4'>
+              {{ jg.get('cost') or '' }}
+            </td>
+          </tr>
+          {% endfor %}
           {% for job in batch['jobs'] %}
           <tr class='border border-collapse hover:bg-slate-100'>
             <td class='font-light pl-4 w-20'>