climate-ref-celery 0.6.3__tar.gz → 0.6.5__tar.gz

{climate_ref_celery-0.6.3 → climate_ref_celery-0.6.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: climate-ref-celery
-Version: 0.6.3
+Version: 0.6.5
 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 → climate_ref_celery-0.6.5}/pyproject.toml

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

climate_ref_celery-0.6.5/src/climate_ref_celery/cli.py

@@ -0,0 +1,172 @@
+"""
+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.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
+from climate_ref_core.providers import DiagnosticProvider
+
+app = typer.Typer(help=__doc__)
+
+
+def import_provider(provider_name: str) -> DiagnosticProvider:
+    """
+    Import the provider using the name of a registered provider.
+
+    Parameters
+    ----------
+    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
+    ------
+    typer.Abort
+        If the provider_package does not define a 'provider' variable
+
+        If the provider_package is not found
+
+    Returns
+    -------
+    :
+        The provider instance
+    """
+    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 = 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:
+        _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()
+    return provider
+
+
+@app.command()
+def start_worker(
+    ctx: typer.Context,
+    loglevel: str = typer.Option("info", help="Log level for the worker"),
+    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 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}'.
+
+    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:
+        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
+
+        queues.append("celery")
+
+    argv = ["worker", "-E", f"--loglevel={loglevel}", f"--queues={','.join(queues)}", *(extra_args or [])]
+    celery_app.worker_main(argv=argv)
+
+
+@app.command()
+def list_config() -> None:
+    """
+    List the celery configuration
+    """
+    celery_app = create_celery_app("climate_ref_celery")
+
+    print(celery_app.conf.humanize())
+
+
+if __name__ == "__main__":  # pragma: no cover
+    app()

climate_ref_celery-0.6.5/tests/unit/test_cli.py

@@ -0,0 +1,196 @@
+import importlib.metadata
+
+import pytest
+from climate_ref_celery.cli import app
+from typer.testing import CliRunner
+
+from climate_ref_core.providers import DiagnosticProvider
+
+runner = CliRunner()
+
+
+def test_cli_help():
+    result = runner.invoke(app, ["--help"])
+    assert result.exit_code == 0
+
+
+@pytest.fixture
+def mock_create_celery_app(mocker):
+    return mocker.patch("climate_ref_celery.cli.create_celery_app")
+
+
+@pytest.fixture
+def mock_register_celery_tasks(mocker):
+    return mocker.patch("climate_ref_celery.cli.register_celery_tasks")
+
+
+@pytest.mark.parametrize("provider", ["test_package", "test_package:provider"])
+def test_start_worker_success(mocker, mock_create_celery_app, mock_register_celery_tasks, provider):
+    mock_celery_app = mock_create_celery_app.return_value
+    mock_provider = mocker.MagicMock(spec=DiagnosticProvider)
+    mock_provider.slug = "example"
+
+    mock_entry_point = mocker.Mock(spec=importlib.metadata.EntryPoint)
+    mock_entry_point.name = "test_package"
+    mock_entry_point.value = "test_package:provider"
+    mock_entry_point.load.return_value = mock_provider
+    mock_entry_points = mocker.patch("importlib.metadata.entry_points", return_value=[mock_entry_point])
+
+    result = runner.invoke(app, ["start-worker", "--provider", provider])
+
+    assert result.exit_code == 0
+    mock_entry_points.assert_called_once_with(group="climate-ref.providers")
+    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", "-E", "--loglevel=info", "--queues=example"]
+    )
+
+
+def test_start_worker_multiple(mocker, mock_create_celery_app, mock_register_celery_tasks):
+    mock_celery_app = mock_create_celery_app.return_value
+
+    mock_provider_a = mocker.MagicMock(spec=DiagnosticProvider)
+    mock_provider_a.slug = "example"
+    mock_entry_point_a = mocker.Mock(spec=importlib.metadata.EntryPoint)
+    mock_entry_point_a.name = "test_package"
+    mock_entry_point_a.value = "test_package:provider"
+    mock_entry_point_a.load.return_value = mock_provider_a
+
+    mock_provider_b = mocker.MagicMock(spec=DiagnosticProvider)
+    mock_provider_b.slug = "other"
+    mock_entry_point_b = mocker.Mock(spec=importlib.metadata.EntryPoint)
+    mock_entry_point_b.name = "other_package"
+    mock_entry_point_b.value = "other_package:provider"
+    mock_entry_point_b.load.return_value = mock_provider_b
+
+    mock_entry_points = mocker.patch(
+        "importlib.metadata.entry_points", return_value=[mock_entry_point_a, mock_entry_point_b]
+    )
+
+    result = runner.invoke(app, ["start-worker", "--provider", "test_package", "--provider", "other_package"])
+
+    assert result.exit_code == 0
+    mock_entry_points.assert_called_with(group="climate-ref.providers")
+    mock_register_celery_tasks.assert_any_call(mock_create_celery_app.return_value, mock_provider_a)
+    mock_register_celery_tasks.assert_any_call(mock_create_celery_app.return_value, mock_provider_b)
+    mock_celery_app.worker_main.assert_called_once_with(
+        argv=["worker", "-E", "--loglevel=info", "--queues=example,other"]
+    )
+
+
+def test_start_core_worker_success(mock_create_celery_app, mock_register_celery_tasks):
+    mock_celery_app = mock_create_celery_app.return_value
+
+    result = runner.invoke(app, ["start-worker"])
+
+    assert result.exit_code == 0
+    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):
+    mock_worker_main = mock_create_celery_app.return_value
+    mock_provider = mocker.MagicMock(spec=DiagnosticProvider)
+    mock_provider.slug = "example"
+
+    mock_entry_point = mocker.Mock(spec=importlib.metadata.EntryPoint)
+    mock_entry_point.name = "test_package"
+    mock_entry_point.load.return_value = mock_provider
+    mocker.patch("importlib.metadata.entry_points", return_value=[mock_entry_point])
+
+    result = runner.invoke(
+        app,
+        [
+            "start-worker",
+            "--loglevel",
+            "error",
+            "--provider",
+            "test_package",
+            "--",
+            "--extra-args",
+            "--concurrency=2",
+        ],
+    )
+
+    assert result.exit_code == 0, result.output
+    mock_worker_main.worker_main.assert_called_once_with(
+        argv=["worker", "-E", "--loglevel=error", "--queues=example", "--extra-args", "--concurrency=2"]
+    )
+
+
+def test_start_worker_package_not_registered(mocker, mock_create_celery_app):
+    mocker.patch("importlib.metadata.entry_points", return_value=[])
+
+    result = runner.invoke(app, ["start-worker", "--provider", "unregistered_package"])
+
+    assert result.exit_code == 1
+    assert "No entry point named 'unregistered_package' was found" in result.output
+    assert "Found entry points: []" in result.output
+    mock_create_celery_app.assert_called_once_with("climate_ref_celery")
+
+
+def test_start_worker_package_not_found(mocker, mock_create_celery_app):
+    mock_entry_point = mocker.Mock(spec=importlib.metadata.EntryPoint)
+    mock_entry_point.name = "missing_package"
+    mock_entry_point.value = "missing_package:provider"
+    mock_entry_point.load.side_effect = ModuleNotFoundError
+    mock_entry_points = mocker.patch("importlib.metadata.entry_points", return_value=[mock_entry_point])
+
+    result = runner.invoke(app, ["start-worker", "--provider", "missing_package"])
+
+    assert result.exit_code == 1
+    assert "Package 'missing_package' not found" in result.output
+    mock_create_celery_app.assert_called_once_with("climate_ref_celery")
+    mock_entry_points.assert_called_once_with(group="climate-ref.providers")
+
+
+def test_start_worker_missing_provider(mocker, mock_create_celery_app):
+    mock_entry_point = mocker.Mock(spec=importlib.metadata.EntryPoint)
+    mock_entry_point.name = "test_package"
+    mock_entry_point.value = "test_package:provider"
+    mock_entry_point.load.side_effect = AttributeError
+    mocker.patch("importlib.metadata.entry_points", return_value=[mock_entry_point])
+
+    result = runner.invoke(app, ["start-worker", "--provider", "test_package"])
+
+    assert result.exit_code == 1, result.output
+    assert "'test_package' does not define a 'provider' attribute" in result.output
+
+
+def test_start_worker_incorrect_provider(mocker, mock_create_celery_app):
+    # Not a DiagnosticProvider
+    mock_provider = mocker.Mock()
+
+    mock_entry_point = mocker.Mock(spec=importlib.metadata.EntryPoint)
+    mock_entry_point.name = "test_package"
+    mock_entry_point.load.return_value = mock_provider
+    mocker.patch("importlib.metadata.entry_points", return_value=[mock_entry_point])
+
+    result = runner.invoke(app, ["start-worker", "--provider", "test_package"])
+
+    assert result.exit_code == 1, result.output
+    assert "Expected DiagnosticProvider, got <class 'unittest.mock.Mock'>" in result.output
+
+
+def test_start_worker_deprecated_package(mocker, mock_create_celery_app):
+    mock_provider = mocker.MagicMock(spec=DiagnosticProvider)
+    mock_provider.slug = "example"
+
+    mock_entry_point = mocker.Mock(spec=importlib.metadata.EntryPoint)
+    mock_entry_point.name = "test_package"
+    mock_entry_point.value = "test_package:provider"
+    mock_entry_point.load.return_value = mock_provider
+    mocker.patch("importlib.metadata.entry_points", return_value=[mock_entry_point])
+
+    result = runner.invoke(app, ["start-worker", "--package", "test_package"])
+
+    assert result.exit_code == 0, result.output
+    assert "The '--package' argument is deprecated. Use '--provider' instead." in result.output
+
+
+def test_list_config():
+    result = runner.invoke(app, ["list-config"])
+
+    assert result.exit_code == 0, result.output
+    assert "broker_url: 'redis://localhost:6379/1'" in result.stdout

climate_ref_celery-0.6.3/src/climate_ref_celery/cli.py

@@ -1,110 +0,0 @@
-"""
-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.
-"""
-
-import importlib
-
-import typer
-
-from climate_ref_celery.app import create_celery_app
-from climate_ref_celery.tasks import register_celery_tasks
-from climate_ref_core.providers import DiagnosticProvider
-
-app = typer.Typer(help=__doc__)
-
-
-def import_provider(provider_package: str) -> DiagnosticProvider:
-    """
-    Import the provider from a given package.
-
-    Parameters
-    ----------
-    provider_package:
-        The package to import the provider from
-
-    Raises
-    ------
-    typer.Abort
-        If the provider_package does not define a 'provider' variable
-
-        If the provider_package is not found
-
-    Returns
-    -------
-    :
-        The provider instance
-    """
-    try:
-        imp = importlib.import_module(provider_package.replace("-", "_"))
-    except ModuleNotFoundError:
-        typer.echo(f"Package '{provider_package}' not found")
-        raise typer.Abort()
-
-    # Get the provider from the provider_package
-    try:
-        provider = imp.provider
-    except AttributeError:
-        typer.echo("The package must define a 'provider' attribute")
-        raise typer.Abort()
-    if not isinstance(provider, DiagnosticProvider):
-        typer.echo(f"Expected DiagnosticProvider, got {type(provider)}")
-        raise typer.Abort()
-    return provider
-
-
-@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"),
-) -> None:
-    """
-    Start a Celery worker for the given package.
-
-    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'.
-    """
-    # 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
-    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"
-
-    argv = ["worker", "-E", f"--loglevel={loglevel}", f"--queues={queue}", *(extra_args or [])]
-    celery_app.worker_main(argv=argv)
-
-
-@app.command()
-def list_config() -> None:
-    """
-    List the celery configuration
-    """
-    celery_app = create_celery_app("climate_ref_celery")
-
-    print(celery_app.conf.humanize())
-
-
-if __name__ == "__main__":  # pragma: no cover
-    app()

climate_ref_celery-0.6.3/tests/unit/test_cli.py

@@ -1,124 +0,0 @@
-import pytest
-from climate_ref_celery.cli import app
-from typer.testing import CliRunner
-
-from climate_ref_core.providers import DiagnosticProvider
-
-runner = CliRunner()
-
-
-def test_cli_help():
-    result = runner.invoke(app, ["--help"])
-    assert result.exit_code == 0
-
-
-@pytest.fixture
-def mock_create_celery_app(mocker):
-    return mocker.patch("climate_ref_celery.cli.create_celery_app")
-
-
-@pytest.fixture
-def mock_register_celery_tasks(mocker):
-    return mocker.patch("climate_ref_celery.cli.register_celery_tasks")
-
-
-def test_start_worker_success(mocker, mock_create_celery_app, mock_register_celery_tasks):
-    mock_celery_app = mock_create_celery_app.return_value
-    mock_provider = mocker.MagicMock(spec=DiagnosticProvider)
-    mock_provider.slug = "example"
-
-    mock_import_module = mocker.patch(
-        "importlib.import_module", return_value=mocker.Mock(provider=mock_provider)
-    )
-
-    result = runner.invoke(app, ["start-worker", "--package", "test_package"])
-
-    assert result.exit_code == 0
-    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", "-E", "--loglevel=info", "--queues=example"]
-    )
-
-
-def test_start_core_worker_success(mock_create_celery_app, mock_register_celery_tasks):
-    mock_celery_app = mock_create_celery_app.return_value
-
-    result = runner.invoke(app, ["start-worker"])
-
-    assert result.exit_code == 0
-    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):
-    mock_worker_main = mock_create_celery_app.return_value
-    mock_provider = mocker.MagicMock(spec=DiagnosticProvider)
-    mock_provider.slug = "example"
-
-    mocker.patch("importlib.import_module", return_value=mocker.Mock(provider=mock_provider))
-
-    result = runner.invoke(
-        app,
-        [
-            "start-worker",
-            "--loglevel",
-            "error",
-            "--package",
-            "test_package",
-            "--",
-            "--extra-args",
-            "--concurrency=2",
-        ],
-    )
-
-    assert result.exit_code == 0, result.output
-    mock_worker_main.worker_main.assert_called_once_with(
-        argv=["worker", "-E", "--loglevel=error", "--queues=example", "--extra-args", "--concurrency=2"]
-    )
-
-
-def test_start_worker_package_not_found(mocker, mock_create_celery_app):
-    mock_import_module = mocker.patch("importlib.import_module", side_effect=ModuleNotFoundError)
-
-    result = runner.invoke(app, ["start-worker", "--package", "missing_package"])
-
-    assert result.exit_code == 1
-    assert "Package 'missing_package' not found" in result.output
-    mock_create_celery_app.assert_called_once_with("climate_ref_celery")
-    mock_import_module.assert_called_once_with("missing_package")
-
-
-def test_start_worker_missing_provider(mocker, mock_create_celery_app):
-    mock_module = mocker.Mock()
-    del mock_module.provider
-    mock_import_module = mocker.patch("importlib.import_module", return_value=mock_module)
-
-    result = runner.invoke(app, ["start-worker", "--package", "test_package"])
-
-    assert result.exit_code == 1, result.output
-    assert "The package must define a 'provider' attribute" in result.output
-    mock_import_module.assert_called_once_with("test_package")
-
-
-def test_start_worker_incorrect_provider(mocker, mock_create_celery_app):
-    # Not a DiagnosticProvider
-    mock_provider = mocker.Mock()
-
-    mock_import_module = mocker.patch(
-        "importlib.import_module", return_value=mocker.Mock(provider=mock_provider)
-    )
-
-    result = runner.invoke(app, ["start-worker", "--package", "test_package"])
-
-    assert result.exit_code == 1, result.output
-    assert "Expected DiagnosticProvider, got <class 'unittest.mock.Mock'>" in result.output
-    mock_import_module.assert_called_once_with("test_package")
-
-
-def test_list_config():
-    result = runner.invoke(app, ["list-config"])
-
-    assert result.exit_code == 0, result.output
-    assert "broker_url: 'redis://localhost:6379/1'" in result.stdout