From 717950c1c6d37c1b3ea19036c383b369ad3b9284 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 +++++++
 batch/pinned-requirements.txt              | 99 +++++++++++++++++++---
 web_common/pinned-requirements.txt         | 55 ++++++++++--
 4 files changed, 186 insertions(+), 24 deletions(-)

diff --git a/batch/batch/front_end/front_end.py b/batch/batch/front_end/front_end.py
index d5348e5364ab..bb2cfed5eaf1 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 caf911f368cb..6d0c0218be3a 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'>
diff --git a/batch/pinned-requirements.txt b/batch/pinned-requirements.txt
index ea4f6eb8b66d..1345bd1fd017 100644
--- a/batch/pinned-requirements.txt
+++ b/batch/pinned-requirements.txt
@@ -1,43 +1,114 @@
 # This file was autogenerated by uv via the following command:
 #    uv pip compile --python-version 3.9 --python-platform linux batch/requirements.txt --output-file=batch/pinned-requirements.txt
 aiodocker==0.21.0
+    # via -r batch/requirements.txt
 aiohttp==3.9.5
-    # via aiodocker
+    # via
+    #   -c batch/../gear/pinned-requirements.txt
+    #   -c batch/../hail/python/dev/pinned-requirements.txt
+    #   -c batch/../hail/python/pinned-requirements.txt
+    #   -c batch/../web_common/pinned-requirements.txt
+    #   aiodocker
 aiorwlock==1.4.0
+    # via -r batch/requirements.txt
 aiosignal==1.3.1
-    # via aiohttp
+    # via
+    #   -c batch/../gear/pinned-requirements.txt
+    #   -c batch/../hail/python/dev/pinned-requirements.txt
+    #   -c batch/../hail/python/pinned-requirements.txt
+    #   -c batch/../web_common/pinned-requirements.txt
+    #   aiohttp
 async-timeout==4.0.3
-    # via aiohttp
+    # via
+    #   -c batch/../gear/pinned-requirements.txt
+    #   -c batch/../hail/python/dev/pinned-requirements.txt
+    #   -c batch/../hail/python/pinned-requirements.txt
+    #   -c batch/../web_common/pinned-requirements.txt
+    #   -r batch/requirements.txt
+    #   aiohttp
 attrs==23.2.0
-    # via aiohttp
+    # via
+    #   -c batch/../gear/pinned-requirements.txt
+    #   -c batch/../hail/python/dev/pinned-requirements.txt
+    #   -c batch/../hail/python/pinned-requirements.txt
+    #   -c batch/../web_common/pinned-requirements.txt
+    #   aiohttp
 dictdiffer==0.9.0
+    # via -r batch/requirements.txt
 frozenlist==1.4.1
     # via
+    #   -c batch/../gear/pinned-requirements.txt
+    #   -c batch/../hail/python/dev/pinned-requirements.txt
+    #   -c batch/../hail/python/pinned-requirements.txt
+    #   -c batch/../web_common/pinned-requirements.txt
     #   aiohttp
     #   aiosignal
 idna==3.7
-    # via yarl
+    # via
+    #   -c batch/../gear/pinned-requirements.txt
+    #   -c batch/../hail/python/dev/pinned-requirements.txt
+    #   -c batch/../hail/python/pinned-requirements.txt
+    #   -c batch/../web_common/pinned-requirements.txt
+    #   yarl
 multidict==6.0.5
     # via
+    #   -c batch/../gear/pinned-requirements.txt
+    #   -c batch/../hail/python/dev/pinned-requirements.txt
+    #   -c batch/../hail/python/pinned-requirements.txt
+    #   -c batch/../web_common/pinned-requirements.txt
     #   aiohttp
     #   yarl
 numpy==1.26.4
-    # via pandas
+    # via
+    #   -c batch/../hail/python/dev/pinned-requirements.txt
+    #   -c batch/../hail/python/pinned-requirements.txt
+    #   pandas
 packaging==24.0
-    # via plotly
+    # via
+    #   -c batch/../hail/python/dev/pinned-requirements.txt
+    #   -c batch/../hail/python/pinned-requirements.txt
+    #   plotly
 pandas==2.2.2
+    # via
+    #   -c batch/../hail/python/pinned-requirements.txt
+    #   -r batch/requirements.txt
 plotly==5.22.0
+    # via
+    #   -c batch/../hail/python/pinned-requirements.txt
+    #   -r batch/requirements.txt
 python-dateutil==2.9.0.post0
-    # via pandas
+    # via
+    #   -c batch/../gear/pinned-requirements.txt
+    #   -c batch/../hail/python/dev/pinned-requirements.txt
+    #   -c batch/../hail/python/pinned-requirements.txt
+    #   pandas
 pytz==2024.1
