From 6534fe1a39d5a649a9f51f53e858a6161d5abffa Mon Sep 17 00:00:00 2001
From: Ryan Chard <rchard@anl.gov>
Date: Sat, 12 Mar 2022 11:28:05 +1300
Subject: [PATCH] Adds exception handling on submit and corrects some status
 reporting

---
 aws/action-status.py |  5 ++--
 aws/funcx-run.py     | 56 ++++++++++++++++++++++++++------------------
 2 files changed, 36 insertions(+), 25 deletions(-)

diff --git a/aws/action-status.py b/aws/action-status.py
index 4e20442..be6953c 100644
--- a/aws/action-status.py
+++ b/aws/action-status.py
@@ -84,13 +84,14 @@ def lambda_handler(event, context):
         )
 
         print("updated_response", update_response)
-        print("failure", failure)
+        
         if failure:
             status = "FAILED"
             details = failure
             display_status = failure
         else:
             status = "ACTIVE"
+            display_status = "Task Pending"
             details = None
 
     # Now check again to see if everything is done
@@ -114,7 +115,7 @@ def lambda_handler(event, context):
     result = {
         "action_id": action_id,
         'status': status,
-        'display_status': 'Function Results Received',
+        'display_status': display_status,
         'details': details
     }
 
diff --git a/aws/funcx-run.py b/aws/funcx-run.py
index 7cb64ce..ab41e7d 100644
--- a/aws/funcx-run.py
+++ b/aws/funcx-run.py
@@ -33,36 +33,46 @@ def lambda_handler(event, context):
     monitor_by = body['monitor_by'] if 'monitor_by' in body else None
     manage_by = body['manage_by'] if 'manage_by' in body else None
 
+    action_status = 'ACTIVE'
+    display_status = 'Function Submitted'
+    details = None
+
+    try:
+        batch = fxc.create_batch()
+        for task in body['body']['tasks']:
+            batch.add(endpoint_id=task['endpoint'], function_id=task['function'],
+                    **task['payload'])
+
+        batch_res = fxc.batch_run(batch)
+        print(batch_res)
+
+        # Create a dynamo record where the primary key is this action's ID
+        # Tasks is a dict by task_id and contains the eventual results from their
+        # execution. Where there are no more None results then the action is complete
+        response = table.put_item(
+            Item={
+                'action-id': action_id,
+                'tasks': json.dumps({task_id: {"result": None, "completed": False} for task_id in batch_res})
+            }
+        )
+        print("Dynamo", response)
+    except Exception as eek:
+        print("Exception", str(eek))
+        action_status = 'FAILED'
+        display_status = str(eek)
+        details = str(eek)
+
     result = {
         "action_id": action_id,
-        'status': 'ACTIVE',
-        'display_status': 'Function Submitted',
-        'details': None,
+        'status': action_status,
+        'display_status': display_status,
+        'details': details,
         'monitor_by': monitor_by,
         'manage_by': manage_by,
         'start_time': now_isoformat(),
     }
 
-    batch = fxc.create_batch()
-
-    for task in body['body']['tasks']:
-        print(task)
-        batch.add(endpoint_id=task['endpoint'], function_id=task['function'],
-                  **task['payload'])
-
-    batch_res = fxc.batch_run(batch)
-    print(batch_res)
-
-    # Create a dynamo record where the primary key is this action's ID
-    # Tasks is a dict by task_id and contains the eventual results from their
-    # execution. Where there are no more None results then the action is complete
-    response = table.put_item(
-        Item={
-            'action-id': action_id,
-            'tasks': json.dumps({task_id: {"result": None, "completed": False} for task_id in batch_res})
-        }
-    )
-    print("Dynamo", response)
+    print('Submit result', result)
     return {
         'statusCode': 202,
         'body': json.dumps(result)