Merge pull request #37 from kartoza/fix-worker-exceptions

Fix worker logs + exceptions
kartoza · May 30, 2024 · 6e025b0 · 6e025b0
2 parents cf0b50b + 012976e
commit 6e025b0
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/django_project/cplus_api/api_views/scenario.py b/django_project/cplus_api/api_views/scenario.py
@@ -6,7 +6,7 @@
 from drf_yasg import openapi
 from drf_yasg.utils import swagger_auto_schema
 from django.shortcuts import get_object_or_404
-from django.core.exceptions import ValidationError
+from rest_framework.exceptions import ValidationError
 from django.contrib.contenttypes.models import ContentType
 from django.core.paginator import Paginator
 from core.celery import cancel_task

diff --git a/django_project/cplus_api/utils/worker_analysis.py b/django_project/cplus_api/utils/worker_analysis.py
@@ -399,6 +399,7 @@ def __init__(self, task_config: TaskConfig,
         self.scenario_task = scenario_task
         self.last_update_progress = None
         self.downloaded_layers = {}
+        self.downloaded_layer_count = 0
 
     def prepare_run(self):
         # clear existing scenario directory if exists
@@ -416,6 +417,8 @@ def prepare_run(self):
     def initialize_input_layers(self, scenario_path: str):
         self.log_message(
             f'Initialize input layers: {self.task_config.total_input_layers}')
+        self.set_custom_progress(0)
+        self.set_status_message('Preparing input layers')
 
         # init priority layers
         priority_layer_paths = {}
@@ -522,6 +525,7 @@ def initialize_input_layers(self, scenario_path: str):
             'Finished copy input layers: '
             f'{self.task_config.total_input_layers}'
         )
+        self.set_custom_progress(100)
 
     def copy_input_layers_by_uuids(
             self, component_type: InputLayer.ComponentTypes,
@@ -534,8 +538,19 @@ def copy_input_layers_by_uuids(
             layers = layers.filter(
                 component_type=component_type
             )
+        total_input_layers = (
+            self.task_config.total_input_layers if
+            self.task_config.total_input_layers > 0 else 1
+        )
         for layer in layers:
             file_path = layer.download_to_working_directory(scenario_path)
+            self.downloaded_layer_count += 1
+            self.set_custom_progress(
+                100 * (
+                    self.downloaded_layer_count /
+                    total_input_layers
+                )
+            )
             if not file_path:
                 continue
             if not os.path.exists(file_path):
@@ -771,7 +786,6 @@ def notify_user(self, is_success: bool):
 
     def finished(self, result: bool):
         if result:
-            self.log_message("Finished from the main task \n")
             self.upload_scenario_outputs()
         else:
             self.log_message(