diff --git a/dspy/clients/lm.py b/dspy/clients/lm.py index b6fa8ffb4a..1a62c06d6a 100644 --- a/dspy/clients/lm.py +++ b/dspy/clients/lm.py @@ -391,10 +391,12 @@ def _run_finetune_job(self, job: TrainingJob): train_kwargs=job.train_kwargs, ) lm = self.copy(model=model) - job.set_result(lm) + if not job.cancelled(): + job.set_result(lm) except Exception as err: logger.error(err) - job.set_result(err) + if not job.cancelled(): + job.set_result(err) def infer_provider(self) -> Provider: if OpenAIProvider.is_provider_model(self.model): diff --git a/tests/clients/test_openai_provider.py b/tests/clients/test_openai_provider.py new file mode 100644 index 0000000000..1b99368658 --- /dev/null +++ b/tests/clients/test_openai_provider.py @@ -0,0 +1,337 @@ +"""Regression tests for the OpenAI fine-tune cancellation race. + +These tests exercise the real ``LM.finetune`` -> ``LM._run_finetune_job`` -> +``OpenAIProvider.finetune`` -> ``TrainingJobOpenAI.cancel`` code paths, mocking +only the ``openai.*`` SDK boundary so no network is required. + +They pin the worker thread at deterministic points (via ``threading.Event`` +pairs) to reproduce the two timing windows in which ``cancel()`` can land while +a fine-tune is in flight, and assert that the worker thread never dies with an +uncaught ``concurrent.futures.InvalidStateError`` and that ``job.result()`` +surfaces ``CancelledError`` cleanly. +""" + +from __future__ import annotations + +import threading +import time +from concurrent.futures import CancelledError +from contextlib import ExitStack +from unittest import mock + +import openai +import pytest + +import dspy +import dspy.clients.openai as oai_mod +from dspy.clients.openai import OpenAIProvider +from dspy.clients.utils_finetune import TrainDataFormat + +# Real openai SDK raises ValueError on a None id (verified against openai>=1.66); +# the mocks below mirror that behavior so the worker observes the same +# exception it would see against a live account. +_NONE_JOB_MSG = "Expected a non-empty value for `fine_tuning_job_id` but received None" +_NONE_FILE_MSG = "Expected a non-empty value for `file_id` but received None" + +TRAIN_DATA = [ + { + "messages": [ + {"role": "user", "content": "Hello, how are you?"}, + {"role": "assistant", "content": "I'm doing great, thank you!"}, + ] + }, + { + "messages": [ + {"role": "user", "content": "What is the capital of France?"}, + {"role": "assistant", "content": "Paris!"}, + ] + }, + { + "messages": [ + {"role": "user", "content": "What is the capital of Germany?"}, + {"role": "assistant", "content": "Berlin!"}, + ] + }, +] + + +class _FakeFile: + def __init__(self, file_id): + self.id = file_id + + +class _FakeJob: + def __init__(self, job_id, *, status="running", estimated_finish=None, fine_tuned_model=None): + self.id = job_id + self.status = status + self.estimated_finish = estimated_finish + self.fine_tuned_model = fine_tuned_model + + +class _FakePage: + def __init__(self, data=None): + self.data = data or [] + + +class _Gate: + """A pair of threading.Events used to pause/release the worker thread.""" + + def __init__(self): + self.paused = threading.Event() + self.release = threading.Event() + + def hold(self): + self.paused.set() + self.release.wait(timeout=30) + + def reset(self): + self.paused.clear() + self.release.clear() + + +@pytest.fixture +def patch_sleep(monkeypatch): + """Make the wait_for_job poll loop instantaneous so tests don't sleep 20s.""" + monkeypatch.setattr(oai_mod.time, "sleep", lambda *_a, **_k: None) + + +@pytest.fixture +def uncaught(): + """Capture any exception that escapes a worker thread via threading.excepthook.""" + captured = {} + original = threading.excepthook + + def hook(args): + captured["exc"] = args.exc_value + + threading.excepthook = hook + try: + yield captured + finally: + threading.excepthook = original + + +def _build_openai_lm(): + return dspy.LM(model="openai/gpt-4o-mini", api_key="sk-fake", cache=False) + + +def _start_finetune(lm): + """Invoke the real LM.finetune and return the live TrainingJobOpenAI.""" + return lm.finetune( + train_data=TRAIN_DATA, + train_data_format=TrainDataFormat.CHAT, + train_kwargs={}, + ) + + +def _patch_boundary(stack, patches): + """Apply ``mock.patch.object`` on each (module, attr, side_effect) tuple.""" + for module, attr, side_effect in patches: + stack.enter_context(mock.patch.object(module, attr, side_effect=side_effect)) + + +def _make_boundary(gate_poll=None, gate_upload=None, terminal_status="succeeded"): + """Build mock callables for the openai.* SDK boundary. + + Args: + gate_poll: ``_Gate`` to hold the worker inside ``jobs.list_events`` + (after ``provider_job_id`` has been set) -- reproduces Window 1. + gate_upload: ``_Gate`` to hold the worker inside ``files.create`` + (before ``provider_file_id`` is set) -- reproduces Window 2a. + terminal_status: status returned by ``jobs.retrieve`` for a real id. + """ + if gate_upload is not None: + gate_upload.reset() + if gate_poll is not None: + gate_poll.reset() + + def fake_files_create(*args, **kwargs): + if gate_upload is not None: + gate_upload.hold() + return _FakeFile("file-real") + + def fake_files_retrieve(file_id): + if file_id is None: + raise ValueError(_NONE_FILE_MSG) + return _FakeFile(file_id) + + def fake_files_delete(file_id): + return None + + def fake_jobs_create(*args, **kwargs): + return _FakeJob( + "ftjob-real", + status="running", + estimated_finish=int(time.time()) + 600, + fine_tuned_model="ft:gpt-4o-mini:custom:ftjob-real", + ) + + def fake_jobs_retrieve(job_id): + if job_id is None: + raise ValueError(_NONE_JOB_MSG) + return _FakeJob( + job_id, + status=terminal_status, + estimated_finish=int(time.time()) + 600, + fine_tuned_model=f"ft:gpt-4o-mini:custom:{job_id}", + ) + + def fake_jobs_list_events(*args, fine_tuning_job_id=None, **kwargs): + if fine_tuning_job_id is None: + raise ValueError(_NONE_JOB_MSG) + if gate_poll is not None: + gate_poll.hold() + return _FakePage([]) + + def fake_jobs_cancel(job_id): + return None + + return [ + (openai.files, "create", fake_files_create), + (openai.files, "retrieve", fake_files_retrieve), + (openai.files, "delete", fake_files_delete), + (openai.fine_tuning.jobs, "create", fake_jobs_create), + (openai.fine_tuning.jobs, "retrieve", fake_jobs_retrieve), + (openai.fine_tuning.jobs, "list_events", fake_jobs_list_events), + (openai.fine_tuning.jobs, "cancel", fake_jobs_cancel), + ] + + +def test_finetune_success_path_delivers_trained_lm(patch_sleep, uncaught): + """Regression: with no cancel(), the trained LM is delivered via result().""" + with ExitStack() as stack: + _patch_boundary(stack, _make_boundary()) + lm = _build_openai_lm() + job = _start_finetune(lm) + job.thread.join(timeout=30) + + assert not job.thread.is_alive() + assert "exc" not in uncaught + result = job.result(timeout=1) + assert isinstance(result, dspy.LM) + assert result.model == "ft:gpt-4o-mini:custom:ftjob-real" + + +def test_finetune_error_path_propagates_exception_when_not_cancelled(patch_sleep, uncaught): + """Regression: an error from the provider (not a cancel) is still stored as + the future's result so result() surfaces it; the cancelled-guard must not + leave the future pending when the job was never cancelled. + + Note: dspy stores the exception via ``set_result(err)`` (not + ``set_exception``), so ``result()`` returns the exception object rather than + re-raising it. The regression guarantee is that the future reaches FINISHED + and the exception is delivered to the caller. + """ + boom = RuntimeError("provider blew up") + + def evil_finetune(job, model, train_data, train_data_format, train_kwargs): + raise boom + + with ExitStack() as stack: + stack.enter_context(mock.patch.object(OpenAIProvider, "finetune", staticmethod(evil_finetune))) + lm = _build_openai_lm() + job = _start_finetune(lm) + job.thread.join(timeout=30) + + assert not job.thread.is_alive() + assert "exc" not in uncaught + assert job.done() + assert job.cancelled() is False + # The exception is surfaced through the future (stored as the result value). + assert job.result(timeout=1) is boom + + +def test_cancel_during_polling_does_not_crash_worker(patch_sleep, uncaught): + """Window 1: cancel() lands while the worker is polling (provider_job_id set). + + Guarantees: + - the worker thread exits cleanly (no uncaught InvalidStateError) + - job.cancelled() is True + - job.result() raises CancelledError (not the inner ValueError) + - the remote cancel API was actually invoked with the live job id + """ + gate = _Gate() + cancel_calls = [] + patches = _make_boundary(gate_poll=gate, terminal_status="running") + patches[-1] = (openai.fine_tuning.jobs, "cancel", lambda job_id: cancel_calls.append(job_id)) + + with ExitStack() as stack: + _patch_boundary(stack, patches) + lm = _build_openai_lm() + job = _start_finetune(lm) + + assert gate.paused.wait(timeout=30) + assert job.provider_job_id == "ftjob-real" + assert job.provider_file_id == "file-real" + assert not job.cancelled() + + job.cancel() + + assert job.cancelled() is True + assert job.provider_job_id is None + assert job.provider_file_id is None + + gate.release.set() + job.thread.join(timeout=30) + + assert not job.thread.is_alive() + assert "exc" not in uncaught, f"worker thread died with uncaught {uncaught.get('exc')!r}" + assert cancel_calls == ["ftjob-real"] + with pytest.raises(CancelledError): + job.result(timeout=1) + + +def test_cancel_during_upload_does_not_crash_worker(patch_sleep, uncaught): + """Window 2a: cancel() lands during upload_data (provider_file_id NOT set). + + Guarantees: + - the worker thread exits cleanly (no uncaught InvalidStateError) even + though it resumes, starts the remote job, and runs it to completion + - job.cancelled() is True + - job.result() raises CancelledError (the trained LM is NOT delivered + through the TrainingJob API) + """ + gate = _Gate() + creates = [] + + def counting_create(*args, **kwargs): + creates.append(kwargs.get("training_file")) + return _FakeJob( + "ftjob-real", + status="running", + estimated_finish=int(time.time()) + 600, + fine_tuned_model="ft:gpt-4o-mini:custom:ftjob-real", + ) + + patches = _make_boundary(gate_upload=gate, terminal_status="succeeded") + # Replace the jobs.create mock with a counting variant. + patches[3] = (openai.fine_tuning.jobs, "create", counting_create) + + with ExitStack() as stack: + _patch_boundary(stack, patches) + lm = _build_openai_lm() + job = _start_finetune(lm) + + assert gate.paused.wait(timeout=30) + assert job.provider_job_id is None + assert job.provider_file_id is None + assert not job.cancelled() + + job.cancel() + + assert job.cancelled() is True + # cancel() during upload_data cannot delete the file (its id is unknown + # to cancel at this moment), so provider_file_id stays None here. + assert job.provider_file_id is None + + gate.release.set() + job.thread.join(timeout=30) + + assert not job.thread.is_alive() + assert "exc" not in uncaught, f"worker thread died with uncaught {uncaught.get('exc')!r}" + # The worker resumed from upload_data and started the remote job with the + # just-uploaded file. The cost-prevention concern is tracked separately; this + # test only locks down the crash fix. The trained LM must NOT be delivered. + assert creates == ["file-real"] + with pytest.raises(CancelledError): + job.result(timeout=1)