Source code for labscheduler.solvers.cp_solver

"""
A solver implementation using google OR-tools to model and solve the JSSP as a constraint program(CP)
"""

import contextlib
import importlib.metadata
import importlib.util
import pickle
import re
import subprocess
import time
from pathlib import Path
from threading import Lock, Thread

from labscheduler.logging_manager import scheduler_logger
from labscheduler.solver_interface import AlgorithmInfo, JSSPSolver
from labscheduler.structures import (
    JSSP,
    Schedule,
    SolutionQuality,
)

if not importlib.util.find_spec("ortools"):
    msg = (
        "The required optional dependency 'ortools' is not installed. Please install it using "
        "'pip install .[cpsolver]' to use the CP solver."
    )
    raise ModuleNotFoundError(msg)

# Check ortools version. If its too low, the solver will not work
ortools_version = importlib.metadata.version("ortools")
ortools_version_parts = []
for part in ortools_version.split(".")[:3]:
    match = re.match(r"\d+", part)
    ortools_version_parts.append(int(match.group()) if match else 0)
ortools_version_tuple = tuple(ortools_version_parts + [0] * (3 - len(ortools_version_parts)))
if ortools_version_tuple < (6, 0, 0):
    msg = f"The CP solver requires ortools>=6.0.0, but ortools=={ortools_version} is installed."
    raise ImportError(msg)




class CPSolver(JSSPSolver):
    def __init__(self):
        # A long-lived worker subprocess pays the (~0.7s) ortools import only once and then serves
        # every solve. It is respawned on demand if it dies (segfault) or hangs past the deadline.
        self._proc: subprocess.Popen | None = None
        self._lock = Lock()  # only one solve at a time may use the pipe
        self._worker_path = Path(__file__).parent / "cp_worker.py"



    def compute_schedule(
        self,
        inst: JSSP,
        time_limit: float,
        offset: float,
        **kwargs,
    ) -> tuple[Schedule | None, SolutionQuality]:
        problem_data = {"inst": inst, "time_limit": time_limit, "offset": offset, **kwargs}

        with self._lock:
            proc = self._ensure_worker()
            worker_start = time.perf_counter()
            result = {}

            # Try to receive the pickled result (blocks until one result frame arrives or the pipe closes)
            def reader():
                with contextlib.suppress(Exception):
                    result.update(pickle.load(proc.stdout))  # noqa: S301

            # Send pickled problem; a broken pipe means the worker just died -> respawn on next call
            try:
                pickle.dump(problem_data, proc.stdin)
                proc.stdin.flush()
            except (BrokenPipeError, OSError):
                self._kill_worker()
                return None, SolutionQuality.INFEASIBLE

            # give the process generous (1 extra second) time to produce a result
            t = Thread(target=reader, daemon=True)
            t.start()
            t.join(time_limit + 1)
            scheduler_logger.info(
                f"[timing] worker (model + solve + result i/o): {time.perf_counter() - worker_start:.3f}s",
            )

            if t.is_alive() or not result:
                # hung past the deadline, or crashed mid-solve (EOF -> empty result):
                # the worker is now in an unknown state, so discard it. Fallbacks cover the caller.
                self._kill_worker()
                return None, SolutionQuality.INFEASIBLE

        return result.get("schedule"), result.get("quality")




    @staticmethod
    def get_algorithm_info() -> AlgorithmInfo:
        return AlgorithmInfo(
            name="CP-Solver",
            is_optimal=True,
            success_guaranty=True,
            max_problem_size=300,
        )




    def is_solvable(self, inst: JSSP) -> bool:
        # model, vars_ = self.create_model(inst, 0)  # noqa: ERA001
        # TODO: there is a method in OR-tools to check solvability
        # state = self.solve_cp(model, vars_, 5)  # noqa: ERA001
        # return state in {OPTIMAL, FEASIBLE}  # noqa: ERA001
        return True




    def _ensure_worker(self) -> subprocess.Popen:
        """Return a live worker, (re)spawning it if it never started or has died."""
        if self._proc is not None and self._proc.poll() is None:
            return self._proc
        t0 = time.perf_counter()
        # stderr is inherited (not captured), so the worker's logs stream straight to our stderr.
        self._proc = subprocess.Popen(  # noqa: S603
            ["python", str(self._worker_path)],  # noqa: S607
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
        )
        scheduler_logger.info(f"[timing] spawned CP worker in {time.perf_counter() - t0:.3f}s")
        return self._proc




    def _kill_worker(self):
        if self._proc is not None:
            with contextlib.suppress(Exception):
                self._proc.kill()
                self._proc.wait()
            self._proc = None