Portfolio_Risk_Management/tests/test_solvers.py

import numpy as np
import pytest
import scipy.sparse as sp

from quantcore.config import Settings
from quantcore.optimization.problem import MILPModel
from quantcore.optimization.solvers import CuOptSolver, HighsSolver, solve_with_fallback, to_cuopt_payload


def readme_model() -> MILPModel:
    # maximise 3x + 2y  s.t. 2x + 3y <= 12, 4x + y <= 10, x,y >= 0 integer  ->  x = 2, y = 2, obj 10
    A = sp.csr_matrix(np.array([[2.0, 3.0], [4.0, 1.0]]))
    return MILPModel(c=np.array([3.0, 2.0]), A=A, row_lb=np.array([-np.inf, -np.inf]), row_ub=np.array([12.0, 10.0]),
                     var_lb=np.zeros(2), var_ub=np.array([np.inf, np.inf]), var_types=["I", "I"], var_names=["x", "y"],
                     row_names=["c1", "c2"], maximize=True)


def test_payload_format():
    payload = to_cuopt_payload(readme_model(), time_limit=5)
    assert payload["csr_constraint_matrix"] == {"offsets": [0, 2, 4], "indices": [0, 1, 0, 1], "values": [2.0, 3.0, 4.0, 1.0]}
    assert payload["constraint_bounds"]["lower_bounds"] == ["ninf", "ninf"]
    assert payload["variable_bounds"]["upper_bounds"] == ["inf", "inf"]
    assert payload["variable_types"] == ["I", "I"] and payload["maximize"] is True
    assert payload["solver_config"]["time_limit"] == 5.0


def test_highs_solves_readme_example():
    res = HighsSolver().solve(readme_model())
    assert res.ok and res.status == "Optimal"
    np.testing.assert_allclose(res.x, [2.0, 2.0])
    assert res.objective == pytest.approx(10.0)


def test_fallback_when_cuopt_unreachable():
    settings = Settings(cuopt_url="http://127.0.0.1:9", solver_fallback=True, cuopt_request_timeout=2)
    res = solve_with_fallback(readme_model(), settings, solver="auto")
    assert res.solver == "highs" and "fallback" in res.message and res.objective == pytest.approx(10.0)
    with pytest.raises(ConnectionError):
        solve_with_fallback(readme_model(), settings, solver="cuopt")
    assert CuOptSolver(settings).health()["reachable"] is False


def test_fallback_when_cuopt_dies_during_the_solve(monkeypatch):
    """cuOpt accepted the job, then its server went away while the client polled for the solution
    (seen with a read-only container whose solver process crashed): 'auto' must still answer, on HiGHS."""
    import requests

    from quantcore.optimization import solvers as S

    settings = Settings(cuopt_url="http://127.0.0.1:9", solver_fallback=True, cuopt_poll_interval=0.01)
    monkeypatch.setattr(S.CuOptSolver, "reachable", lambda self, ttl=30.0: True)
    monkeypatch.setattr(S.CuOptSolver, "_post_with_retry", lambda self, payload: "job-1")

    def lost(*args, **kwargs):
        raise requests.ConnectionError("Max retries exceeded")

    monkeypatch.setattr(S.requests, "get", lost)
    monkeypatch.setattr(S.requests, "delete", lost)
    res = solve_with_fallback(readme_model(), settings, solver="auto")
    assert res.solver == "highs" and "connection lost" in res.message and res.objective == pytest.approx(10.0)
    with pytest.raises(ConnectionError):
        solve_with_fallback(readme_model(), settings, solver="cuopt")


@pytest.mark.skipif(not CuOptSolver(Settings()).health().get("reachable"), reason="cuOpt server not reachable")
def test_cuopt_solves_readme_example():
    res = CuOptSolver(Settings()).solve(readme_model(), time_limit=10)
    assert res.ok and res.solver == "cuopt"
    np.testing.assert_allclose(res.x, [2.0, 2.0])
    assert res.objective == pytest.approx(10.0)


def test_cuopt_retries_down_the_method_ladder_after_an_abort(monkeypatch):
    """cuOpt's default method aborts on some CVaR-objective models; the client works down
    ``CUOPT_RETRY_METHODS`` and reports which method finally solved it."""
    from quantcore.optimization import solvers as S

    calls = []

    class Response:
        def __init__(self, payload, status=200):
            self._payload = payload
            self.status_code = status
            self.text = str(payload)

        def json(self):
            return self._payload

        def raise_for_status(self):
            pass

    def fake_post(url, data=None, headers=None, timeout=None, **kw):  # noqa: ARG001
        import msgpack

        calls.append(msgpack.unpackb(data)["solver_config"])
        return Response({"reqId": f"req{len(calls)}"})

    def fake_get(url, headers=None, timeout=None, **kw):
        if len(calls) == 1:  # first attempt: aborted
            return Response({"error_result": True, "error": "req1 was aborted"})
        return Response({"response": {"solver_response": {
            "status": "Optimal",
            "solution": {"primal_solution": [1.0, 0.0], "primal_objective": 1.0, "solver_time": 0.1,
                         "milp_statistics": {"mip_gap": 0.0}},
        }}})

    class FakeSession:  # the client opens a fresh session per submit attempt
        def __enter__(self):
            return self

        def __exit__(self, *exc):
            return False

        post = staticmethod(fake_post)

    monkeypatch.setattr(S.requests, "Session", FakeSession)
    monkeypatch.setattr(S.requests, "get", fake_get)
    monkeypatch.setattr(S.requests, "delete", lambda *a, **k: Response({}))

    model = readme_model()
    res = S.CuOptSolver(Settings(cuopt_url="http://example.invalid")).solve(model, time_limit=5)
    first = S.CUOPT_RETRY_METHODS[0]
    assert res.status == "Optimal" and f"method={first}" in res.message
    assert calls[0].get("method") is None and calls[1]["method"] == first   # default, then the ladder


def test_cuopt_reports_when_every_method_aborts(monkeypatch):
    """All rungs exhausted: the caller is told that, not just that the first attempt aborted."""
    from quantcore.optimization import solvers as S

    calls = []

    class Response:
        def __init__(self, payload):
            self._payload, self.status_code, self.text = payload, 200, str(payload)

        def json(self):
            return self._payload

        def raise_for_status(self):
            pass

    class FakeSession:
        def __enter__(self):
            return self

        def __exit__(self, *exc):
            return False

        @staticmethod
        def post(url, data=None, headers=None, timeout=None, **kw):
            import msgpack

            calls.append(msgpack.unpackb(data)["solver_config"])
            return Response({"reqId": f"req{len(calls)}"})

    monkeypatch.setattr(S.requests, "Session", FakeSession)
    monkeypatch.setattr(S.requests, "get",
                        lambda *a, **k: Response({"error_result": True, "error": "req was aborted"}))
    monkeypatch.setattr(S.requests, "delete", lambda *a, **k: Response({}))

    res = S.CuOptSolver(Settings(cuopt_url="http://example.invalid")).solve(readme_model(), time_limit=5)
    assert res.status == "Error"
    assert str(S.CUOPT_RETRY_METHODS) in res.message
    assert len(calls) == 1 + len(S.CUOPT_RETRY_METHODS)      # the default plus every rung