add workflow cap for all users (#140)

add workflow cap for all users

Author: tianj7

src/argowrapper/constants.py

@@ -42,6 +42,7 @@
 GEN3_WORKFLOW_PHASE_LABEL: Final = "phase"
 GEN3_SUBMIT_TIMESTAMP_LABEL: Final = "submittedAt"
 GEN3_NON_VA_WORKFLOW_MONTHLY_CAP: Final = 20
+GEN3_DEFAULT_WORKFLOW_MONTHLY_CAP: Final = 50
 
 
 class POD_COMPLETION_STRATEGY(Enum):

src/argowrapper/engine/argo_engine.py

@@ -405,7 +405,11 @@ def get_user_workflows_for_current_month(self, auth_header: str) -> List[Dict]:
         )
         user_monthly_workflows = []
         for workflow in all_user_workflows:
-            if workflow[GEN3_WORKFLOW_PHASE_LABEL] in {"Running", "Succeeded"}:
+            if workflow[GEN3_WORKFLOW_PHASE_LABEL] in {
+                "Running",
+                "Succeeded",
+                "Failed",
+            }:
                 submitted_time_str = workflow[GEN3_SUBMIT_TIMESTAMP_LABEL]
                 submitted_time = datetime.strptime(
                     submitted_time_str, "%Y-%m-%dT%H:%M:%SZ"

src/argowrapper/routes/routes.py

@@ -17,6 +17,7 @@
     GEN3_TEAM_PROJECT_METADATA_LABEL,
     GEN3_USER_METADATA_LABEL,
     GEN3_NON_VA_WORKFLOW_MONTHLY_CAP,
+    GEN3_DEFAULT_WORKFLOW_MONTHLY_CAP,
 )
 
 from argowrapper import logger
@@ -209,7 +210,7 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-def check_user_billing_id(request):
+def check_user_info_for_billing_id_and_workflow_limit(request):
     """
     Check whether user is non-VA user
     if user is VA-user, do nothing and proceed
@@ -232,16 +233,26 @@ def check_user_billing_id(request):
         raise exception
     logger.info("Got user info successfully. Checking for billing id..")
 
-    if "tags" in user_info and "billing_id" in user_info["tags"]:
-        billing_id = user_info["tags"]["billing_id"]
-        logger.info("billing id found in user tags: " + billing_id)
-        return billing_id
+    if "tags" in user_info:
+        if "billing_id" in user_info["tags"]:
+            billing_id = user_info["tags"]["billing_id"]
+            logger.info("billing id found in user tags: " + billing_id)
+        else:
+            billing_id = None
+
+        if "workflow_limit" in user_info["tags"]:
+            workflow_limit = user_info["tags"]["workflow_limit"]
+            logger.info("Workflow limit found in user tags: " + workflow_limit)
+        else:
+            workflow_limit = None
+
+        return billing_id, workflow_limit
     else:
-        logger.info("billing id not found.")
-        return None
+        logger.info("User info does not have tags")
+        return None, None
 
 
-def check_user_reached_monthly_workflow_cap(request_token):
+def check_user_reached_monthly_workflow_cap(request_token, billing_id, custom_limit):
     """
     Query Argo service to see how many successful run user already
     have in the current calendar month. If the number is greater than
@@ -252,20 +263,28 @@ def check_user_reached_monthly_workflow_cap(request_token):
         current_month_workflows = argo_engine.get_user_workflows_for_current_month(
             request_token
         )
+        username = argo_engine_helper.get_username_from_token(request_token)
+        if custom_limit:
+            limit = custom_limit
+        else:
+            if billing_id:
+                limit = GEN3_NON_VA_WORKFLOW_MONTHLY_CAP
+            else:
+                limit = GEN3_DEFAULT_WORKFLOW_MONTHLY_CAP
 
