test_training.py

# -*- coding: utf-8 -*-
import logging
import sys

import pytest
import responses

from arkindex.mock import MockApiClient
from arkindex_worker.worker import BaseWorker
from arkindex_worker.worker.training import TrainingMixin, create_archive


@pytest.fixture
def mock_training_worker(monkeypatch):
    class TrainingWorker(BaseWorker, TrainingMixin):
        """
        This class is needed to run tests in the context of a training worker
        """

    monkeypatch.setattr(sys, "argv", ["worker"])
    training_worker = TrainingWorker()
    training_worker.api_client = MockApiClient()
    training_worker.args = training_worker.parser.parse_args()
    return training_worker


@pytest.fixture
def default_model_version():
    return {
        "id": "model_version_id",
        "model_id": "model_id",
        "state": "created",
        "parent": "42" * 16,
        "tag": "A simple tag",
        "description": "A description",
        "configuration": {"test": "value"},
        "s3_url": None,
        "s3_put_url": "http://upload.archive",
        "hash": None,
        "archive_hash": None,
        "size": None,
        "created": "2000-01-01T00:00:00Z",
    }


def test_create_archive(model_file_dir):
    """Create an archive with all base attributes"""

    with create_archive(path=model_file_dir) as (
        zst_archive_path,
        hash,
        size,
        archive_hash,
    ):
        assert zst_archive_path.exists(), "The archive was not created"
        assert (
            hash == "c5aedde18a768757351068b840c8c8f9"
        ), "Hash was not properly computed"
        assert 300 < size < 700

    assert not zst_archive_path.exists(), "Auto removal failed"


def test_create_archive_with_subfolder(model_file_dir_with_subfolder):
    """Create an archive when the model's file is in a folder containing a subfolder"""

    with create_archive(path=model_file_dir_with_subfolder) as (
        zst_archive_path,
        hash,
        size,
        archive_hash,
    ):
        assert zst_archive_path.exists(), "The archive was not created"
        assert (
            hash == "3e453881404689e6e125144d2db3e605"
        ), "Hash was not properly computed"
        assert 300 < size < 1500

    assert not zst_archive_path.exists(), "Auto removal failed"


def test_handle_s3_uploading_errors(mock_training_worker, model_file_dir):
    s3_endpoint_url = "http://s3.localhost.com"
    responses.add_passthru(s3_endpoint_url)
    responses.add(responses.Response(method="PUT", url=s3_endpoint_url, status=400))
    file_path = model_file_dir / "model_file.pth"
    with pytest.raises(Exception):
        mock_training_worker.upload_to_s3(file_path, {"s3_put_url": s3_endpoint_url})


@pytest.mark.parametrize(
    "method",
    [
        ("publish_model_version"),
        ("create_model_version"),
        ("update_model_version"),
        ("upload_to_s3"),
        ("validate_model_version"),
    ],
)
def test_training_mixin_read_only(mock_training_worker, method, caplog):
    """All operations related to models versions returns early if the worker is configured as read only"""
    # Set worker in read_only mode
    mock_training_worker.worker_run_id = None
    assert mock_training_worker.is_read_only

    assert mock_training_worker.model_version is None
    getattr(mock_training_worker, method)()
    assert mock_training_worker.model_version is None
    assert [(level, message) for _, level, message in caplog.record_tuples] == [
        (
            logging.WARNING,
            "Cannot perform this operation as the worker is in read-only mode",
        ),
    ]


def test_create_model_version_already_created(mock_training_worker):
    mock_training_worker.model_version = {"id": "model_version_id"}
    with pytest.raises(
        AssertionError, match="A model version has already been created."
    ):
        mock_training_worker.create_model_version(model_id="model_id")


@pytest.mark.parametrize("set_tag", [True, False])
def test_create_model_version(mock_training_worker, default_model_version, set_tag):
    args = {
        "parent": "42" * 16,
        "tag": "A simple tag",
        "description": "A description",
        "configuration": {"test": "value"},
    }
    if not set_tag:
        del args["tag"]
        default_model_version["tag"] = None
    mock_training_worker.api_client.add_response(
        "CreateModelVersion",
        id="model_id",
        response=default_model_version,
        body=args,
    )
    assert mock_training_worker.model_version is None
    mock_training_worker.create_model_version(model_id="model_id", **args)
    assert mock_training_worker.model_version == default_model_version


def test_update_model_version_not_created(mock_training_worker):
    with pytest.raises(AssertionError, match="No model version has been created yet."):
        mock_training_worker.update_model_version()


def test_update_model_version(mock_training_worker, default_model_version):
    mock_training_worker.model_version = default_model_version
    args = {"tag": "A new tag"}
    new_model_version = {**default_model_version, "tag": "A new tag"}
    mock_training_worker.api_client.add_response(
        "UpdateModelVersion",
        id="model_version_id",
        response=new_model_version,
        body=args,
    )
    mock_training_worker.update_model_version(**args)
    assert mock_training_worker.model_version == new_model_version


def test_validate_model_version_not_created(mock_training_worker):
    with pytest.raises(
        AssertionError,
        match="You must create the model version and upload its archive before validating it.",
    ):
        mock_training_worker.validate_model_version(hash="a", size=1, archive_hash="b")


@pytest.mark.parametrize("deletion_failed", [True, False])
def test_validate_model_version_hash_conflict(
    mock_training_worker, default_model_version, caplog, deletion_failed
):
    mock_training_worker.model_version = {"id": "another_id"}
    args = {
        "hash": "hash",
        "archive_hash": "archive_hash",
        "size": 30,
    }
    mock_training_worker.api_client.add_error_response(
        "ValidateModelVersion",
        id="another_id",
        status_code=409,
        body=args,
        content=default_model_version,
    )
    if deletion_failed:
        mock_training_worker.api_client.add_error_response(
            "DestroyModelVersion",
            id="another_id",
            status_code=403,
            content="Not admin",
        )
    else:
        mock_training_worker.api_client.add_response(
            "DestroyModelVersion",
            id="another_id",
            response="No content",
        )

    mock_training_worker.validate_model_version(**args)
    assert mock_training_worker.model_version == default_model_version
    error_msg = []
    if deletion_failed:
        error_msg = [
            (
                logging.ERROR,
                "An error occurred removing the pending version another_id: Not admin.",
            )
        ]
    assert [
        (level, message)
        for module, level, message in caplog.record_tuples
        if module == "arkindex_worker"
    ] == [
        (
            logging.WARNING,
            "An available model version exists with hash hash, using it instead of the pending version.",
        ),
        (logging.WARNING, "Removing the pending model version."),
        *error_msg,
        (logging.INFO, "Model version model_version_id is now available."),
    ]


def test_validate_model_version(mock_training_worker, default_model_version, caplog):
    mock_training_worker.model_version = {"id": "model_version_id"}
    args = {
        "hash": "hash",
        "archive_hash": "archive_hash",
        "size": 30,
    }
    mock_training_worker.api_client.add_response(
        "ValidateModelVersion",
        id="model_version_id",
        body=args,
        response=default_model_version,
    )

    mock_training_worker.validate_model_version(**args)
    assert mock_training_worker.model_version == default_model_version
    assert [
        (level, message)
        for module, level, message in caplog.record_tuples
        if module == "arkindex_worker"
    ] == [
        (logging.INFO, "Model version model_version_id is now available."),
    ]