From 5696e607f8b42dce50fa450e58d67c88cfee548a Mon Sep 17 00:00:00 2001 From: Trevor McKay Date: Fri, 4 Sep 2026 14:08:41 -0400 Subject: [PATCH] change Python grpc client so wait loop happens in Python If we use the C++ wait_result() method, the GIL is locked and no other threads run. --- docs/cuopt/source/cuopt-grpc/routing.rst | 1 + .../cuopt/cuopt/grpc/client/grpc_client.pyx | 65 +++++- .../linear_programming/test_grpc_client.py | 198 +++++++++++++++++- 3 files changed, 252 insertions(+), 12 deletions(-) diff --git a/docs/cuopt/source/cuopt-grpc/routing.rst b/docs/cuopt/source/cuopt-grpc/routing.rst index 2827301058..e911054bf9 100644 --- a/docs/cuopt/source/cuopt-grpc/routing.rst +++ b/docs/cuopt/source/cuopt-grpc/routing.rst @@ -134,6 +134,7 @@ Import path: ``cuopt.grpc.routing``. .. autoclass:: cuopt.grpc.routing.RoutingClient :members: :undoc-members: + :exclude-members: _status .. autoexception:: cuopt.grpc.routing.RoutingSolveError :members: diff --git a/python/cuopt/cuopt/grpc/client/grpc_client.pyx b/python/cuopt/cuopt/grpc/client/grpc_client.pyx index 3f9cdcb0e5..d1ced65eec 100644 --- a/python/cuopt/cuopt/grpc/client/grpc_client.pyx +++ b/python/cuopt/cuopt/grpc/client/grpc_client.pyx @@ -93,6 +93,38 @@ class JobNotReadyError(GrpcError): pass +# Matches the previous C++ wait() poll cadence. Keep this in Python +# (time.sleep) so the GIL is released between short status RPCs. +_WAIT_POLL_INTERVAL_S = 1.0 + + +def _wait_poll_loop( + get_status, + job_id, + timeout_seconds, + error_cls, + poll_interval_s=_WAIT_POLL_INTERVAL_S, +): + """Poll ``get_status(job_id)`` until the job is terminal or the timeout. + + The loop and ``time.sleep`` run in Python so the GIL is released between + short status RPCs. Concurrent incumbent/log stream threads can therefore + run during ``Client.wait``. ``timeout_seconds <= 0`` waits indefinitely. + """ + if timeout_seconds < 0: + raise error_cls("timeout_seconds must be non-negative") + deadline = ( + time.monotonic() + timeout_seconds if timeout_seconds > 0 else None + ) + while True: + status = get_status(job_id) + if status not in (JobStatus.QUEUED, JobStatus.PROCESSING): + return status + if deadline is not None and time.monotonic() >= deadline: + raise error_cls("Timeout waiting for job completion") + time.sleep(poll_interval_s) + + cdef int _invoke_log_callback( const char* line, size_t line_len, @@ -320,14 +352,17 @@ cdef class Client: becomes ``0`` and waits indefinitely). Positive timeouts poll about once per second and raise :class:`GrpcError` if the deadline expires (they do not return a non-terminal :class:`JobStatus`). + + The wait loop runs in Python and only calls :meth:`status` for each + poll, so the GIL is released between checks. Concurrent + :meth:`start_incumbent_stream` and + :meth:`start_log_stream` threads can therefore make progress during + the wait. """ - cdef int timeout_seconds = 0 if timeout is None else int(timeout) - cdef grpc_status_result_t wait_result = self._client.get().wait( - job_id.encode("utf-8"), timeout_seconds + timeout_seconds = 0 if timeout is None else int(timeout) + return _wait_poll_loop( + self.status, job_id, timeout_seconds, GrpcError ) - if not wait_result.success: - raise GrpcError(wait_result.error_message.decode("utf-8")) - return JobStatus(wait_result.status) def cancel(self, str job_id): """ @@ -1095,18 +1130,26 @@ cdef class RoutingClient: raise RoutingSolveError(sub.error_message.decode("utf-8")) return sub.job_id.decode("utf-8") + def _status(self, str job_id): + cdef grpc_status_result_t st = self._client.get().status( + job_id.encode("utf-8") + ) + if not st.success: + raise RoutingSolveError(st.error_message.decode("utf-8")) + return JobStatus(st.status) + def wait(self, str job_id, int timeout=0): """Block until the job finishes; return the terminal status int. Raises ``RoutingSolveError`` if the wait itself fails (e.g. transport error or unknown job), mirroring the LP/MILP client. + + Polls job status from Python so the GIL is released between short + status RPCs. ``timeout <= 0`` waits indefinitely. """ - cdef grpc_status_result_t st = self._client.get().wait( - job_id.encode("utf-8"), timeout + return _wait_poll_loop( + self._status, job_id, timeout, RoutingSolveError ) - if not st.success: - raise RoutingSolveError(st.error_message.decode("utf-8")) - return st.status def result(self, str job_id): """Fetch and parse the routing solution for a completed job. diff --git a/python/cuopt/cuopt/tests/linear_programming/test_grpc_client.py b/python/cuopt/cuopt/tests/linear_programming/test_grpc_client.py index 7807c757ff..1365226d01 100644 --- a/python/cuopt/cuopt/tests/linear_programming/test_grpc_client.py +++ b/python/cuopt/cuopt/tests/linear_programming/test_grpc_client.py @@ -2,10 +2,15 @@ # SPDX-License-Identifier: Apache-2.0 import os +import threading import time import pytest +from cuopt.grpc.client.grpc_client import ( + _WAIT_POLL_INTERVAL_S, + _wait_poll_loop, +) from cuopt.grpc.linear_programming import ( Client, GrpcError, @@ -75,6 +80,133 @@ def _assert_demo_lp_solution(client): client.delete(job_id) +class TestWaitPollLoop: + def test_returns_immediately_when_already_terminal(self, monkeypatch): + sleeps = [] + monkeypatch.setattr(time, "sleep", sleeps.append) + status = _wait_poll_loop( + lambda job_id: JobStatus.COMPLETED, "job", 0, GrpcError + ) + assert status is JobStatus.COMPLETED + assert sleeps == [] + + def test_sleeps_between_in_flight_polls(self, monkeypatch): + polls = [] + sleeps = [] + + def get_status(job_id): + polls.append(job_id) + if len(polls) < 3: + return JobStatus.PROCESSING + return JobStatus.CANCELLED + + monkeypatch.setattr(time, "sleep", sleeps.append) + status = _wait_poll_loop(get_status, "abc", 0, GrpcError) + assert status is JobStatus.CANCELLED + assert polls == ["abc", "abc", "abc"] + assert sleeps == [_WAIT_POLL_INTERVAL_S, _WAIT_POLL_INTERVAL_S] + + def test_timeout_raises_after_deadline(self, monkeypatch): + ticks = iter([100.0, 100.0, 101.0]) + monkeypatch.setattr(time, "monotonic", lambda: next(ticks)) + monkeypatch.setattr(time, "sleep", lambda seconds: None) + with pytest.raises( + GrpcError, match="Timeout waiting for job completion" + ): + _wait_poll_loop( + lambda job_id: JobStatus.QUEUED, "job", 1, GrpcError + ) + + def test_rejects_negative_timeout(self): + with pytest.raises( + GrpcError, match="timeout_seconds must be non-negative" + ): + _wait_poll_loop( + lambda job_id: JobStatus.PROCESSING, "job", -1, GrpcError + ) + + def test_other_thread_runs_during_sleep(self): + started = threading.Event() + progressed = threading.Event() + polls = {"n": 0} + + def get_status(job_id): + polls["n"] += 1 + if polls["n"] == 1: + started.set() + return JobStatus.PROCESSING + assert progressed.is_set() + return JobStatus.COMPLETED + + def worker(): + assert started.wait(timeout=2) + progressed.set() + + thread = threading.Thread(target=worker) + thread.start() + status = _wait_poll_loop( + get_status, + "job", + 0, + GrpcError, + poll_interval_s=0.05, + ) + thread.join(timeout=2) + assert status is JobStatus.COMPLETED + assert progressed.is_set() + + def test_client_wait_releases_gil(self, monkeypatch): + """A GIL-bound spinner must progress during Client.wait's poll sleep. + + Hits are counted only inside the sleep, so this fails if wait() still + called the C++ WaitForCompletion/sleep path that holds the GIL. + """ + import cuopt.grpc.client.grpc_client as grpc_mod + + polls = {"n": 0} + + def fake_status(self, job_id): + polls["n"] += 1 + if polls["n"] == 1: + return JobStatus.PROCESSING + return JobStatus.COMPLETED + + monkeypatch.setattr(Client, "status", fake_status) + client = Client.__new__(Client) + + hits = {"n": 0} + during_sleep = [] + stop = threading.Event() + + def spinner(): + while not stop.is_set(): + hits["n"] += 1 + + real_sleep = time.sleep + + def tracking_sleep(seconds): + assert seconds == _WAIT_POLL_INTERVAL_S + before = hits["n"] + real_sleep(0.15) + during_sleep.append(hits["n"] - before) + + monkeypatch.setattr(grpc_mod.time, "sleep", tracking_sleep) + thread = threading.Thread(target=spinner) + thread.start() + try: + status = client.wait("job") + finally: + stop.set() + thread.join(timeout=2) + + assert status is JobStatus.COMPLETED + assert during_sleep, "wait() never slept between status polls" + assert during_sleep[0] > 1000, ( + "spinner made no progress during wait sleep; GIL likely held " + f"(hits={during_sleep[0]})" + ) + + class TestTlsConfig: def test_mtls_requires_both_client_materials(self): pem = "-----BEGIN CERTIFICATE-----\nabc\n-----END CERTIFICATE-----" @@ -272,7 +404,7 @@ def get_solution( job_id = client.submit(problem, settings) client.start_incumbent_stream(job_id, settings=settings) - terminal = _poll_until_complete(client, job_id, _MIP_NAMES) + terminal = client.wait(job_id, timeout=120) assert terminal == JobStatus.COMPLETED client.join_incumbent_stream(job_id) @@ -284,6 +416,70 @@ def get_solution( assert solution is not None client.delete(job_id) + def test_mip_incumbent_stream_live_during_wait(self, grpc_server): + """Incumbent callbacks must fire during wait(), not in a burst after. + + The 2-variable MIP in test_mip_incumbent_stream finishes too fast to + tell. swath1 with a time limit stays PROCESSING long enough that a + GIL-holding wait() would delay every callback until join(). + """ + if not os.path.isfile(_SWATH1_MPS): + pytest.skip(f"dataset not found: {_SWATH1_MPS}") + + class TimedIncumbents(GetSolutionCallback): + def __init__(self): + super().__init__() + self.times = [] + self.costs = [] + + def get_solution( + self, solution, solution_cost, solution_bound, user_data + ): + self.times.append(time.monotonic()) + self.costs.append(float(solution_cost[0])) + + collector = TimedIncumbents() + settings = SolverSettings() + settings.set_mip_callback(collector, None) + settings.set_parameter(CUOPT_TIME_LIMIT, 8) + + client = Client("localhost", grpc_server) + job_id = client.submit(Read(_SWATH1_MPS), settings) + client.start_incumbent_stream( + job_id, settings=settings, poll_interval_ms=200 + ) + try: + terminal = client.wait(job_id, timeout=30) + wait_end = time.monotonic() + client.join_incumbent_stream(job_id) + finally: + client.delete(job_id) + + if terminal != JobStatus.COMPLETED: + pytest.skip(f"job did not complete ({terminal.name})") + if len(collector.times) < 2: + pytest.skip( + "need >=2 incumbents to test live delivery, got " + f"{len(collector.times)}" + ) + + n_before = sum(t < wait_end for t in collector.times) + spread = max(collector.times) - min(collector.times) + lag = wait_end - min(collector.times) + print( + f"incumbents={len(collector.times)} before_wait={n_before} " + f"spread={spread:.3f}s first_to_wait_end={lag:.3f}s" + ) + assert n_before >= 1, ( + f"all {len(collector.times)} incumbents arrived at/after " + f"wait() returned (spread={spread:.4f}s); GIL likely held" + ) + assert spread > 0.15, ( + f"incumbent timestamps clustered in {spread:.4f}s " + f"(n={len(collector.times)}, lag_to_wait_end={lag:.4f}s); " + "likely dumped as a burst at completion" + ) + @pytest.mark.xdist_group(name="grpc_server") @pytest.mark.filterwarnings("ignore::DeprecationWarning")