diff --git a/mixtape/core/tasks/inference_tasks.py b/mixtape/core/tasks/inference_tasks.py index a069339..ef178a8 100644 --- a/mixtape/core/tasks/inference_tasks.py +++ b/mixtape/core/tasks/inference_tasks.py @@ -15,8 +15,7 @@ from pettingzoo.utils import ParallelEnv from ray.rllib.algorithms.algorithm import Algorithm -from mixtape.core.models import AgentStep, Episode, Inference -from mixtape.core.models.step import Step +from mixtape.core.models import AgentStep, Episode, Inference, Step from mixtape.core.ray_utils.environments import register_environment if TYPE_CHECKING: diff --git a/mixtape/core/tests/conftest.py b/mixtape/core/tests/conftest.py index 0ed32f4..8b6ca8a 100644 --- a/mixtape/core/tests/conftest.py +++ b/mixtape/core/tests/conftest.py @@ -1,3 +1,4 @@ +from click.testing import CliRunner from playwright.sync_api import BrowserContext import pytest from pytest_django.live_server_helper import LiveServer @@ -10,6 +11,12 @@ def api_client() -> APIClient: return APIClient() +@pytest.fixture +def cli_runner() -> CliRunner: + # Don't catch exceptions, so they'll be raised in the test case + return CliRunner(catch_exceptions=False) + + # This intentionally overrides the built-in fixture from pytest_playwright. # This will also cause other built-in fixtures like "page" to have a base URL set. @pytest.fixture diff --git a/mixtape/core/tests/data/checkpoint_archive.tar.bz2 b/mixtape/core/tests/data/checkpoint_archive.tar.bz2 new file mode 100644 index 0000000..b75d691 Binary files /dev/null and b/mixtape/core/tests/data/checkpoint_archive.tar.bz2 differ diff --git a/mixtape/core/tests/factories.py b/mixtape/core/tests/factories.py index b87ffc7..d6ca73d 100644 --- a/mixtape/core/tests/factories.py +++ b/mixtape/core/tests/factories.py @@ -1,5 +1,6 @@ from collections import OrderedDict from io import BytesIO +from pathlib import Path import PIL.Image import PIL.ImageDraw @@ -37,6 +38,11 @@ class Meta: training = factory.SubFactory(TrainingFactory) last = True + archive = factory.django.FileField( + from_path=Path(__file__).parent / 'data' / 'checkpoint_archive.tar.bz2', + # TODO: Change to f'checkpoint/{uuid4()}.tar.bz2' + filename=f'checkpoint/archive.tar.bz2' + ) class InferenceFactory(factory.django.DjangoModelFactory[Inference]): diff --git a/mixtape/core/tests/test_cli.py b/mixtape/core/tests/test_cli.py new file mode 100644 index 0000000..4da77ad --- /dev/null +++ b/mixtape/core/tests/test_cli.py @@ -0,0 +1,59 @@ +from click.testing import CliRunner +import pytest + +from mixtape.core.management.commands.inference import inference as inference_command +from mixtape.core.management.commands.training import training as training_command +from mixtape.core.models import Checkpoint, Episode, Training + +from .factories import CheckpointFactory + + +@pytest.mark.parametrize( + ('env_name', 'parallel'), + [ + ('knights_archers_zombies_v10', True), + # ('knights_archers_zombies_v10', False), + # ('BattleZone-v5', False), + ], + ids=[ + 'pettingzoo_parallel', + # 'pettingzoo_aec', + # 'gymnasium', + ], +) +@pytest.mark.django_db +def test_cli_training(cli_runner: CliRunner, env_name: str, parallel: bool): + training_result = cli_runner.invoke( + training_command, + [ + '--env_name', + env_name, + '--algorithm', + 'PPO', + *(['--parallel'] if parallel else []), + '--training_iteration', + '2', + '--immediate', + ], + ) + assert training_result.exit_code == 0 + training = Training.objects.get() + assert training.environment == env_name + assert training.algorithm == 'PPO' + assert training.parallel is parallel + assert training.iterations == 2 + + +@pytest.mark.django_db +def test_cli_inference(cli_runner: CliRunner): + checkpoint = CheckpointFactory.create() + + inference_result = cli_runner.invoke( + inference_command, + [ + str(checkpoint.id), + '--immediate', + ], + ) + assert inference_result.exit_code == 0 + assert Episode.objects.filter(inference__checkpoint=checkpoint).exists() diff --git a/mixtape/settings/testing.py b/mixtape/settings/testing.py index e7337d3..12ee44d 100644 --- a/mixtape/settings/testing.py +++ b/mixtape/settings/testing.py @@ -16,3 +16,5 @@ MINIO_STORAGE_MEDIA_BUCKET_NAME = f'test-django-storage-{randbelow(1_000_000):06d}' # Testing will set EMAIL_BACKEND to use the memory backend + +CELERY_TASK_EAGER_PROPAGATES = True diff --git a/pyproject.toml b/pyproject.toml index fdd846c..0010e52 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -168,4 +168,12 @@ filterwarnings = [ "ignore:Tensorflow not installed; ParametricUMAP will be unavailable:ImportWarning:umap", # Fixed by https://github.com/ray-project/ray/pull/49004 in Ray 2.43 "ignore:invalid escape sequence:DeprecationWarning", + "ignore:pkg_resources is deprecated:UserWarning:pygame.pkgdata", + "ignore:Deprecated call to `pkg_resources.declare_namespace:DeprecationWarning:pkg_resources", + "ignore:checkpoint_freq is deprecated:DeprecationWarning:ray.tune.tune", + "ignore:checkpoint_at_end is deprecated:DeprecationWarning:ray.tune.tune", + "ignore:.*Box bound precision lowered by casting:UserWarning:gymnasium.spaces.box", + "ignore:.*The obs returned by the `reset:UserWarning:gymnasium.utils.passive_env_checker", + "ignore:unclosed file:ResourceWarning:ray._private.node", + "ignore::ray.util.annotations.RayDeprecationWarning", ]