climate-ref-celery 0.5.0__tar.gz → 0.5.2__tar.gz

{climate_ref_celery-0.5.0 → climate_ref_celery-0.5.2}/.gitignore

@@ -74,7 +74,6 @@ coverage.xml
 *.pot
 
 # Django stuff:
-*.log
 local_settings.py
 db.sqlite3
 db.sqlite3-journal

{climate_ref_celery-0.5.0 → climate_ref_celery-0.5.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: climate-ref-celery
-Version: 0.5.0
+Version: 0.5.2
 Summary: Celery app for mananging tasks and workers
 Author-email: Jared Lewis <jared.lewis@climate-resource.com>
 License-File: LICENCE

{climate_ref_celery-0.5.0 → climate_ref_celery-0.5.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "climate-ref-celery"
-version = "0.5.0"
+version = "0.5.2"
 description = "Celery app for mananging tasks and workers"
 readme = "README.md"
 authors = [

{climate_ref_celery-0.5.0 → climate_ref_celery-0.5.2}/src/climate_ref_celery/cli.py

@@ -54,6 +54,7 @@ def import_provider(provider_package: str) -> DiagnosticProvider:
 
 @app.command()
 def start_worker(
+    ctx: typer.Context,
     loglevel: str = typer.Option("info", help="Log level for the worker"),
     package: str | None = typer.Option(help="Package to import tasks from", default=None),
     extra_args: list[str] = typer.Argument(None, help="Additional arguments for the worker"),
@@ -75,6 +76,10 @@ def start_worker(
         # Attempt to import the provider
         provider = import_provider(package)
 
+        if hasattr(ctx.obj, "config"):
+            # Configure the provider so that it knows where the conda environments are
+            provider.configure(ctx.obj.config)
+
         # Wrap each diagnostics in the provider with a celery tasks
         register_celery_tasks(celery_app, provider)
         queue = provider.slug
@@ -84,7 +89,7 @@ def start_worker(
 
         queue = "celery"
 
-    argv = ["worker", f"--loglevel={loglevel}", f"--queues={queue}", *(extra_args or [])]
+    argv = ["worker", "-E", f"--loglevel={loglevel}", f"--queues={queue}", *(extra_args or [])]
     celery_app.worker_main(argv=argv)
 
 

{climate_ref_celery-0.5.0 → climate_ref_celery-0.5.2}/src/climate_ref_celery/executor.py

@@ -14,9 +14,8 @@ from climate_ref.config import Config
 from climate_ref.models import Execution
 from climate_ref_celery.app import app
 from climate_ref_celery.tasks import generate_task_name
-from climate_ref_core.diagnostics import Diagnostic, ExecutionDefinition, ExecutionResult
+from climate_ref_core.diagnostics import ExecutionDefinition, ExecutionResult
 from climate_ref_core.executor import Executor
-from climate_ref_core.providers import DiagnosticProvider
 
 
 class CeleryExecutor(Executor):
@@ -43,8 +42,6 @@ class CeleryExecutor(Executor):
 
     def run(
         self,
-        provider: DiagnosticProvider,
-        diagnostic: Diagnostic,
         definition: ExecutionDefinition,
         execution: Execution | None = None,
     ) -> None:
@@ -62,10 +59,6 @@ class CeleryExecutor(Executor):
 
         Parameters
         ----------
-        provider
-            Provider for the diagnostic
-        diagnostic
-            Diagnostic to run
         definition
             A description of the information needed for this execution of the diagnostic
             This includes relative paths to the data files,
@@ -78,12 +71,14 @@ class CeleryExecutor(Executor):
         """
         from climate_ref_celery.worker_tasks import handle_result
 
-        name = generate_task_name(provider, diagnostic)
+        diagnostic = definition.diagnostic
+
+        name = generate_task_name(diagnostic.provider, diagnostic)
 
         async_result = app.send_task(
             name,
             args=[definition, self.config.log_level],
-            queue=provider.slug,
+            queue=diagnostic.provider.slug,
             link=handle_result.s(execution_id=execution.id).set(queue="celery") if execution else None,
         )
         logger.debug(f"Celery task {async_result.id} submitted")

{climate_ref_celery-0.5.0 → climate_ref_celery-0.5.2}/src/climate_ref_celery/tasks.py

@@ -13,13 +13,10 @@ The main process is responsible for tracking what diagnostics have been register
 and to respond to new workers coming online.
 """
 
-from collections.abc import Callable
-
 from celery import Celery
-from loguru import logger
 
-from climate_ref_core.diagnostics import Diagnostic, ExecutionDefinition, ExecutionResult
-from climate_ref_core.logging import redirect_logs
+from climate_ref_core.diagnostics import Diagnostic
+from climate_ref_core.executor import execute_locally
 from climate_ref_core.providers import DiagnosticProvider
 
 
@@ -30,32 +27,6 @@ def generate_task_name(provider: DiagnosticProvider, diagnostic: Diagnostic) ->
     return f"{provider.slug}.{diagnostic.slug}"
 
 
-def _diagnostic_task_factory(
-    diagnostic: Diagnostic,
-) -> Callable[
-    [ExecutionDefinition, str],
-    ExecutionResult,
-]:
-    """
-    Create a new task for the given diagnostic
-    """
-
-    def task(definition: ExecutionDefinition, log_level: str) -> ExecutionResult:
-        """
-        Task to run the diagnostic
-        """
-        logger.info(f"Running diagnostic {diagnostic.name} with definition {definition}")
-        try:
-            with redirect_logs(definition, log_level):
-                return diagnostic.run(definition)
-        except Exception:
-            logger.exception(f"Error running diagnostic {diagnostic.slug}:{definition.key}")
-            # TODO: This exception should be caught and a unsuccessful result returned.
-            raise
-
-    return task
-
-
 def register_celery_tasks(app: Celery, provider: DiagnosticProvider) -> None:
     """
     Register all tasks for the given provider
@@ -69,10 +40,14 @@ def register_celery_tasks(app: Celery, provider: DiagnosticProvider) -> None:
     provider
         The provider to register tasks for
     """
-    for metric in provider.diagnostics():
-        print(f"Registering task for diagnostic {metric.name}")
+    for diagnostic in provider.diagnostics():
+        print(f"Registering task for diagnostic {diagnostic.name}")
+
+        # The task function is the same for all diagnostics
+        # The diagnostic is included in the definition
+        # The queue is important to ensure that the task is run in the correct worker
         app.task(  # type: ignore
-            _diagnostic_task_factory(metric),
-            name=generate_task_name(provider, metric),
+            execute_locally,
+            name=generate_task_name(provider, diagnostic),
             queue=provider.slug,
         )

{climate_ref_celery-0.5.0 → climate_ref_celery-0.5.2}/tests/unit/test_cli.py

@@ -37,7 +37,7 @@ def test_start_worker_success(mocker, mock_create_celery_app, mock_register_cele
     mock_import_module.assert_called_once_with("test_package")
     mock_register_celery_tasks.assert_called_once_with(mock_create_celery_app.return_value, mock_provider)
     mock_celery_app.worker_main.assert_called_once_with(
-        argv=["worker", "--loglevel=info", "--queues=example"]
+        argv=["worker", "-E", "--loglevel=info", "--queues=example"]
     )
 
 
@@ -47,7 +47,9 @@ def test_start_core_worker_success(mock_create_celery_app, mock_register_celery_
     result = runner.invoke(app, ["start-worker"])
 
     assert result.exit_code == 0
-    mock_celery_app.worker_main.assert_called_once_with(argv=["worker", "--loglevel=info", "--queues=celery"])
+    mock_celery_app.worker_main.assert_called_once_with(
+        argv=["worker", "-E", "--loglevel=info", "--queues=celery"]
+    )
 
 
 def test_start_worker_success_extra_args(mocker, mock_create_celery_app, mock_register_celery_tasks):
@@ -73,7 +75,7 @@ def test_start_worker_success_extra_args(mocker, mock_create_celery_app, mock_re
 
     assert result.exit_code == 0, result.output
     mock_worker_main.worker_main.assert_called_once_with(
-        argv=["worker", "--loglevel=error", "--queues=example", "--extra-args", "--concurrency=2"]
+        argv=["worker", "-E", "--loglevel=error", "--queues=example", "--extra-args", "--concurrency=2"]
     )
 
 

{climate_ref_celery-0.5.0 → climate_ref_celery-0.5.2}/tests/unit/test_executor.py

@@ -10,7 +10,7 @@ def test_run_metric(provider, config, mock_diagnostic, metric_definition, mocker
     mock_execution_result = mocker.MagicMock()
 
     if include_execution_result:
-        executor.run(provider, mock_diagnostic, metric_definition, mock_execution_result)
+        executor.run(metric_definition, mock_execution_result)
 
         mock_app.send_task.assert_called_once_with(
             "mock_provider.mock",
@@ -19,7 +19,7 @@ def test_run_metric(provider, config, mock_diagnostic, metric_definition, mocker
             queue="mock_provider",
         )
     else:
-        executor.run(provider, mock_diagnostic, metric_definition, None)
+        executor.run(metric_definition, None)
 
         mock_app.send_task.assert_called_once_with(
             "mock_provider.mock",

climate_ref_celery-0.5.2/tests/unit/test_tasks.py

@@ -0,0 +1,23 @@
+from celery import Celery
+from climate_ref_celery.tasks import generate_task_name, register_celery_tasks
+
+
+def test_generate_task_name(mock_diagnostic):
+    assert mock_diagnostic.slug == "mock"
+    assert mock_diagnostic.provider.slug == "mock_provider"
+    assert generate_task_name(mock_diagnostic.provider, mock_diagnostic) == "mock_provider.mock"
+
+
+def test_registry_celery_tasks(provider, mocker):
+    """
+    Test that the tasks are registered correctly
+    """
+    assert len(provider) == 2
+
+    mock_app = mocker.MagicMock(spec=Celery)
+    register_celery_tasks(mock_app, provider)
+
+    assert mock_app.task.call_count == 2
+    assert mock_app.task.call_args_list[0].kwargs["name"] == "mock_provider.mock"
+    assert mock_app.task.call_args_list[1].kwargs["name"] == "mock_provider.failed"
+    assert mock_app.task.call_args_list[0].kwargs["queue"] == "mock_provider"

climate_ref_celery-0.5.0/tests/unit/test_tasks.py

@@ -1,47 +0,0 @@
-from unittest.mock import Mock
-
-from celery import Celery
-from climate_ref_celery.tasks import _diagnostic_task_factory, register_celery_tasks
-
-from climate_ref_core.diagnostics import ExecutionDefinition
-from climate_ref_core.providers import DiagnosticProvider
-
-
-def test_diagnostic_task_factory(tmp_path, caplog):
-    # Mock Diagnostic and ExecutionDefinition
-    mock_metric = Mock()
-
-    definition = ExecutionDefinition(
-        key="test", datasets=None, output_directory=tmp_path / "output", root_directory=tmp_path
-    )
-
-    # Create task using factory
-    task = _diagnostic_task_factory(mock_metric)
-
-    # Run task and check result
-    result = task(definition, "INFO")
-    assert result == mock_metric.run.return_value
-    mock_metric.run.assert_called_once_with(definition)
-
-
-def test_register_celery_tasks(mocker):
-    mock_task_factory = mocker.patch("climate_ref_celery.tasks._diagnostic_task_factory")
-    # Mock Celery app and DiagnosticProvider
-    mock_app = Mock(spec=Celery)
-    mock_provider = Mock(spec=DiagnosticProvider)
-    mock_provider.slug = "test_provider"
-    mock_provider.diagnostics.return_value = [Mock(), Mock()]
-    mock_provider.diagnostics.return_value[0].slug = "metric1"
-    mock_provider.diagnostics.return_value[1].slug = "metric2"
-
-    # Register tasks
-    register_celery_tasks(mock_app, mock_provider)
-
-    # Check that tasks are registered
-    assert mock_app.task.call_count == 2
-    mock_app.task.assert_any_call(
-        mock_task_factory(mock_provider.diagnostics()[0]), name="test_provider.metric1", queue="test_provider"
-    )
-    mock_app.task.assert_any_call(
-        mock_task_factory(mock_provider.diagnostics()[1]), name="test_provider.metric2", queue="test_provider"
-    )