Merge pull request #493 from xchem/m2ms-1226-xca-uploader

Target upload tweaks for better cooperation with XCA uploader

Author: kaliif

viewer/target_loader.py

@@ -72,6 +72,7 @@ class UploadState(str, Enum):
     REPORTING = "REPORTING"
     SUCCESS = "SUCCESS"
     FAILED = "FAILED"
+    CANCELED = "CANCELED"
 
 
 class Level(str, Enum):
@@ -150,13 +151,14 @@ def log(self, level: Level, message: str) -> None:
         self.stack.append(UploadReportEntry(level=level, message=message))
         self._update_task(message)
 
-    def final(self, archive_name):
-        if self.upload_state == UploadState.PROCESSING:
-            self.upload_state = UploadState.SUCCESS
-            message = f"{archive_name} uploaded successfully."
+    def final(self, message, upload_state=None):
+        if upload_state:
+            self.upload_state = upload_state
         else:
-            self.upload_state = UploadState.FAILED
-            message = f"Uploading {archive_name} failed."
+            if self.upload_state == UploadState.PROCESSING:
+                self.upload_state = UploadState.SUCCESS
+            else:
+                self.upload_state = UploadState.FAILED
 
         self.stack.append(UploadReportEntry(message=message))
         self._update_task(self.json())
@@ -167,6 +169,7 @@ def json(self):
     def _update_task(self, message: str | list) -> None:
         if self.task:
             try:
+                logger.debug("taskstuff %s", dir(self.task))
                 self.task.update_state(
                     state=self.upload_state,
                     meta={
@@ -1134,7 +1137,9 @@ def process_metadata(
         # moved this bit from init
         self.target, target_created = Target.objects.get_or_create(
             title=self.target_name,
-            display_name=self.target_name,
+            defaults={
+                "display_name": self.target_name,
+            },
         )
 
         # TODO: original target loader's function get_create_projects
@@ -1153,7 +1158,7 @@ def process_metadata(
             # remove uploaded file
             Path(self.bundle_path).unlink()
             msg = f"{self.bundle_name} already uploaded, skipping."
-            self.report.log(Level.INFO, msg)
+            self.report.final(msg, upload_state=UploadState.CANCELED)
             raise FileExistsError(msg)
 
         if project_created and committer.pk == settings.ANONYMOUS_USER:
@@ -1596,7 +1601,7 @@ def load_target(
                     )
         except IntegrityError as exc:
             logger.error(exc, exc_info=True)
-            target_loader.report.final(target_loader.data_bundle)
+            target_loader.report.final(f"Uploading {target_loader.data_bundle} failed")
             target_loader.experiment_upload.message = target_loader.report.json()
             raise IntegrityError from exc
 
@@ -1618,7 +1623,7 @@ def load_target(
 
         set_directory_permissions(target_loader.abs_final_path, 0o755)
 
-        target_loader.report.final(target_loader.data_bundle)
+        target_loader.report.final(f"{target_loader.data_bundle} uploaded successfully")
         target_loader.experiment_upload.message = target_loader.report.json()
 
         # logger.debug("%s", upload_report)

viewer/views.py

@@ -12,6 +12,7 @@
 
 import pandas as pd
 import pytz
+from celery import Celery
 from celery.result import AsyncResult
 from dateutil.parser import parse
 from django.conf import settings
@@ -1578,15 +1579,39 @@ def get(self, request, task_id, *args, **kwargs):
 
         # task_id is (will be) a UUID, but Celery expects a string
         task_id_str = str(task_id)
-        try:
-            result = AsyncResult(task_id_str)
-        except TimeoutError:
-            error = {'error': 'Task result query timed out'}
-            return Response(error, status=status.HTTP_408_REQUEST_TIMEOUT)
-
-        if result.state == 'PENDING':
-            error = {'error': 'Unknown task'}
-            return Response(error, status=status.HTTP_400_BAD_REQUEST)
+
+        celery_app = Celery("fragalysis")
+        inspect = celery_app.control.inspect()
+        ping = inspect.ping()
+
+        if ping:
+            active_tasks = inspect.active()
+
+            # active_tasks.values is a list of tasks for every worker
+            if task_id_str in [
+                k['id'] for worker in active_tasks.values() for k in worker
+            ]:
+                # celery confirms task exists
+                try:
+                    result = AsyncResult(task_id_str)
+                except TimeoutError:
+                    error = {'error': 'Task result query timed out'}
+                    return Response(error, status=status.HTTP_408_REQUEST_TIMEOUT)
+            else:
+                # no such task
+                error = {'error': 'Unknown task'}
+                return Response(error, status=status.HTTP_400_BAD_REQUEST)
+
+        else:
+            # task scheduler not running. This may be the case in local
+            # development, but this means there's really no way to
+            # validate whether the task exists
+            logger.warning('Celery not running!')
+            try:
+                result = AsyncResult(task_id_str)
+            except TimeoutError:
+                error = {'error': 'Task result query timed out'}
+                return Response(error, status=status.HTTP_408_REQUEST_TIMEOUT)
 
         # Extract messages (from task info)
         # Assuming the task has some info.