diff --git a/src/logits/_service_client.py b/src/logits/_service_client.py index fa044d2..cef7597 100644 --- a/src/logits/_service_client.py +++ b/src/logits/_service_client.py @@ -39,6 +39,8 @@ def __init__( @staticmethod def _wrap_training_client(training_client: tinker.TrainingClient) -> LogitsTrainingClient: + if isinstance(training_client, LogitsTrainingClient): + return training_client # tinker.TrainingClient stores the model seq id in a private field. # We only need it to construct an equivalent wrapper instance. return LogitsTrainingClient( diff --git a/tests/test_training_client_state.py b/tests/test_training_client_state.py new file mode 100644 index 0000000..2a30f37 --- /dev/null +++ b/tests/test_training_client_state.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +import asyncio +import json + +import httpx +import pytest +import respx +import tinker + +import logits + + +@pytest.fixture +def wire_requests(): + requests = [] + results = {} + + def handle(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) if request.content else {} + endpoint = request.url.path + requests.append((endpoint, body)) + if endpoint == "/api/v1/client/config": + result = {} + elif endpoint == "/api/v1/create_session": + result = {"session_id": "local-session"} + elif endpoint == "/api/v1/weights_info": + result = {"base_model": "Qwen/Qwen3-8B", "is_lora": True, "lora_rank": 32} + elif endpoint == "/api/v1/create_model": + results["create"] = {"model_id": "local-model", "type": "create_model"} + result = {"request_id": "create"} + elif endpoint == "/api/v1/load_weights": + results["load"] = {"path": body["path"], "type": "load_weights"} + result = {"request_id": "load"} + elif endpoint == "/api/v1/optim_step": + results["optim"] = {"metrics": {}} + result = {"request_id": "optim"} + elif endpoint == "/api/v1/retrieve_future": + result = results[body["request_id"]] + else: + raise AssertionError(f"Unexpected request: {endpoint}") + return httpx.Response(200, json=result) + + with respx.mock(assert_all_called=False) as router: + router.route(host="logits.test").mock(side_effect=handle) + yield requests + + +@pytest.mark.parametrize("sdk", [tinker, logits], ids=["tinker-control", "logits"]) +def test_fresh_training_client_starts_at_sequence_one(sdk, wire_requests): + service = sdk.ServiceClient(base_url="http://logits.test", api_key="tml-fake-key") + try: + training = service.create_lora_training_client(base_model="Qwen/Qwen3-8B") + training.optim_step(sdk.AdamParams(learning_rate=1e-4)).result() + assert [ + body["seq_id"] + for endpoint, body in wire_requests + if endpoint == "/api/v1/optim_step" + ] == [1] + finally: + service.holder.close() + + +@pytest.mark.parametrize("sdk", [tinker, logits], ids=["tinker-control", "logits"]) +@pytest.mark.parametrize( + "optimizer", [False, True], ids=["weights", "weights-and-optimizer"] +) +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("scheme", ["tinker", "logits"]) +def test_restored_training_client_continues_sequence( + sdk, optimizer, asynchronous, scheme, wire_requests +): + service = sdk.ServiceClient(base_url="http://logits.test", api_key="tml-fake-key") + method = "create_training_client_from_state" + ( + "_with_optimizer" if optimizer else "" + ) + checkpoint = f"{scheme}://previous-run/weights/checkpoint" + try: + if asynchronous: + + async def run(): + training = await getattr(service, method + "_async")(checkpoint) + future = await training.optim_step_async( + sdk.AdamParams(learning_rate=1e-4) + ) + await future.result_async() + + asyncio.run(run()) + else: + training = getattr(service, method)(checkpoint) + training.optim_step(sdk.AdamParams(learning_rate=1e-4)).result() + sequences = [ + (endpoint, body["seq_id"]) + for endpoint, body in wire_requests + if endpoint in ("/api/v1/load_weights", "/api/v1/optim_step") + ] + assert sequences == [("/api/v1/load_weights", 1), ("/api/v1/optim_step", 2)] + finally: + service.holder.close()