-        if len(current_month_workflows) >= GEN3_NON_VA_WORKFLOW_MONTHLY_CAP:
-            logger.info(
-                "User already executed {} workflows this month and cannot create new ones anymore.".format(
-                    len(current_month_workflows)
+        if len(current_month_workflows) >= limit:
+            logger.warn(
+                "This user {} already executed {} workflows this month and cannot create new ones anymore. The currently monthly cap for this user is {}.".format(
+                    username, len(current_month_workflows), limit
                 )
             )
+            return True
+        else:
             logger.info(
-                "The currently monthly cap is {}.".format(
-                    GEN3_NON_VA_WORKFLOW_MONTHLY_CAP
+                "This user {} executed {} workflows this month. The currently monthly cap for this user is {}.".format(
+                    username, len(current_month_workflows), limit
                 )
             )
-            return True
-
         return False
     except Exception as e:
         logger.error(e)
@@ -292,13 +311,14 @@ def submit_workflow(
         reached_monthly_cap = False
 
         # check if user has a billing id tag:
-        billing_id = check_user_billing_id(request)
+        billing_id, workflow_limit = check_user_info_for_billing_id_and_workflow_limit(
+            request
+        )
 
         # if user has billing_id (non-VA user), check if they already reached the monthly cap
-        if billing_id:
-            reached_monthly_cap = check_user_reached_monthly_workflow_cap(
-                request.headers.get("Authorization")
-            )
+        reached_monthly_cap = check_user_reached_monthly_workflow_cap(
+            request.headers.get("Authorization"), billing_id, workflow_limit
+        )
 
         # submit workflow:
         if not reached_monthly_cap:
@@ -311,8 +331,9 @@ def submit_workflow(
                 status_code=HTTP_403_FORBIDDEN,
             )
     except Exception as exception:
+        logger.error(str(exception))
         return HTMLResponse(
-            content=str(exception),
+            content="Unexpected Error Occurred",
             status_code=HTTP_500_INTERNAL_SERVER_ERROR,
         )
 
@@ -331,8 +352,9 @@ def get_workflow_details(
         return argo_engine.get_workflow_details(workflow_name, uid)
 
     except Exception as exception:
+        logger.error(str(exception))
         return HTMLResponse(
-            content=str(exception),
+            content="Unexpected Error Occurred",
             status_code=HTTP_500_INTERNAL_SERVER_ERROR,
         )
 
@@ -354,7 +376,7 @@ def retry_workflow(
         logger.error(traceback.format_exc())
         logger.error(f"could not retry {workflow_name}, failed with error {exception}")
         return HTMLResponse(
-            content=str(exception),
+            content="Could not retry workflow, error occurred",
             status_code=HTTP_500_INTERNAL_SERVER_ERROR,
         )
 
@@ -372,8 +394,9 @@ def cancel_workflow(
         return argo_engine.cancel_workflow(workflow_name)
 
     except Exception as exception:
+        logger.error(str(exception))
         return HTMLResponse(
-            content=str(exception),
+            content="Unexpected Error Occurred",
             status_code=HTTP_500_INTERNAL_SERVER_ERROR,
         )
 
@@ -400,8 +423,9 @@ def get_workflows(
             )
 
     except Exception as exception:
+        logger.error(str(exception))
         return HTMLResponse(
-            content=exception,
+            content="Unexpected Error Occurred",
             status_code=HTTP_500_INTERNAL_SERVER_ERROR,
         )
 
@@ -419,7 +443,8 @@ def get_workflow_logs(
         return argo_engine.get_workflow_logs(workflow_name, uid)
 
     except Exception as exception:
+        logger.error(str(exception))
         return HTMLResponse(
-            content=exception,
+            content="Unexpected Error Occurred",
             status_code=HTTP_500_INTERNAL_SERVER_ERROR,
         )

test/test_argo_engine.py

@@ -825,6 +825,11 @@ def test_get_user_workflows_for_current_month(monkeypatch):
             "phase": "Succeeded",
             "submittedAt": "2023-11-15T17:52:52Z",
         },
+        {
+            "uid": "uid_3",
+            "phase": "Failed",
+            "submittedAt": "2023-11-02T00:00:00Z",
+        },
     ]
     engine.get_workflows_for_label_selector = mock.MagicMock(
         return_value=workflows_mock_response

test/test_routes.py

@@ -106,14 +106,17 @@ def test_submit_workflow(client):
     ) as mock_engine, patch(
         "argowrapper.routes.routes.log_auth_check_type"
     ) as mock_log, patch(
-        "argowrapper.routes.routes.check_user_billing_id"
+        "argowrapper.routes.routes.check_user_info_for_billing_id_and_workflow_limit"
     ) as mock_check_billing_id, patch(
         "requests.get"
-    ) as mock_requests:
+    ) as mock_requests, patch(
+        "argowrapper.routes.routes.check_user_reached_monthly_workflow_cap"
+    ) as mock_check_monthly_cap:
         mock_auth.return_value = True
         mock_engine.return_value = "workflow_123"
-        mock_check_billing_id.return_value = None
+        mock_check_billing_id.return_value = None, None
         mock_requests.side_effect = mocked_requests_get
+        mock_check_monthly_cap.return_value = False
 
         response = client.post(
             "/submit",
@@ -516,10 +519,16 @@ def test_submit_workflow_with_user_billing_id(client):
         "argowrapper.routes.routes.log_auth_check_type"
     ) as mock_log, patch(
         "requests.get"
-    ) as mock_requests:
+    ) as mock_requests, patch(
+        "argowrapper.engine.argo_engine.ArgoEngine.get_user_workflows_for_current_month"
+    ) as mock_get_workflow:
         mock_auth.return_value = True
         mock_engine.return_value = "workflow_123"
         mock_requests.side_effect = mocked_requests_get
+        mock_get_workflow.return_value = [
+            {"wf_name": "workflow1"},
+            {"wf_name": "workflow2"},
+        ]
 
         data["user_tags"] = {"tags": {}}
         response = client.post(
@@ -562,7 +571,7 @@ def test_submit_workflow_with_user_billing_id(client):
             )
             assert mock_engine.call_args.args[2] == "1234"
 
-        data["user_tags"] == 500
+        data["user_tags"] = 500
 
         response = client.post(
             "/submit",
@@ -588,17 +597,45 @@ def test_check_user_reached_monthly_workflow_cap():
             {"wf_name": "workflow1"},
             {"wf_name": "workflow2"},
         ]
+
+        # Test Under Default Limit
+        assert (
+            check_user_reached_monthly_workflow_cap(
+                headers["Authorization"], None, None
+            )
+            == False
+        )
+
+        # Test Custom Limit
         assert (
-            check_user_reached_monthly_workflow_cap(headers["Authorization"]) == False
+            check_user_reached_monthly_workflow_cap(headers["Authorization"], None, 2)
+            == True
         )
 
+        # Test Billing Id User Exceeding Limit
         workflows = []
         for index in range(GEN3_NON_VA_WORKFLOW_MONTHLY_CAP + 1):
             workflows.append({"wf_name": "workflow" + str(index)})
-
         mock_get_workflow.return_value = workflows
 
-        assert check_user_reached_monthly_workflow_cap(headers["Authorization"]) == True
+        assert (
+            check_user_reached_monthly_workflow_cap(
+                headers["Authorization"], "1234", None
+            )
+            == True
+        )
+
+        # Test VA User Exceeding Limit
+        workflows = []
+        for index in range(GEN3_DEFAULT_WORKFLOW_MONTHLY_CAP + 1):
+            workflows.append({"wf_name": "workflow" + str(index)})
+        mock_get_workflow.return_value = workflows
+        assert (
+            check_user_reached_monthly_workflow_cap(
+                headers["Authorization"], None, None
+            )
+            == True
+        )
 
 
 def test_submit_workflow_with_billing_id_and_over_monthly_cap(client):
@@ -607,15 +644,44 @@ def test_submit_workflow_with_billing_id_and_over_monthly_cap(client):
     ) as mock_engine, patch(
         "argowrapper.routes.routes.log_auth_check_type"
     ) as mock_log, patch(
-        "argowrapper.routes.routes.check_user_billing_id"
+        "argowrapper.routes.routes.check_user_info_for_billing_id_and_workflow_limit"
+    ) as mock_check_billing_id, patch(
+        "argowrapper.routes.routes.check_user_reached_monthly_workflow_cap"
+    ) as mock_check_monthly_cap, patch(
+        "requests.get"
+    ) as mock_requests:
+        mock_auth.return_value = True
+        mock_engine.return_value = "workflow_123"
+        mock_check_billing_id.return_value = "1234", None
+        mock_check_monthly_cap.return_value = True
+        mock_requests.side_effect = mocked_requests_get
+
+        response = client.post(
+            "/submit",
+            data=json.dumps(data),
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": EXAMPLE_AUTH_HEADER,
+            },
+        )
+        assert response.status_code == 403
+
+
+def test_submit_workflow_over_monthly_cap(client):
+    with patch("argowrapper.routes.routes.auth.authenticate") as mock_auth, patch(
+        "argowrapper.routes.routes.argo_engine.workflow_submission"
+    ) as mock_engine, patch(
+        "argowrapper.routes.routes.log_auth_check_type"
+    ) as mock_log, patch(
+        "argowrapper.routes.routes.check_user_info_for_billing_id_and_workflow_limit"
     ) as mock_check_billing_id, patch(
         "argowrapper.routes.routes.check_user_reached_monthly_workflow_cap"
     ) as mock_check_monthly_cap, patch(
         "requests.get"
     ) as mock_requests:
         mock_auth.return_value = True
         mock_engine.return_value = "workflow_123"
-        mock_check_billing_id.return_value = "1234"
+        mock_check_billing_id.return_value = None, None
         mock_check_monthly_cap.return_value = True
         mock_requests.side_effect = mocked_requests_get
 
@@ -636,13 +702,13 @@ def test_submit_workflow_with_non_team_project_cohort(client):
     ) as mock_engine, patch(
         "argowrapper.routes.routes.log_auth_check_type"
     ) as mock_log, patch(
-        "argowrapper.routes.routes.check_user_billing_id"
+        "argowrapper.routes.routes.check_user_info_for_billing_id_and_workflow_limit"
     ) as mock_check_billing_id, patch(
         "requests.get"
     ) as mock_requests:
         mock_auth.return_value = True
         mock_engine.return_value = "workflow_123"
-        mock_check_billing_id.return_value = None
+        mock_check_billing_id.return_value = None, None
         mock_requests.side_effect = mocked_requests_get
 
         data["outcome"]["cohort_ids"] = [400]