Task status v2 (#1558)

* Various updates to task management * Updates * Updates * Updates to task manager * Update unit test * Minor updates to use celery states
google · Oct 1, 2024 · c290576 · c290576
1 parent 37889ef
commit c290576
Show file tree

Hide file tree

Showing 7 changed files with 76 additions and 74 deletions.
diff --git a/turbinia/state_manager.py b/turbinia/state_manager.py
@@ -22,6 +22,7 @@
 import json
 import logging
 import sys
+from celery import states as celery_states
 from datetime import datetime
 from datetime import timedelta
 from typing import Any, List, Dict, Optional
@@ -269,7 +270,6 @@ def update_request_task(self, task) -> None:
       task (TurbiniaTask): Turbinia task object.
     """
     request_key = self.redis_client.build_key_name('request', task.request_id)
-    task_key = self.redis_client.build_key_name('task', task.id)
     try:
       self.redis_client.add_to_list(request_key, 'task_ids', task.id)
       request_last_update = datetime.strptime(
@@ -291,12 +291,12 @@ def update_request_task(self, task) -> None:
         elif task.result.successful is False:
           self.redis_client.add_to_list(request_key, 'failed_tasks', task.id)
           statuses_to_remove.remove('failed_tasks')
-      task_status = self.redis_client.get_attribute(task_key, 'status')
-      if task_status and 'Task is running' in task_status:
+      task_status = task.celery_state
+      if task_status == celery_states.STARTED:
         self.redis_client.add_to_list(request_key, 'running_tasks', task.id)
         statuses_to_remove.remove('running_tasks')
-      elif (task_status is None or task_status == 'queued' or
-            task_status == 'pending'):
+      elif (task_status is None or task_status == celery_states.RECEIVED or
+            task_status == celery_states.PENDING):
         self.redis_client.add_to_list(request_key, 'queued_tasks', task.id)
         statuses_to_remove.remove('queued_tasks')
       for status_name in statuses_to_remove:

diff --git a/turbinia/task_manager.py b/turbinia/task_manager.py
@@ -687,6 +687,14 @@ def _backend_setup(self, *args, **kwargs):
       self.celery_runner = self.celery.app.task(
           task_utils.task_runner, name='task_runner')
 
+  def _get_worker_name(self, celery_stub):
+    """Gets the Celery worker name from the AsyncResult object."""
+    worker_name = celery_stub.result.get('hostname', None)
+    if worker_name:
+      # example hostname: celery@turbinia-hostname
+      worker_name = worker_name.split('@')[1]
+    return worker_name
+
   def process_tasks(self):
     """Determine the current state of our tasks.
 
@@ -700,28 +708,33 @@ def process_tasks(self):
       celery_task = task.stub
       # ref: https://docs.celeryq.dev/en/stable/reference/celery.states.html
       if not celery_task:
-        log.debug(f'Task {task.stub.task_id:s} not yet created.')
+        log.info(f'Task {task.stub.task_id:s} not yet created.')
         check_timeout = True
       elif celery_task.status == celery_states.STARTED:
-        # Task status will be set to running when the worker executes run_wrapper()
-        log.debug(f'Task {celery_task.id:s} not finished.')
+        log.warning(f'Task {celery_task.id:s} {task.id} started.')
+        task.worker_name = self._get_worker_name(celery_task)
+        task.celery_state = celery_states.STARTED
         check_timeout = True
       elif celery_task.status == celery_states.FAILURE:
-        log.warning(f'Task {celery_task.id:s} failed.')
+        log.info(f'Task {celery_task.id:s} failed.')
+        task.celery_state = celery_states.FAILURE
         self.close_failed_task(task)
         completed_tasks.append(task)
       elif celery_task.status == celery_states.SUCCESS:
+        task.celery_state = celery_states.SUCCESS
         task.result = workers.TurbiniaTaskResult.deserialize(celery_task.result)
         completed_tasks.append(task)
       elif celery_task.status == celery_states.PENDING:
-        task.status = 'pending'
+        log.info(f'Task {celery_task.id:s} is pending.')
+        task.celery_state = celery_states.PENDING
         check_timeout = True
-        log.debug(f'Task {celery_task.id:s} is pending.')
       elif celery_task.status == celery_states.RECEIVED:
-        task.status = 'queued'
+        log.info(f'Task {celery_task.id:s} is queued.')
+        task.worker_name = self._get_worker_name(celery_task)
+        task.celery_state = celery_states.RECEIVED
         check_timeout = True
-        log.debug(f'Task {celery_task.id:s} is queued.')
       elif celery_task.status == celery_states.REVOKED:
+        task.celery_state = celery_states.REVOKED
         message = (
             f'Celery task {celery_task.id:s} associated with Turbinia '
             f'task {task.id} was revoked. This could be caused if the task is '
@@ -746,6 +759,9 @@ def process_tasks(self):
           task = self.timeout_task(task, timeout)
           completed_tasks.append(task)
 
+      # Update task metadata so we have an up to date state.
+      self.state_manager.update_task(task)
+
     outstanding_task_count = len(self.tasks) - len(completed_tasks)
     if outstanding_task_count > 0:
       log.info(f'{outstanding_task_count:d} Tasks still outstanding.')

diff --git a/turbinia/tcelery.py b/turbinia/tcelery.py
@@ -57,6 +57,7 @@ def setup(self):
         task_default_queue=config.INSTANCE_ID,
         # Re-queue task if Celery worker abruptly exists
         task_reject_on_worker_lost=True,
+        task_track_started=True,
         worker_cancel_long_running_tasks_on_connection_loss=True,
         worker_concurrency=1,
         worker_prefetch_multiplier=1,

diff --git a/turbinia/workers/__init__.py b/turbinia/workers/__init__.py
@@ -129,8 +129,8 @@ class TurbiniaTaskResult:
 
   # The list of attributes that we will persist into storage
   STORED_ATTRIBUTES = [
-      'worker_name', 'report_data', 'report_priority', 'run_time', 'status',
-      'saved_paths', 'successful', 'evidence_size'
+      'report_data', 'report_priority', 'run_time', 'status', 'saved_paths',
+      'successful', 'evidence_size', 'worker_name'
   ]
 
   def __init__(
@@ -454,7 +454,7 @@ class TurbiniaTask:
   STORED_ATTRIBUTES = [
       'id', 'job_id', 'job_name', 'start_time', 'last_update', 'name',
       'evidence_name', 'evidence_id', 'request_id', 'requester', 'group_name',
-      'reason', 'group_id', 'celery_id'
+      'reason', 'group_id', 'celery_id', 'celery_state', 'worker_name'
   ]
 
   # The list of evidence states that are required by a Task in order to run.
@@ -483,6 +483,7 @@ def __init__(
 
     self.id = uuid.uuid4().hex
     self.celery_id = None
+    self.celery_state = None
     self.is_finalize_task = False
     self.job_id = None
     self.job_name = None
@@ -508,6 +509,7 @@ def __init__(
     self.group_name = group_name
     self.reason = reason
     self.group_id = group_id
+    self.worker_name = None
 
   def serialize(self):
     """Converts the TurbiniaTask object into a serializable dict.
@@ -821,6 +823,7 @@ def setup(self, evidence):
       TurbiniaException: If the evidence can not be found.
     """
     self.setup_metrics()
+    self.state_manager = state_manager.get_state_manager()
     self.output_manager.setup(self.name, self.id, self.request_id)
     self.tmp_dir, self.output_dir = self.output_manager.get_local_output_dirs()
     if not self.result:
@@ -1039,21 +1042,10 @@ def run_wrapper(self, evidence):
 
     log.debug(f'Task {self.name:s} {self.id:s} awaiting execution')
     failure_message = None
-    worker_name = platform.node()
+    self.worker_name = platform.node()
     try:
       evidence = evidence_decode(evidence)
       self.result = self.setup(evidence)
-      # Call update_task_status to update status
-      # We cannot call update_task() here since it will clobber previously
-      # stored data by the Turbinia server when the task was created, which is
-      # not present in the TurbiniaTask object the worker currently has in its
-      # runtime.
-      self.update_task_status(self, 'queued')
-      # Beucase of the same reason, we perform a single attribute update
-      # for the worker name.
-      task_key = self.state_manager.redis_client.build_key_name('task', self.id)
-      self.state_manager.redis_client.set_attribute(
-          task_key, 'worker_name', json.dumps(worker_name))
       turbinia_worker_tasks_queued_total.inc()
       task_runtime_metrics = self.get_metrics()
     except TurbiniaException as exception:
@@ -1108,9 +1100,6 @@ def run_wrapper(self, evidence):
         self._evidence_config = evidence.config
         self.task_config = self.get_task_recipe(evidence.config)
         self.worker_start_time = datetime.now()
-        # Update task status so we know which worker the task executed on.
-        updated_status = f'Task is running on {worker_name}'
-        self.update_task_status(self, updated_status)
         self.result = self.run(evidence, self.result)
 
       # pylint: disable=broad-except
@@ -1161,7 +1150,7 @@ def run_wrapper(self, evidence):
         failure_message = (
             'Task Result was auto-closed from task executor on {0:s} likely '
             'due to previous failures.  Previous status: [{1:s}]'.format(
-                self.result.worker_name, status))
+                self.worker_name, status))
       self.result.log(failure_message)
       try:
         self.result.close(self, False, failure_message)
@@ -1197,23 +1186,3 @@ def run(self, evidence, result):
         TurbiniaTaskResult object.
     """
     raise NotImplementedError
-
-  def update_task_status(self, task, status=None):
-    """Updates the task status and pushes it directly to datastore.
-
-    Args:
-      task (TurbiniaTask): The calling Task object
-      status (str): Brief word or phrase for Task state. If not supplied, the
-          existing Task status will be used.
-    """
-    if not status:
-      return
-    if not self.state_manager:
-      self.state_manager = state_manager.get_state_manager()
-    if self.state_manager:
-      task_key = self.state_manager.redis_client.build_key_name('task', task.id)
-      self.state_manager.redis_client.set_attribute(
-          task_key, 'status', json.dumps(status))
-      self.state_manager.update_request_task(task)
-    else:
-      log.info('No state_manager initialized, not updating Task info')
diff --git a/turbinia/workers/workers_test.py b/turbinia/workers/workers_test.py
@@ -106,7 +106,6 @@ def setResults(
 
     self.result.input_evidence = evidence.RawDisk()
     self.result.status = 'TestStatus'
-    self.update_task_status = mock.MagicMock()
     self.result.close = mock.MagicMock()
     self.task.setup = mock.MagicMock(return_value=setup)
     self.result.worker_name = 'worker1'
@@ -147,8 +146,10 @@ def testTurbiniaTaskSerialize(self):
     self.assertEqual(out_obj.__dict__, self.plaso_task.__dict__)
 
   @mock.patch('turbinia.redis_client.RedisClient.set_attribute')
-  @mock.patch('turbinia.workers.TurbiniaTask.update_task_status')
-  def testTurbiniaTaskRunWrapper(self, _, __):
+  def testTurbiniaTaskRunWrapper(
+      self,
+      _,
+  ):
     """Test that the run wrapper executes task run."""
     self.setResults()
     self.result.closed = True
@@ -159,8 +160,10 @@ def testTurbiniaTaskRunWrapper(self, _, __):
     self.result.close.assert_not_called()
 
   @mock.patch('turbinia.redis_client.RedisClient.set_attribute')
-  @mock.patch('turbinia.workers.TurbiniaTask.update_task_status')
-  def testTurbiniaTaskRunWrapperAutoClose(self, _, __):
+  def testTurbiniaTaskRunWrapperAutoClose(
+      self,
+      _,
+  ):
     """Test that the run wrapper closes the task."""
     self.setResults()
     new_result = self.task.run_wrapper(self.evidence.__dict__)
@@ -170,8 +173,7 @@ def testTurbiniaTaskRunWrapperAutoClose(self, _, __):
 
   @mock.patch('turbinia.state_manager.get_state_manager')
   @mock.patch('turbinia.redis_client.RedisClient.set_attribute')
-  @mock.patch('turbinia.workers.TurbiniaTask.update_task_status')
-  def testTurbiniaTaskRunWrapperBadResult(self, _, __, ___):
+  def testTurbiniaTaskRunWrapperBadResult(self, _, __):
     """Test that the run wrapper recovers from run returning bad result."""
     bad_result = 'Not a TurbiniaTaskResult'
     checked_result = TurbiniaTaskResult(base_output_dir=self.base_output_dir)
@@ -185,8 +187,7 @@ def testTurbiniaTaskRunWrapperBadResult(self, _, __, ___):
     self.assertIn('CheckedResult', new_result.status)
 
   @mock.patch('turbinia.redis_client.RedisClient.set_attribute')
-  @mock.patch('turbinia.workers.TurbiniaTask.update_task_status')
-  def testTurbiniaTaskJobUnavailable(self, _, __):
+  def testTurbiniaTaskJobUnavailable(self, _):
     """Test that the run wrapper can fail if the job doesn't exist."""
     self.setResults()
     self.task.job_name = 'non_exist'
@@ -198,8 +199,7 @@ def testTurbiniaTaskJobUnavailable(self, _, __):
     self.assertEqual(new_result.status, canary_status)
 
   @mock.patch('turbinia.redis_client.RedisClient.set_attribute')
-  @mock.patch('turbinia.workers.TurbiniaTask.update_task_status')
-  def testTurbiniaTaskRunWrapperExceptionThrown(self, _, __):
+  def testTurbiniaTaskRunWrapperExceptionThrown(self, _):
     """Test that the run wrapper recovers from run throwing an exception."""
     self.setResults()
     self.task.run = mock.MagicMock(side_effect=TurbiniaException)
@@ -253,9 +253,7 @@ def testTurbiniaTaskValidateResultBadResult(self, _, __):
 
   @mock.patch('turbinia.workers.evidence_decode')
   @mock.patch('turbinia.redis_client.RedisClient.set_attribute')
-  @mock.patch('turbinia.workers.TurbiniaTask.update_task_status')
-  def testTurbiniaTaskEvidenceValidationFailure(
-      self, _, __, evidence_decode_mock):
+  def testTurbiniaTaskEvidenceValidationFailure(self, _, evidence_decode_mock):
     """Tests Task fails when evidence validation fails."""
     self.setResults()
     test_evidence = evidence.RawDisk()

diff --git a/web/src/components/TaskDetails.vue b/web/src/components/TaskDetails.vue
@@ -39,11 +39,17 @@ limitations under the License.
     <v-alert v-if="taskDetails.successful === true" type="success" prominent>
       {{ taskDetails.status }}
     </v-alert>
-    <v-alert v-else-if="taskDetails.successful === null" type="info" prominent>
-      <div v-if="taskDetails.status">
-        {{ taskDetails.status }}
-      </div>
-      <div v-else>Task {{ taskDetails.id }} is pending</div>
+    <v-alert v-else-if="taskDetails.successful === false" type="error" prominent>
+      {{ taskDetails.status }}
+    </v-alert>
+    <v-alert v-else-if="taskDetails.celery_state === 'STARTED'" type="info" prominent>
+      Task {{ taskDetails.id }} is running on {{ taskDetails.worker_name }}
+    </v-alert>
+    <v-alert v-else-if="taskDetails.celery_state === 'PENDING'" type="info" prominent>
+      Task {{ taskDetails.id }} is pending.
+    </v-alert>
+    <v-alert v-else-if="taskDetails.celery_state === 'RECEIVED'" type="info" prominent>
+      Task {{ taskDetails.id }} is queued.
     </v-alert>
     <v-alert v-else type="error" prominent>
       {{ taskDetails.status }}
@@ -72,6 +78,12 @@ limitations under the License.
             </div>
             <div v-else>N/A</div>
           </v-list-item>
+          <v-list-item title="Celery State:">
+            <div v-if="taskDetails.celery_state">
+              {{ taskDetails.celery_state }}
+            </div>
+            <div v-else>N/A</div>
+          </v-list-item>
           <v-list-item title="Evidence ID:">
             <div v-if="taskDetails.evidence_id">
               {{ taskDetails.evidence_id }}

diff --git a/web/src/components/TaskList.vue b/web/src/components/TaskList.vue
@@ -73,14 +73,20 @@ export default {
           let data = response.data['tasks']
           for (const task in data) {
             let task_dict = data[task]
-            let taskStatusTemp = task_dict.status
+            let taskStatusTemp = task_dict.celery_state
             // As pending status requests show as null or pending
-            if (taskStatusTemp === null || taskStatusTemp === "pending") {
+            if (taskStatusTemp === null || taskStatusTemp === "PENDING") {
               taskStatusTemp = 'is pending on server.'
             }
-            else if (taskStatusTemp == "queued") {
+            else if (taskStatusTemp == "RECEIVED") {
               taskStatusTemp = 'is queued for execution.'
             }
+            else if (taskStatusTemp == "STARTED") {
+              taskStatusTemp = 'is running on ' + task_dict.worker_name
+            }
+            else {
+              taskStatusTemp = task_dict.status
+            }
             if (this.filterJobs.length > 0) {
               let jobName = task_dict.job_name.toLowerCase()
               if (this.radioFilter && !this.filterJobs.includes(jobName)) {