Terminate job if a mandatory task fails

cfa_azure/clients.py

@@ -1045,7 +1045,6 @@ def add_job(
         pool_name: str | None = None,
         save_logs_to_blob: str | None = None,
         logs_folder: str | None = None,
-        end_job_on_task_failure: bool = False,
         task_retries: int = 0,
         mark_complete_after_tasks_run: bool = False,
     ) -> None:
@@ -1056,7 +1055,6 @@ def add_job(
             pool_name (str|None): pool to use for job. If None, will used self.pool_name from client. Default None.
             save_logs_to_blob (str): the name of the blob container. Must be mounted to the pool. Default None for no saving.
             logs_folder (str): the folder structure to use when saving logs to blob. Default None will save to /stdout_stderr/ folder in specified blob container.
-            end_job_on_task_failure (bool): whether to end the job if a task fails. Default False.
             task_retries (int): number of times to retry a task that fails. Default 0.
             mark_complete_after_tasks_run (bool): whether to mark the job as completed when all tasks finish running. Default False.
         """
@@ -1089,7 +1087,6 @@ def add_job(
         helpers.add_job(
             job_id=job_id_r,
             pool_id=p_name,
-            end_job_on_task_failure=end_job_on_task_failure,
             batch_client=self.batch_client,
             task_retries=task_retries,
             mark_complete=mark_complete_after_tasks_run,

cfa_azure/helpers.py

@@ -24,6 +24,7 @@
     ExitOptions,
     JobAction,
     JobConstraints,
+    MetadataItem,
     OnAllTasksComplete,
     OnTaskFailure,
 )
@@ -728,7 +729,6 @@ def list_nodes_by_pool(pool_name: str, config: dict, node_state: str = None):
 def add_job(
     job_id: str,
     pool_id: str,
-    end_job_on_task_failure: bool,
     batch_client: object,
     task_retries: int = 0,
     mark_complete: bool = False,
@@ -738,7 +738,6 @@ def add_job(
     Args:
         job_id (str): name of the job to run
         pool_id (str): name of pool
-        end_job_on_task_failure (bool): whether to end a running job if a task fails
         batch_client (object): batch client object
         task_retries (int): number of times to retry the task if it fails. Default 3.
         mark_complete (bool): whether to mark the job complete after tasks finish running. Default False.
@@ -758,6 +757,7 @@ def add_job(
         on_all_tasks_complete=on_all_tasks_complete,
         on_task_failure=OnTaskFailure.perform_exit_options_job_action,
         constraints=job_constraints,
+        metadata=[MetadataItem(name="mark_complete", value=mark_complete)],
     )
     logger.debug("Attempting to add job.")
     try:
@@ -833,8 +833,20 @@ def add_task_to_job(
             logger.debug("Adding task dependency.")
         task_deps = batchmodels.TaskDependencies(task_ids=depends_on)
 
+    job_action = JobAction.none
+    if check_job_exists(job_id, batch_client):
+        job_details = batch_client.job.get(job_id)
+        if job_details and job_details.metadata:
+            for metadata in job_details.metadata:
+                if (
+                    metadata.name == "mark_complete"
+                    and bool(metadata.value) is True
+                ):
+                    job_action = JobAction.terminate
+                    break
+
     no_exit_options = ExitOptions(
-        dependency_action=DependencyAction.satisfy, job_action=JobAction.none
+        dependency_action=DependencyAction.satisfy, job_action=job_action
     )
     if run_dependent_tasks_on_fail:
         exit_conditions = ExitConditions(
@@ -849,7 +861,7 @@ def add_task_to_job(
     else:
         terminate_exit_options = ExitOptions(
             dependency_action=DependencyAction.block,
-            job_action=JobAction.none,
+            job_action=job_action,
         )
         exit_conditions = ExitConditions(
             exit_codes=[

tests/fake_client.py

@@ -114,6 +114,12 @@ def delete(self, *args):
         def add(self, job):
             return True
 
+        def get(self, job):
+            return True
+
+        def list(self):
+            return []
+
     class FakeTag:
         def __init__(self, tag):
             self.name = tag

tests/helpers_tests.py

@@ -724,10 +724,7 @@ def test_add_job(self, mock_logger):
         batch_client = FakeClient()
         job_id = "my_job_id"
         cfa_azure.helpers.add_job(
-            job_id,
-            FAKE_BATCH_POOL,
-            batch_client=batch_client,
-            end_job_on_task_failure=False,
+            job_id, FAKE_BATCH_POOL, batch_client=batch_client
         )
         mock_logger.info.assert_called_with(
             f"Job '{job_id}' created successfully."
@@ -738,10 +735,7 @@ def test_add_job_task_failure(self, mock_logger):
         batch_client = FakeClient()
         job_id = "my_job_id"
         cfa_azure.helpers.add_job(
-            job_id,
-            FAKE_BATCH_POOL,
-            batch_client=batch_client,
-            end_job_on_task_failure=False,
+            job_id, FAKE_BATCH_POOL, batch_client=batch_client
         )
         mock_logger.debug.assert_called_with("Attempting to add job.")