swf-testbed/episodes/prompt_processing.py

0001 """Prompt-processing workflow episodes.
0002
0003 Agent messages are recorded live through the shared testbed mapping;
0004 the completion pass joins the PanDA side — the tasks named by the run
0005 number and the jobs born within them, which are the workflow's worker
0006 lanes (docs/agentic-workflow-view.md).
0007 """
0008
0009 import logging
0010
0011 from .common import TestbedEpisodeDefinition
0012
0013 logger = logging.getLogger(__name__)
0014
0015 TERMINAL_TASK_STATUSES = {'done', 'finished', 'failed', 'broken',
0016                           'aborted', 'exhausted'}
0017
0018
0019 def _aware(value):
0020     """PanDA REST timestamps are naive UTC; stamp the offset."""
0021     if not value:
0022         return None
0023     value = str(value)
0024     if value.endswith('Z') or '+' in value[10:]:
0025         return value
0026     return value + '+00:00'
0027
0028
0029 class PromptProcessingEpisodes(TestbedEpisodeDefinition):
0030     workflow_name = 'prompt_processing'
0031     completion_deadline_seconds = 1800
0032
0033     def completion_poll(self, context, ingest):
0034         run_id = ((context.last_message or {}).get('run_id')
0035                   or (context.first_message or {}).get('run_id'))
0036         if not run_id:
0037             logger.warning('episode %s has no run id; closing without '
0038                            'a PanDA join', context.episode_id)
0039             return True
0040
0041         session = ingest.session
0042         base = ingest.base_url
0043         response = session.get(
0044             f'{base}/api/panda/tasks/',
0045             params={'taskname': f'swf.{run_id}.processed', 'days': 2},
0046             timeout=30)
0047         response.raise_for_status()
0048         tasks = response.json().get('items') or []
0049         if not tasks:
0050             # Tasks appear seconds after submission; none yet means the
0051             # join is early, not empty. The deadline bounds the wait.
0052             return False
0053         if any((t.get('status') or '') not in TERMINAL_TASK_STATUSES
0054                for t in tasks):
0055             return False
0056
0057         events, participants = [], []
0058         for task in tasks:
0059             taskid = task['jeditaskid']
0060             task_pid = f'task-{taskid}'
0061             participants.append({
0062                 'id': task_pid,
0063                 'label': task.get('taskname') or task_pid,
0064                 'kind': 'panda_task',
0065                 'born_at': _aware(task.get('creationdate')),
0066                 'died_at': _aware(task.get('endtime')),
0067             })
0068             events.append({
0069                 'time': _aware(task.get('creationdate')),
0070                 'kind': 'task_created',
0071                 'participant': task_pid,
0072                 'payload': {'jeditaskid': taskid,
0073                             'site': task.get('site'),
0074                             'status': task.get('status')},
0075             })
0076             jobs_response = session.get(
0077                 f'{base}/api/panda/jobs/',
0078                 params={'taskid': taskid, 'days': 2}, timeout=60)
0079             jobs_response.raise_for_status()
0080             for job in jobs_response.json().get('items') or []:
0081                 job_pid = f"job-{job['pandaid']}"
0082                 participants.append({
0083                     'id': job_pid,
0084                     'label': f"job {job['pandaid']}",
0085                     'kind': 'panda_job',
0086                     'born_at': _aware(job.get('creationtime')),
0087                     'died_at': _aware(job.get('endtime')),
0088                 })
0089                 for field, kind in (('creationtime', 'job_created'),
0090                                     ('starttime', 'job_started'),
0091                                     ('endtime', 'job_ended')):
0092                     if job.get(field):
0093                         events.append({
0094                             'time': _aware(job[field]),
0095                             'kind': kind,
0096                             'participant': job_pid,
0097                             'payload': {'site': job.get('computingsite'),
0098                                         'status': job.get('jobstatus'),
0099                                         'jeditaskid': taskid},
0100                         })
0101
0102         ingest.append(scope=self.scope, episode_id=context.episode_id,
0103                       events=events, participants=participants)
0104         context.notes['panda_tasks'] = [t['jeditaskid'] for t in tasks]
0105         return True
0106
0107     def summary(self, context):
0108         return {'run_id': ((context.last_message or {}).get('run_id')),
0109                 'panda_tasks': context.notes.get('panda_tasks', [])}