climate-ref-celery 0.6.3__py3-none-any.whl → 0.6.4__py3-none-any.whl

climate_ref_celery/cli.py

@@ -3,11 +3,15 @@ Managing remote celery workers
 
 This module is used to manage remote execution workers for the Climate REF project.
 It is added to the `ref` command line interface if the `climate-ref-celery` package is installed.
+
+A celery worker should be run for each diagnostic provider.
 """
 
-import importlib
+import importlib.metadata
+from warnings import warn
 
 import typer
+from loguru import logger
 
 from climate_ref_celery.app import create_celery_app
 from climate_ref_celery.tasks import register_celery_tasks
@@ -16,14 +20,32 @@ from climate_ref_core.providers import DiagnosticProvider
 app = typer.Typer(help=__doc__)
 
 
-def import_provider(provider_package: str) -> DiagnosticProvider:
+def import_provider(provider_name: str) -> DiagnosticProvider:
     """
-    Import the provider from a given package.
+    Import the provider using the name of a registered provider.
 
     Parameters
     ----------
-    provider_package:
-        The package to import the provider from
+    provider_name:
+        The name of a registered provider.
+
+        Packages can register a provider by defining an
+        [entry point](https://packaging.python.org/en/latest/specifications/entry-points/)
+         in its `pyproject.toml` file under the group `"climate-ref.providers"`.
+
+        Example: 'climate_ref_esmvaltool:provider' would require a section in the `pyproject.toml` for the
+        `climate_ref_esmvaltool` package like this:
+
+        ```
+        [project.entry-points."climate-ref.providers"]
+        esmvaltool = "climate_ref_esmvaltool:provider"
+        ```
+
+        `"esmvaltool"` or ("climate_ref_esmvaltool:provider")
+        can then be used as the `provider_name` argument.
+
+        If the entry point is not found, an error will be raised
+        and the list of available providers will be shown.
 
     Raises
     ------
@@ -37,18 +59,37 @@ def import_provider(provider_package: str) -> DiagnosticProvider:
     :
         The provider instance
     """
-    try:
-        imp = importlib.import_module(provider_package.replace("-", "_"))
-    except ModuleNotFoundError:
-        typer.echo(f"Package '{provider_package}' not found")
+    provider_entry_points = importlib.metadata.entry_points(group="climate-ref.providers")
+    for entry_point in provider_entry_points:
+        logger.debug(f"found entry point: {entry_point}")
+
+        # Also support the case where the entrypoint definition ('name:provider') is supplied
+        if entry_point.name == provider_name or entry_point.value == provider_name:  # noqa: PLR1714
+            break
+    else:
+        found_entry_points = ", ".join(f"{ep.name} ({ep.value})" for ep in provider_entry_points)
+        if len(found_entry_points) == 0:
+            found_entry_points = "[]"
+
+        typer.echo(
+            f"No entry point named {provider_name!r} was found. Found entry points: {found_entry_points}."
+        )
         raise typer.Abort()
 
     # Get the provider from the provider_package
     try:
-        provider = imp.provider
+        provider = entry_point.load()
+    except ModuleNotFoundError:
+        _split = entry_point.value.split(":", 1)
+        typer.echo(f"Invalid entrypoint {entry_point}: Package {_split[0]!r} not found.")
+        raise typer.Abort()
     except AttributeError:
-        typer.echo("The package must define a 'provider' attribute")
+        _split = entry_point.value.split(":", 1)
+        typer.echo(
+            f"Invalid entrypoint {entry_point}: {_split[0]!r} does not define a {_split[1]!r} attribute."
+        )
         raise typer.Abort()
+
     if not isinstance(provider, DiagnosticProvider):
         typer.echo(f"Expected DiagnosticProvider, got {type(provider)}")
         raise typer.Abort()
@@ -59,40 +100,61 @@ def import_provider(provider_package: str) -> DiagnosticProvider:
 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),
+    provider: list[str] | None = typer.Option(
+        help="Name of the provider to start a worker for. This argument may be supplied multiple times. "
+        "If no provider is given, the worker will consume the default queue.",
+        default=None,
+    ),
+    package: str | None = typer.Option(help="Deprecated. Use provider instead", default=None),
     extra_args: list[str] = typer.Argument(None, help="Additional arguments for the worker"),
 ) -> None:
     """
-    Start a Celery worker for the given package.
+    Start a Celery worker for the given provider.
 
     A celery worker enables the execution of tasks in the background on multiple different nodes.
     This worker will register a celery task for each diagnostic in the provider.
     The worker tasks can be executed by sending a celery task with the name
     '{package_slug}_{diagnostic_slug}'.
 
-    The package must define a 'provider' variable that is an instance of 'ref_core.DiagnosticProvider'.
+    Providers must be registered as entry points in the `pyproject.toml` file of the package.
+    The entry point should be defined under the group `climate-ref.providers`
+    (See `import_provider` for details).
     """
     # Create a new celery app
     celery_app = create_celery_app("climate_ref_celery")
 
     if package:
-        # 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
+        msg = "The '--package' argument is deprecated. Use '--provider' instead."
+        # Deprecation warning for package argument
+        warn(
+            msg,
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        typer.echo(msg)
+        # Assume the package is the provider
+        provider = [package + ":provider"]
+
+    queues = []
+    if provider:
+        for p in provider:
+            # Attempt to import the provider
+            provider_instance = import_provider(p)
+
+            if hasattr(ctx.obj, "config"):
+                # Configure the provider so that it knows where the conda environments are
+                provider_instance.configure(ctx.obj.config)
+
+            # Wrap each diagnostics in the provider with a celery tasks
+            register_celery_tasks(celery_app, provider_instance)
+            queues.append(provider_instance.slug)
     else:
         # This might need some tweaking in later PRs to pull in the appropriate tasks
         import climate_ref_celery.worker_tasks  # noqa: F401
 
-        queue = "celery"
+        queues.append("celery")
 
-    argv = ["worker", "-E", f"--loglevel={loglevel}", f"--queues={queue}", *(extra_args or [])]
+    argv = ["worker", "-E", f"--loglevel={loglevel}", f"--queues={','.join(queues)}", *(extra_args or [])]
     celery_app.worker_main(argv=argv)
 
 

{climate_ref_celery-0.6.3.dist-info → climate_ref_celery-0.6.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: climate-ref-celery
-Version: 0.6.3
+Version: 0.6.4
 Summary: Celery app for mananging tasks and workers
 Author-email: Jared Lewis <jared.lewis@climate-resource.com>
 License-Expression: Apache-2.0

{climate_ref_celery-0.6.3.dist-info → climate_ref_celery-0.6.4.dist-info}/RECORD

@@ -1,6 +1,6 @@
 climate_ref_celery/__init__.py,sha256=f-dS4JxMLmTCobZtftGDi6bUbtYCv_9iWqxSpC74tWw,129
 climate_ref_celery/app.py,sha256=KtDX8-6cye9RzJL5rXbJH1Us_pGYONawYAkOnxTcBmI,1527
-climate_ref_celery/cli.py,sha256=7eQ-ubng-dGQoZqSmLAUeM04Ucr2lOm772hzAlAeiD8,3376
+climate_ref_celery/cli.py,sha256=7cQjB836xIuQ8flwnh8_gJmoVJUcHhqwcaH5d1N9Yg4,6022
 climate_ref_celery/executor.py,sha256=Qsx1bvM_-H3eMeg5VsJAV6ghHuGK3WnYyIRAd0-1DG4,4764
 climate_ref_celery/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 climate_ref_celery/tasks.py,sha256=6Ra2H48edbcklsf6pAj1JNfzi67n6WQCUXJs-tloE80,1941
@@ -9,9 +9,9 @@ climate_ref_celery/celeryconf/__init__.py,sha256=xKcHhnECGZAxrjnMtfaPBU_vgqwRbRL
 climate_ref_celery/celeryconf/base.py,sha256=uXaz9S_1hCdwKpvyWBDWE69YVIXHD3Ar7egNKLk9Gx8,484
 climate_ref_celery/celeryconf/dev.py,sha256=um7hvEhlTBfwey9ler_vAHAln_DrpqAP3fbbO2wdmUQ,217
 climate_ref_celery/celeryconf/prod.py,sha256=HyGTfwTZD8MjIkfkb3qyfzjjTPE9sKAoXNq3fJDFO9Q,231
-climate_ref_celery-0.6.3.dist-info/METADATA,sha256=BdXzmcTQzA-VNHUabapu7rLnw5Z-GIUxRqO2NZPTgIk,3693
-climate_ref_celery-0.6.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-climate_ref_celery-0.6.3.dist-info/entry_points.txt,sha256=U9b-T6EpLV3ZXmHUpEzp8x5TZnCjQ1ynncIkFMwDuPE,58
-climate_ref_celery-0.6.3.dist-info/licenses/LICENCE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-climate_ref_celery-0.6.3.dist-info/licenses/NOTICE,sha256=4qTlax9aX2-mswYJuVrLqJ9jK1IkN5kSBqfVvYLF3Ws,128
-climate_ref_celery-0.6.3.dist-info/RECORD,,
+climate_ref_celery-0.6.4.dist-info/METADATA,sha256=k-EzKWFU77223JpuGYjHWRjpR6hJTGBBmyYB38eBjzs,3693
+climate_ref_celery-0.6.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+climate_ref_celery-0.6.4.dist-info/entry_points.txt,sha256=U9b-T6EpLV3ZXmHUpEzp8x5TZnCjQ1ynncIkFMwDuPE,58
+climate_ref_celery-0.6.4.dist-info/licenses/LICENCE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+climate_ref_celery-0.6.4.dist-info/licenses/NOTICE,sha256=4qTlax9aX2-mswYJuVrLqJ9jK1IkN5kSBqfVvYLF3Ws,128
+climate_ref_celery-0.6.4.dist-info/RECORD,,