-    # via pandas
+    # via
+    #   -c batch/../hail/python/pinned-requirements.txt
+    #   pandas
 six==1.16.0
-    # via python-dateutil
+    # via
+    #   -c batch/../gear/pinned-requirements.txt
+    #   -c batch/../hail/python/dev/pinned-requirements.txt
+    #   -c batch/../hail/python/pinned-requirements.txt
+    #   python-dateutil
 tenacity==8.3.0
-    # via plotly
+    # via
+    #   -c batch/../hail/python/pinned-requirements.txt
+    #   plotly
 typing-extensions==4.11.0
-    # via aiodocker
+    # via
+    #   -c batch/../hail/python/dev/pinned-requirements.txt
+    #   -c batch/../hail/python/pinned-requirements.txt
+    #   aiodocker
 tzdata==2024.1
-    # via pandas
+    # via
+    #   -c batch/../hail/python/pinned-requirements.txt
+    #   pandas
 yarl==1.9.4
-    # via aiohttp
+    # via
+    #   -c batch/../gear/pinned-requirements.txt
+    #   -c batch/../hail/python/dev/pinned-requirements.txt
+    #   -c batch/../hail/python/pinned-requirements.txt
+    #   -c batch/../web_common/pinned-requirements.txt
+    #   aiohttp
diff --git a/web_common/pinned-requirements.txt b/web_common/pinned-requirements.txt
index 5d6c3d12f36a..66eb2640d5bc 100644
--- a/web_common/pinned-requirements.txt
+++ b/web_common/pinned-requirements.txt
@@ -1,28 +1,67 @@
 # This file was autogenerated by uv via the following command:
 #    uv pip compile --python-version 3.9 --python-platform linux web_common/requirements.txt --output-file=web_common/pinned-requirements.txt
 aiohttp==3.9.5
-    # via aiohttp-jinja2
+    # via
+    #   -c web_common/../gear/pinned-requirements.txt
+    #   -c web_common/../hail/python/dev/pinned-requirements.txt
+    #   -c web_common/../hail/python/pinned-requirements.txt
+    #   aiohttp-jinja2
 aiohttp-jinja2==1.6
+    # via -r web_common/requirements.txt
 aiosignal==1.3.1
-    # via aiohttp
+    # via
+    #   -c web_common/../gear/pinned-requirements.txt
+    #   -c web_common/../hail/python/dev/pinned-requirements.txt
+    #   -c web_common/../hail/python/pinned-requirements.txt
+    #   aiohttp
 async-timeout==4.0.3
-    # via aiohttp
+    # via
+    #   -c web_common/../gear/pinned-requirements.txt
+    #   -c web_common/../hail/python/dev/pinned-requirements.txt
+    #   -c web_common/../hail/python/pinned-requirements.txt
+    #   aiohttp
 attrs==23.2.0
-    # via aiohttp
+    # via
+    #   -c web_common/../gear/pinned-requirements.txt
+    #   -c web_common/../hail/python/dev/pinned-requirements.txt
+    #   -c web_common/../hail/python/pinned-requirements.txt
+    #   aiohttp
 frozenlist==1.4.1
     # via
+    #   -c web_common/../gear/pinned-requirements.txt
+    #   -c web_common/../hail/python/dev/pinned-requirements.txt
+    #   -c web_common/../hail/python/pinned-requirements.txt
     #   aiohttp
     #   aiosignal
 idna==3.7
-    # via yarl
+    # via
+    #   -c web_common/../gear/pinned-requirements.txt
+    #   -c web_common/../hail/python/dev/pinned-requirements.txt
+    #   -c web_common/../hail/python/pinned-requirements.txt
+    #   yarl
 jinja2==3.1.4
-    # via aiohttp-jinja2
+    # via
+    #   -c web_common/../hail/python/dev/pinned-requirements.txt
+    #   -c web_common/../hail/python/pinned-requirements.txt
+    #   -r web_common/requirements.txt
+    #   aiohttp-jinja2
 libsass==0.23.0
+    # via -r web_common/requirements.txt
 markupsafe==2.1.5
-    # via jinja2
+    # via
+    #   -c web_common/../hail/python/dev/pinned-requirements.txt
+    #   -c web_common/../hail/python/pinned-requirements.txt
+    #   jinja2
 multidict==6.0.5
     # via
+    #   -c web_common/../gear/pinned-requirements.txt
+    #   -c web_common/../hail/python/dev/pinned-requirements.txt
+    #   -c web_common/../hail/python/pinned-requirements.txt
     #   aiohttp
     #   yarl
 yarl==1.9.4
-    # via aiohttp
+    # via
+    #   -c web_common/../gear/pinned-requirements.txt
+    #   -c web_common/../hail/python/dev/pinned-requirements.txt
+    #   -c web_common/../hail/python/pinned-requirements.txt
+    #   aiohttp