mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-15 04:19:08 +00:00
test(integration): cover fine_tuning endpoint contracts (#44697)
* test(integration): cover fine_tuning endpoint contracts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): prove fine_tuning ownership hook, header spelling and drain order Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(fine_tuning): assert unified job id and retrieve ownership through route-granted keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(fine_tuning): assert managed file and job ownership rows in Postgres Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(fine_tuning): cover validation_file ownership, default retrieve ownership and full job objects Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mateo <mateo@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
b37eac29f0
commit
b22a0bd9e5
2 changed files with 934 additions and 0 deletions
442
tests/integration/compatibility/test_fine_tuning_jobs.py
Normal file
442
tests/integration/compatibility/test_fine_tuning_jobs.py
Normal file
|
|
@ -0,0 +1,442 @@
|
|||
import base64
|
||||
import json
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from email.message import Message
|
||||
from email.parser import BytesParser
|
||||
from email.policy import HTTP
|
||||
from typing import Final
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually, gateway_from_environment
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from openai import OpenAI
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
|
||||
|
||||
OPENAI_KEY: Final = "sk-fine-tuning-openai"
|
||||
AZURE_KEY: Final = "azure-fine-tuning-key"
|
||||
AMBIENT_OPENAI_KEY: Final = "sk-fine-tuning-ambient-environment"
|
||||
AMBIENT_AZURE_KEY: Final = "azure-fine-tuning-ambient-environment"
|
||||
AZURE_API_VERSION: Final = "2024-10-21"
|
||||
MANAGED_MODEL: Final = "gpt-4.1-openai"
|
||||
DEPLOYMENT_KEYS: Final = ("sk-fine-tuning-deployment-a", "sk-fine-tuning-deployment-b")
|
||||
PROVIDER_JOB: Final = "ftjob-provider-123"
|
||||
RUN: Final = uuid.uuid4().hex[:12]
|
||||
MANAGED_JOB: Final = f"ftjob-managed-{RUN}"
|
||||
TRAINING_LINE: Final = b'{"messages": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}]}\n'
|
||||
JOB_PATH: Final = re.compile(r"^/(?:openai/)?(?:v1/)?fine_tuning/jobs/([^/]+)(/cancel)?$")
|
||||
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
PREFIXES: Final = (pytest.param("/v1", id="v1"), pytest.param("", id="unversioned"))
|
||||
|
||||
|
||||
class Job(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
id: str
|
||||
object: str
|
||||
created_at: int
|
||||
error: JsonValue
|
||||
fine_tuned_model: str | None
|
||||
finished_at: int | None
|
||||
hyperparameters: dict[str, JsonValue]
|
||||
model: str
|
||||
organization_id: str
|
||||
result_files: list[str]
|
||||
seed: int
|
||||
status: str
|
||||
trained_tokens: int | None
|
||||
training_file: str
|
||||
validation_file: str | None
|
||||
estimated_finish: int | None = None
|
||||
integrations: JsonValue = None
|
||||
metadata: JsonValue = None
|
||||
method: JsonValue = None
|
||||
|
||||
|
||||
def _client_job(job_id: str, status: str, training_file: str = "file-abc123") -> Job:
|
||||
return Job.model_validate(
|
||||
{
|
||||
**_job(job_id, status, training_file),
|
||||
"hyperparameters": {"batch_size": None, "learning_rate_multiplier": None, "n_epochs": 3},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _job(job_id: str, status: str, training_file: str = "file-abc123", **fields: JsonValue) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": job_id,
|
||||
"object": "fine_tuning.job",
|
||||
"created_at": 1700000000,
|
||||
"error": None,
|
||||
"fine_tuned_model": None,
|
||||
"finished_at": None,
|
||||
"hyperparameters": {"n_epochs": 3},
|
||||
"model": "gpt-4o-mini-2024-07-18",
|
||||
"organization_id": "org-fine-tuning",
|
||||
"result_files": [],
|
||||
"seed": 42,
|
||||
"status": status,
|
||||
"trained_tokens": None,
|
||||
"training_file": training_file,
|
||||
"validation_file": None,
|
||||
**fields,
|
||||
}
|
||||
|
||||
|
||||
def _provider_file_id(bearer: str) -> str:
|
||||
return f"file-{RUN}-" + bearer.removeprefix("sk-fine-tuning-")
|
||||
|
||||
|
||||
def _bearer(request: Request) -> str:
|
||||
return request.headers.get("authorization", "").removeprefix("Bearer ")
|
||||
|
||||
|
||||
def _environment(wire: Wire) -> dict[str, str]:
|
||||
return {
|
||||
"OPENAI_BASE_URL": f"{wire.url}/v1",
|
||||
"OPENAI_API_BASE": f"{wire.url}/v1",
|
||||
"OPENAI_API_KEY": AMBIENT_OPENAI_KEY,
|
||||
"AZURE_API_BASE": wire.url,
|
||||
"AZURE_API_KEY": AMBIENT_AZURE_KEY,
|
||||
"AZURE_API_VERSION": AZURE_API_VERSION,
|
||||
}
|
||||
|
||||
|
||||
def _reply(body: Mapping[str, JsonValue]) -> Reply:
|
||||
return Reply(body=json.dumps(body).encode())
|
||||
|
||||
|
||||
def _respond(request: Request) -> Reply:
|
||||
path: Final = urlsplit(request.target).path
|
||||
if (request.method, path) == ("GET", "/v1/models"):
|
||||
return _reply({"object": "list", "data": []})
|
||||
if (request.method, path) == ("POST", "/v1/files"):
|
||||
return _reply(
|
||||
{
|
||||
"id": _provider_file_id(_bearer(request)),
|
||||
"object": "file",
|
||||
"bytes": len(TRAINING_LINE),
|
||||
"created_at": 1700000000,
|
||||
"filename": "train.jsonl",
|
||||
"purpose": "fine-tune",
|
||||
"status": "processed",
|
||||
}
|
||||
)
|
||||
if (request.method, path) in {("POST", "/v1/fine_tuning/jobs"), ("POST", "/openai/fine_tuning/jobs")}:
|
||||
sent: Final = JSON_OBJECT.validate_json(request.body)
|
||||
job_id: Final = MANAGED_JOB if _bearer(request) in DEPLOYMENT_KEYS else PROVIDER_JOB
|
||||
return _reply(_job(job_id, "validating_files", str(sent["training_file"])))
|
||||
matched: Final = JOB_PATH.match(path)
|
||||
if matched and matched.group(1) and request.method == "GET":
|
||||
return _reply(_job(matched.group(1), "running"))
|
||||
if matched and matched.group(2) and request.method == "POST":
|
||||
return _reply(_job(matched.group(1), "cancelled"))
|
||||
return Reply(
|
||||
status=404, body=json.dumps({"error": {"message": f"unscripted {request.method} {request.target}"}}).encode()
|
||||
)
|
||||
|
||||
|
||||
def _config(wire: Wire) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model_list": [
|
||||
{
|
||||
"model_name": MANAGED_MODEL,
|
||||
"litellm_params": {"model": "openai/gpt-4.1", "api_key": key, "api_base": f"{wire.url}/v1"},
|
||||
"model_info": {"supported_endpoints": ["/chat/completions", "/fine_tuning"]},
|
||||
}
|
||||
for key in DEPLOYMENT_KEYS
|
||||
],
|
||||
"finetune_settings": [
|
||||
{"custom_llm_provider": "openai", "api_key": OPENAI_KEY, "api_base": f"{wire.url}/v1"},
|
||||
{
|
||||
"custom_llm_provider": "azure",
|
||||
"api_key": AZURE_KEY,
|
||||
"api_base": wire.url,
|
||||
"api_version": AZURE_API_VERSION,
|
||||
},
|
||||
],
|
||||
"general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY", "database_url": "os.environ/DATABASE_URL"},
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
proxy: Gateway
|
||||
wire: Wire
|
||||
|
||||
def sdk(self, prefix: str) -> OpenAI:
|
||||
return OpenAI(
|
||||
base_url=f"{str(self.proxy.client.base_url).rstrip('/')}{prefix}", api_key=self.proxy.key, max_retries=0
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp("fine-tuning")
|
||||
with gateway_from_environment() as shared, wire_server(_respond) as wire:
|
||||
config: Final = directory / "config.yaml"
|
||||
config.write_text(yaml.safe_dump(_config(wire)))
|
||||
with owned_proxy(shared, directory, _environment(wire), config=config) as proxy:
|
||||
boot: Final[list[Request]] = []
|
||||
eventually(lambda: boot.extend(wire.drain()) or len(boot), lambda count: count >= len(DEPLOYMENT_KEYS))
|
||||
assert sorted((request.method, request.target, _bearer(request)) for request in boot) == [
|
||||
("GET", "/v1/models", key) for key in DEPLOYMENT_KEYS
|
||||
]
|
||||
yield Rig(proxy, wire)
|
||||
assert wire.drain() == ()
|
||||
|
||||
|
||||
def _drained(rig: Rig) -> tuple[Request, ...]:
|
||||
return rig.wire.drain()
|
||||
|
||||
|
||||
def _only(requests: tuple[Request, ...], method: str, target: str) -> Request:
|
||||
assert [(request.method, request.target) for request in requests] == [(method, target)]
|
||||
return requests[0]
|
||||
|
||||
|
||||
def _decoded(unified_id: str) -> str:
|
||||
return base64.urlsafe_b64decode(unified_id + "=" * (-len(unified_id) % 4)).decode()
|
||||
|
||||
|
||||
def _multipart_parts(request: Request) -> tuple[Message, ...]:
|
||||
envelope: Final = f"content-type: {request.headers['content-type']}\r\n\r\n".encode() + request.body
|
||||
parsed: Final = BytesParser(policy=HTTP).parsebytes(envelope)
|
||||
assert parsed.is_multipart(), request.headers["content-type"]
|
||||
return tuple(parsed.iter_parts())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", PREFIXES)
|
||||
def test_create_forwards_every_documented_optional_field(rig: Rig, prefix: str) -> None:
|
||||
with rig.sdk(prefix) as client:
|
||||
raw: Final = client.fine_tuning.jobs.with_raw_response.create(
|
||||
model="gpt-4o-mini",
|
||||
training_file="file-abc123",
|
||||
validation_file="file-xyz789",
|
||||
hyperparameters={"n_epochs": 3, "batch_size": "auto", "learning_rate_multiplier": 0.1},
|
||||
suffix="custom-model",
|
||||
seed=42,
|
||||
extra_body={"custom_llm_provider": "openai"},
|
||||
)
|
||||
sent: Final = _only(_drained(rig), "POST", "/v1/fine_tuning/jobs")
|
||||
assert sent.headers["authorization"] == f"Bearer {OPENAI_KEY}"
|
||||
assert JSON_OBJECT.validate_json(sent.body) == {
|
||||
"model": "gpt-4o-mini",
|
||||
"training_file": "file-abc123",
|
||||
"validation_file": "file-xyz789",
|
||||
"hyperparameters": {"n_epochs": 3, "batch_size": "auto", "learning_rate_multiplier": 0.1},
|
||||
"suffix": "custom-model",
|
||||
"seed": 42,
|
||||
}, sent.body
|
||||
assert raw.status_code == 200, raw.text
|
||||
assert Job.model_validate_json(raw.text) == _client_job(PROVIDER_JOB, "validating_files"), raw.text
|
||||
|
||||
|
||||
def test_create_forwards_the_method_object(rig: Rig) -> None:
|
||||
pytest.skip(
|
||||
"BUG: create drops the documented method object, the provider receives only model, training_file and hyperparameters {}"
|
||||
)
|
||||
with rig.sdk("/v1") as client:
|
||||
raw: Final = client.fine_tuning.jobs.with_raw_response.create(
|
||||
model="gpt-4o-mini",
|
||||
training_file="file-abc123",
|
||||
method={"type": "supervised", "supervised": {"hyperparameters": {"n_epochs": 2}}},
|
||||
extra_body={"custom_llm_provider": "openai"},
|
||||
)
|
||||
sent: Final = _only(_drained(rig), "POST", "/v1/fine_tuning/jobs")
|
||||
assert raw.status_code == 200, raw.text
|
||||
assert JSON_OBJECT.validate_json(sent.body) == {
|
||||
"model": "gpt-4o-mini",
|
||||
"training_file": "file-abc123",
|
||||
"hyperparameters": {},
|
||||
"method": {"type": "supervised", "supervised": {"hyperparameters": {"n_epochs": 2}}},
|
||||
}, sent.body
|
||||
|
||||
|
||||
def test_create_forwards_integrations_as_the_sdk_sends_them(rig: Rig) -> None:
|
||||
pytest.skip(
|
||||
"BUG: create rejects the SDK integrations objects with 422 because the proxy types integrations as list[str]"
|
||||
)
|
||||
with rig.sdk("/v1") as client:
|
||||
raw: Final = client.fine_tuning.jobs.with_raw_response.create(
|
||||
model="gpt-4o-mini",
|
||||
training_file="file-abc123",
|
||||
integrations=[{"type": "wandb", "wandb": {"project": "fine-tuning"}}],
|
||||
extra_body={"custom_llm_provider": "openai"},
|
||||
)
|
||||
sent: Final = _only(_drained(rig), "POST", "/v1/fine_tuning/jobs")
|
||||
assert raw.status_code == 200, raw.text
|
||||
assert JSON_OBJECT.validate_json(sent.body) == {
|
||||
"model": "gpt-4o-mini",
|
||||
"training_file": "file-abc123",
|
||||
"hyperparameters": {},
|
||||
"integrations": [{"type": "wandb", "wandb": {"project": "fine-tuning"}}],
|
||||
}, sent.body
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", PREFIXES)
|
||||
def test_retrieve_forwards_the_raw_provider_job_id(rig: Rig, prefix: str) -> None:
|
||||
with rig.sdk(prefix) as client:
|
||||
raw: Final = client.fine_tuning.jobs.with_raw_response.retrieve(
|
||||
PROVIDER_JOB, extra_query={"custom_llm_provider": "openai"}
|
||||
)
|
||||
sent: Final = _only(_drained(rig), "GET", f"/v1/fine_tuning/jobs/{PROVIDER_JOB}")
|
||||
assert sent.headers["authorization"] == f"Bearer {OPENAI_KEY}"
|
||||
assert sent.body == b""
|
||||
assert raw.status_code == 200, raw.text
|
||||
assert Job.model_validate_json(raw.text) == _client_job(PROVIDER_JOB, "running"), raw.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", PREFIXES)
|
||||
def test_cancel_sends_the_cancel_to_the_provider(rig: Rig, prefix: str) -> None:
|
||||
with rig.sdk(prefix) as client:
|
||||
raw: Final = client.fine_tuning.jobs.with_raw_response.cancel(
|
||||
PROVIDER_JOB, extra_body={"custom_llm_provider": "openai"}
|
||||
)
|
||||
sent: Final = _only(_drained(rig), "POST", f"/v1/fine_tuning/jobs/{PROVIDER_JOB}/cancel")
|
||||
assert sent.headers["authorization"] == f"Bearer {OPENAI_KEY}"
|
||||
assert sent.body == b""
|
||||
assert raw.status_code == 200, raw.text
|
||||
assert Job.model_validate_json(raw.text) == _client_job(PROVIDER_JOB, "cancelled"), raw.text
|
||||
|
||||
|
||||
def test_azure_create_merges_finetune_settings_and_keeps_azure_fields(rig: Rig) -> None:
|
||||
with rig.sdk("/v1") as client:
|
||||
raw: Final = client.fine_tuning.jobs.with_raw_response.create(
|
||||
model="gpt-35-turbo-1106",
|
||||
training_file="file-abc",
|
||||
extra_body={
|
||||
"custom_llm_provider": "azure",
|
||||
"hyperparameters": {"n_epochs": 3, "prompt_loss_weight": 0.01},
|
||||
"trainingType": 2,
|
||||
},
|
||||
)
|
||||
sent: Final = _only(_drained(rig), "POST", f"/openai/fine_tuning/jobs?api-version={AZURE_API_VERSION}")
|
||||
assert sent.headers["api-key"] == AZURE_KEY
|
||||
assert JSON_OBJECT.validate_json(sent.body) == {
|
||||
"model": "gpt-35-turbo-1106",
|
||||
"training_file": "file-abc",
|
||||
"hyperparameters": {"n_epochs": 3},
|
||||
"trainingType": 2,
|
||||
"prompt_loss_weight": 0.01,
|
||||
}, sent.body
|
||||
assert raw.status_code == 200, raw.text
|
||||
assert Job.model_validate_json(raw.text) == _client_job(PROVIDER_JOB, "validating_files", "file-abc"), raw.text
|
||||
|
||||
|
||||
def test_azure_cancel_merges_finetune_settings(rig: Rig) -> None:
|
||||
with rig.sdk("/v1") as client:
|
||||
raw: Final = client.fine_tuning.jobs.with_raw_response.cancel(
|
||||
PROVIDER_JOB, extra_body={"custom_llm_provider": "azure"}
|
||||
)
|
||||
sent: Final = _only(
|
||||
_drained(rig), "POST", f"/openai/fine_tuning/jobs/{PROVIDER_JOB}/cancel?api-version={AZURE_API_VERSION}"
|
||||
)
|
||||
assert sent.headers["api-key"] == AZURE_KEY
|
||||
assert sent.body == b""
|
||||
assert raw.status_code == 200, raw.text
|
||||
assert Job.model_validate_json(raw.text) == _client_job(PROVIDER_JOB, "cancelled"), raw.text
|
||||
|
||||
|
||||
def test_create_reads_the_documented_custom_llm_provider_header(rig: Rig) -> None:
|
||||
pytest.skip(
|
||||
"BUG: create ignores the documented custom-llm-provider header and returns 500"
|
||||
" Invalid request, No litellm managed file id or custom_llm_provider provided."
|
||||
)
|
||||
with rig.sdk("/v1") as client:
|
||||
raw: Final = client.fine_tuning.jobs.with_raw_response.create(
|
||||
model="gpt-35-turbo-1106", training_file="file-abc", extra_headers={"custom-llm-provider": "azure"}
|
||||
)
|
||||
sent: Final = _only(_drained(rig), "POST", f"/openai/fine_tuning/jobs?api-version={AZURE_API_VERSION}")
|
||||
assert sent.headers["api-key"] == AZURE_KEY
|
||||
assert JSON_OBJECT.validate_json(sent.body) == {
|
||||
"model": "gpt-35-turbo-1106",
|
||||
"training_file": "file-abc",
|
||||
"hyperparameters": {},
|
||||
}, sent.body
|
||||
assert raw.status_code == 200, raw.text
|
||||
assert Job.model_validate_json(raw.text) == _client_job(PROVIDER_JOB, "validating_files", "file-abc"), raw.text
|
||||
|
||||
|
||||
def test_cancel_reads_the_documented_custom_llm_provider_header(rig: Rig) -> None:
|
||||
pytest.skip(
|
||||
"BUG: cancel ignores the documented custom-llm-provider header and the azure finetune_settings entry"
|
||||
" and sends the cancel to the default OpenAI base URL and key"
|
||||
)
|
||||
with rig.sdk("/v1") as client:
|
||||
raw: Final = client.fine_tuning.jobs.with_raw_response.cancel(
|
||||
PROVIDER_JOB, extra_headers={"custom-llm-provider": "azure"}
|
||||
)
|
||||
sent: Final = _only(
|
||||
_drained(rig), "POST", f"/openai/fine_tuning/jobs/{PROVIDER_JOB}/cancel?api-version={AZURE_API_VERSION}"
|
||||
)
|
||||
assert sent.headers["api-key"] == AZURE_KEY
|
||||
assert sent.body == b""
|
||||
assert raw.status_code == 200, raw.text
|
||||
assert Job.model_validate_json(raw.text) == _client_job(PROVIDER_JOB, "cancelled"), raw.text
|
||||
|
||||
|
||||
def test_retrieve_reads_the_documented_custom_llm_provider_header(rig: Rig) -> None:
|
||||
pytest.skip(
|
||||
"BUG: retrieve ignores the documented custom-llm-provider header and returns 400"
|
||||
" Invalid request, No litellm managed file id or custom_llm_provider provided."
|
||||
)
|
||||
with rig.sdk("/v1") as client:
|
||||
raw: Final = client.fine_tuning.jobs.with_raw_response.retrieve(
|
||||
PROVIDER_JOB, extra_headers={"custom-llm-provider": "azure"}
|
||||
)
|
||||
sent: Final = _only(
|
||||
_drained(rig), "GET", f"/openai/fine_tuning/jobs/{PROVIDER_JOB}?api-version={AZURE_API_VERSION}"
|
||||
)
|
||||
assert sent.headers["api-key"] == AZURE_KEY
|
||||
assert sent.body == b""
|
||||
assert raw.status_code == 200, raw.text
|
||||
assert Job.model_validate_json(raw.text) == _client_job(PROVIDER_JOB, "running"), raw.text
|
||||
|
||||
|
||||
def test_managed_create_sends_the_provider_file_id_of_the_pinned_deployment(rig: Rig) -> None:
|
||||
pytest.skip(
|
||||
"BUG: managed create sends the base64 LiteLLM file id and the prefixed model openai/gpt-4.1 to the provider"
|
||||
" and returns the decoded unified id as training_file"
|
||||
)
|
||||
with rig.sdk("/v1") as client:
|
||||
uploaded: Final = client.files.with_raw_response.create(
|
||||
file=("train.jsonl", TRAINING_LINE, "application/jsonl"),
|
||||
purpose="fine-tune",
|
||||
extra_body={"target_model_names": MANAGED_MODEL},
|
||||
)
|
||||
assert uploaded.status_code == 200, uploaded.text
|
||||
managed_file: Final = JSON_OBJECT.validate_json(uploaded.text)["id"]
|
||||
assert isinstance(managed_file, str) and _decoded(managed_file).startswith("litellm_proxy:"), uploaded.text
|
||||
uploads: Final = _drained(rig)
|
||||
assert sorted((request.method, request.target, _bearer(request)) for request in uploads) == [
|
||||
("POST", "/v1/files", key) for key in DEPLOYMENT_KEYS
|
||||
]
|
||||
for upload in uploads:
|
||||
parts = {part.get_param("name", header="content-disposition"): part for part in _multipart_parts(upload)}
|
||||
assert sorted(parts) == ["file", "purpose"], upload.body[:300]
|
||||
assert parts["file"].get_filename() == "train.jsonl"
|
||||
assert parts["file"].get_payload(decode=True) == TRAINING_LINE
|
||||
assert parts["purpose"].get_filename() is None
|
||||
assert parts["purpose"].get_content() == "fine-tune"
|
||||
|
||||
raw: Final = client.fine_tuning.jobs.with_raw_response.create(model=MANAGED_MODEL, training_file=managed_file)
|
||||
sent: Final = _only(_drained(rig), "POST", "/v1/fine_tuning/jobs")
|
||||
deployment_key: Final = _bearer(sent)
|
||||
assert deployment_key in DEPLOYMENT_KEYS, sent.headers
|
||||
assert JSON_OBJECT.validate_json(sent.body) == {
|
||||
"model": "gpt-4.1",
|
||||
"training_file": _provider_file_id(deployment_key),
|
||||
"hyperparameters": {},
|
||||
}, sent.body
|
||||
assert raw.status_code == 200, raw.text
|
||||
job: Final = Job.model_validate_json(raw.text)
|
||||
assert re.fullmatch(
|
||||
rf"litellm_proxy;model_id:[0-9a-f]{{64}};generic_response_id:{MANAGED_JOB}", _decoded(job.id)
|
||||
), raw.text
|
||||
assert job.training_file == managed_file, raw.text
|
||||
492
tests/integration/security/test_fine_tuning_managed_ids.py
Normal file
492
tests/integration/security/test_fine_tuning_managed_ids.py
Normal file
|
|
@ -0,0 +1,492 @@
|
|||
import base64
|
||||
import json
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.client import Gateway, eventually, gateway_from_environment, object_value, string_value
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from openai import OpenAI
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
|
||||
|
||||
MANAGED_MODEL: Final = "gpt-4.1-managed-ids"
|
||||
DEPLOYMENT_KEY: Final = "sk-fine-tuning-managed-deployment"
|
||||
RUN: Final = uuid.uuid4().hex[:12]
|
||||
PROVIDER_FILE: Final = f"file-managed-{RUN}"
|
||||
PROVIDER_JOB: Final = f"ftjob-managed-{RUN}"
|
||||
DEFAULT_PROVIDER_FILE: Final = f"file-managed-default-{RUN}"
|
||||
DEFAULT_PROVIDER_JOB: Final = f"ftjob-managed-default-{RUN}"
|
||||
TRAINING_LINE: Final = b'{"messages": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}]}\n'
|
||||
JOB_PATH: Final = re.compile(r"^/v1/fine_tuning/jobs/([^/]+?)(/cancel)?$")
|
||||
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
class Job(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
id: str
|
||||
object: str
|
||||
created_at: int
|
||||
error: JsonValue
|
||||
fine_tuned_model: str | None
|
||||
finished_at: int | None
|
||||
hyperparameters: dict[str, JsonValue]
|
||||
status: str
|
||||
model: str
|
||||
organization_id: str
|
||||
result_files: list[str]
|
||||
seed: int
|
||||
trained_tokens: int | None
|
||||
training_file: str
|
||||
validation_file: str | None
|
||||
estimated_finish: int | None = None
|
||||
integrations: JsonValue = None
|
||||
metadata: JsonValue = None
|
||||
method: JsonValue = None
|
||||
|
||||
|
||||
class ProxyError(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
error: dict[str, JsonValue]
|
||||
|
||||
|
||||
def _job(job_id: str, status: str, file_id: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": job_id,
|
||||
"object": "fine_tuning.job",
|
||||
"created_at": 1700000000,
|
||||
"error": None,
|
||||
"fine_tuned_model": None,
|
||||
"finished_at": None,
|
||||
"hyperparameters": {"n_epochs": 3},
|
||||
"model": "gpt-4.1",
|
||||
"organization_id": "org-fine-tuning",
|
||||
"result_files": [],
|
||||
"seed": 42,
|
||||
"status": status,
|
||||
"trained_tokens": None,
|
||||
"training_file": file_id,
|
||||
"validation_file": None,
|
||||
}
|
||||
|
||||
|
||||
def _client_job(job_id: str, status: str, file_id: str) -> Job:
|
||||
return Job.model_validate(
|
||||
{
|
||||
**_job(job_id, status, file_id),
|
||||
"hyperparameters": {"batch_size": None, "learning_rate_multiplier": None, "n_epochs": 3},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _reply(body: Mapping[str, JsonValue]) -> Reply:
|
||||
return Reply(body=json.dumps(body).encode())
|
||||
|
||||
|
||||
def _responder(file_id: str, job_id: str) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
return _respond(request, file_id, job_id)
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _respond(request: Request, file_id: str, job_id: str) -> Reply:
|
||||
path: Final = urlsplit(request.target).path
|
||||
if (request.method, path) == ("GET", "/v1/models"):
|
||||
return _reply({"object": "list", "data": []})
|
||||
if (request.method, path) == ("POST", "/v1/files"):
|
||||
return _reply(
|
||||
{
|
||||
"id": file_id,
|
||||
"object": "file",
|
||||
"bytes": len(TRAINING_LINE),
|
||||
"created_at": 1700000000,
|
||||
"filename": "train.jsonl",
|
||||
"purpose": "fine-tune",
|
||||
"status": "processed",
|
||||
}
|
||||
)
|
||||
if (request.method, path) == ("POST", "/v1/fine_tuning/jobs"):
|
||||
return _reply(_job(job_id, "validating_files", file_id))
|
||||
matched: Final = JOB_PATH.match(path)
|
||||
if matched and request.method == "GET" and not matched.group(2):
|
||||
return _reply(_job(matched.group(1), "running", file_id))
|
||||
if matched and request.method == "POST" and matched.group(2):
|
||||
return _reply(_job(matched.group(1), "cancelled", file_id))
|
||||
return Reply(
|
||||
status=404, body=json.dumps({"error": {"message": f"unscripted {request.method} {request.target}"}}).encode()
|
||||
)
|
||||
|
||||
|
||||
def _config(wire: Wire, litellm_settings: dict[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model_list": [
|
||||
{
|
||||
"model_name": MANAGED_MODEL,
|
||||
"litellm_params": {"model": "openai/gpt-4.1", "api_key": DEPLOYMENT_KEY, "api_base": f"{wire.url}/v1"},
|
||||
"model_info": {"supported_endpoints": ["/chat/completions", "/fine_tuning"]},
|
||||
}
|
||||
],
|
||||
"finetune_settings": [
|
||||
{"custom_llm_provider": "openai", "api_key": "sk-fine-tuning-raw-provider", "api_base": f"{wire.url}/v1"}
|
||||
],
|
||||
"litellm_settings": litellm_settings,
|
||||
"general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY", "database_url": "os.environ/DATABASE_URL"},
|
||||
}
|
||||
|
||||
|
||||
def _unified_job_id(model_id: str, job_id: str) -> str:
|
||||
return (
|
||||
base64.urlsafe_b64encode(f"litellm_proxy;model_id:{model_id};generic_response_id:{job_id}".encode())
|
||||
.decode()
|
||||
.rstrip("=")
|
||||
)
|
||||
|
||||
|
||||
def _deployment_id(proxy: Gateway) -> str:
|
||||
listed: Final = proxy.get("/v1/model/info")["data"]
|
||||
assert isinstance(listed, list), listed
|
||||
ids: Final = [
|
||||
string_value(object_value(object_value(entry)["model_info"])["id"])
|
||||
for entry in listed
|
||||
if object_value(entry)["model_name"] == MANAGED_MODEL
|
||||
]
|
||||
assert len(ids) == 1, listed
|
||||
return ids[0]
|
||||
|
||||
|
||||
def _targets(requests: tuple[Request, ...]) -> list[tuple[str, str, str]]:
|
||||
return [(request.method, request.target, request.headers.get("authorization", "")) for request in requests]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
proxy: Gateway
|
||||
wire: Wire
|
||||
owner: str
|
||||
owner_user: str
|
||||
intruder: str
|
||||
intruder_user: str
|
||||
managed_file: str
|
||||
provider_file: str
|
||||
managed_job: str
|
||||
created_text: str
|
||||
deployment_id: str
|
||||
provider_job: str
|
||||
|
||||
def sdk(self, key: str) -> OpenAI:
|
||||
return OpenAI(base_url=f"{str(self.proxy.client.base_url).rstrip('/')}/v1", api_key=key, max_retries=0)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _rig(directory: Path, litellm_settings: dict[str, JsonValue], file_id: str, job_id: str) -> Iterator[Rig]:
|
||||
with gateway_from_environment() as shared, wire_server(_responder(file_id, job_id)) as wire:
|
||||
config: Final = directory / "config.yaml"
|
||||
config.write_text(yaml.safe_dump(_config(wire, litellm_settings)))
|
||||
with owned_proxy(shared, directory, {}, config=config) as proxy, proxy.scenario() as scenario:
|
||||
eventually(
|
||||
lambda: _targets(wire.drain()), lambda seen: seen == [("GET", "/v1/models", f"Bearer {DEPLOYMENT_KEY}")]
|
||||
)
|
||||
deployment_id: Final = _deployment_id(proxy)
|
||||
owner_user: Final = scenario.user()
|
||||
owner: Final = scenario.key(user_id=owner_user)
|
||||
intruder_user: Final = scenario.user()
|
||||
intruder: Final = scenario.key(user_id=intruder_user)
|
||||
with OpenAI(
|
||||
base_url=f"{str(proxy.client.base_url).rstrip('/')}/v1", api_key=owner, max_retries=0
|
||||
) as client:
|
||||
uploaded: Final = client.files.with_raw_response.create(
|
||||
file=("train.jsonl", TRAINING_LINE, "application/jsonl"),
|
||||
purpose="fine-tune",
|
||||
extra_body={"target_model_names": MANAGED_MODEL},
|
||||
)
|
||||
assert uploaded.status_code == 200, uploaded.text
|
||||
managed_file: Final = uploaded.parse().id
|
||||
created: Final = client.fine_tuning.jobs.with_raw_response.create(
|
||||
model=MANAGED_MODEL, training_file=managed_file
|
||||
)
|
||||
assert _targets(wire.drain()) == [
|
||||
("POST", "/v1/files", f"Bearer {DEPLOYMENT_KEY}"),
|
||||
("POST", "/v1/fine_tuning/jobs", f"Bearer {DEPLOYMENT_KEY}"),
|
||||
]
|
||||
yield Rig(
|
||||
proxy,
|
||||
wire,
|
||||
owner,
|
||||
owner_user,
|
||||
intruder,
|
||||
intruder_user,
|
||||
managed_file,
|
||||
file_id,
|
||||
Job.model_validate_json(created.text).id,
|
||||
created.text,
|
||||
deployment_id,
|
||||
job_id,
|
||||
)
|
||||
assert wire.drain() == ()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
with _rig(
|
||||
tmp_path_factory.mktemp("fine-tuning-managed-ids"), {"require_managed_files": True}, PROVIDER_FILE, PROVIDER_JOB
|
||||
) as built:
|
||||
yield built
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def default_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
with _rig(
|
||||
tmp_path_factory.mktemp("fine-tuning-managed-ids-default"), {}, DEFAULT_PROVIDER_FILE, DEFAULT_PROVIDER_JOB
|
||||
) as built:
|
||||
yield built
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fixture", ["rig", "default_rig"])
|
||||
def test_managed_file_and_job_rows_record_the_owner(fixture: str, request: pytest.FixtureRequest) -> None:
|
||||
built: Final[Rig] = request.getfixturevalue(fixture)
|
||||
file_rows: Final = read_rows(
|
||||
'SELECT model_mappings, flat_model_file_ids, created_by, team_id FROM "LiteLLM_ManagedFileTable" WHERE unified_file_id = %s',
|
||||
(built.managed_file,),
|
||||
)
|
||||
assert file_rows == [
|
||||
{
|
||||
"model_mappings": {built.deployment_id: built.provider_file},
|
||||
"flat_model_file_ids": [built.provider_file],
|
||||
"created_by": built.owner_user,
|
||||
"team_id": None,
|
||||
}
|
||||
], file_rows
|
||||
job_rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT unified_object_id, model_object_id, file_purpose, created_by, team_id FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s',
|
||||
(built.managed_job,),
|
||||
),
|
||||
lambda rows: len(rows) == 1,
|
||||
)
|
||||
assert job_rows == [
|
||||
{
|
||||
"unified_object_id": built.managed_job,
|
||||
"model_object_id": built.provider_job,
|
||||
"file_purpose": "fine-tune",
|
||||
"created_by": built.owner_user,
|
||||
"team_id": None,
|
||||
}
|
||||
], job_rows
|
||||
|
||||
|
||||
def _refused(response: httpx.Response, wire: Wire, detail: str) -> None:
|
||||
reached: Final = _targets(wire.drain())
|
||||
assert response.status_code == 403, response.text
|
||||
assert ProxyError.model_validate_json(response.text).error["message"] == detail, response.text
|
||||
assert reached == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fixture", ["rig", "default_rig"])
|
||||
def test_owner_create_returns_the_unified_job_id(fixture: str, request: pytest.FixtureRequest) -> None:
|
||||
built: Final[Rig] = request.getfixturevalue(fixture)
|
||||
created: Final = Job.model_validate_json(built.created_text)
|
||||
expected: Final = _client_job(
|
||||
_unified_job_id(built.deployment_id, built.provider_job), "validating_files", built.provider_file
|
||||
)
|
||||
assert created.model_dump(exclude={"training_file"}) == expected.model_dump(exclude={"training_file"}), (
|
||||
"training_file is excluded because test_managed_create_sends_the_provider_file_id_of_the_pinned_deployment "
|
||||
"documents the known managed-create bug",
|
||||
built.created_text,
|
||||
)
|
||||
|
||||
|
||||
def test_another_key_cannot_create_with_or_cancel_the_owners_managed_ids(rig: Rig) -> None:
|
||||
created: Final = rig.proxy.request(
|
||||
"POST",
|
||||
"/v1/fine_tuning/jobs",
|
||||
{"model": MANAGED_MODEL, "training_file": rig.managed_file},
|
||||
key=rig.intruder,
|
||||
)
|
||||
_refused(
|
||||
created,
|
||||
rig.wire,
|
||||
f"User {rig.intruder_user} does not have access to the file {rig.managed_file}",
|
||||
)
|
||||
cancelled: Final = rig.proxy.request("POST", f"/v1/fine_tuning/jobs/{rig.managed_job}/cancel", key=rig.intruder)
|
||||
_refused(cancelled, rig.wire, "The caller does not have access to this managed fine-tuning job id.")
|
||||
|
||||
|
||||
def test_another_key_cannot_cancel_the_owners_managed_job_by_default(default_rig: Rig) -> None:
|
||||
cancelled: Final = default_rig.proxy.request(
|
||||
"POST", f"/v1/fine_tuning/jobs/{default_rig.managed_job}/cancel", key=default_rig.intruder
|
||||
)
|
||||
_refused(
|
||||
cancelled,
|
||||
default_rig.wire,
|
||||
f"User {default_rig.intruder_user} does not have access to the object {default_rig.managed_job}",
|
||||
)
|
||||
|
||||
|
||||
def test_another_key_cannot_create_with_the_owners_managed_file_by_default(default_rig: Rig) -> None:
|
||||
pytest.skip(
|
||||
"BUG: another key can create a fine-tuning job with the owner's managed training_file"
|
||||
" when require_managed_files is off"
|
||||
)
|
||||
created: Final = default_rig.proxy.request(
|
||||
"POST",
|
||||
"/v1/fine_tuning/jobs",
|
||||
{"model": MANAGED_MODEL, "training_file": default_rig.managed_file},
|
||||
key=default_rig.intruder,
|
||||
)
|
||||
_refused(
|
||||
created,
|
||||
default_rig.wire,
|
||||
f"User {default_rig.intruder_user} does not have access to the file {default_rig.managed_file}",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("method", "path", "body", "kind"),
|
||||
[
|
||||
pytest.param(
|
||||
"POST",
|
||||
"/v1/fine_tuning/jobs",
|
||||
{"model": "gpt-4o-mini", "training_file": "file-raw123", "custom_llm_provider": "openai"},
|
||||
"file",
|
||||
id="create",
|
||||
),
|
||||
pytest.param("GET", "/v1/fine_tuning/jobs/ftjob-raw123", None, "fine-tuning job", id="retrieve"),
|
||||
pytest.param("POST", "/v1/fine_tuning/jobs/ftjob-raw123/cancel", None, "fine-tuning job", id="cancel"),
|
||||
],
|
||||
)
|
||||
def test_raw_provider_ids_are_refused_when_managed_files_are_required(
|
||||
rig: Rig, method: str, path: str, body: dict[str, JsonValue] | None, kind: str
|
||||
) -> None:
|
||||
response: Final = rig.proxy.request(
|
||||
method, path, body, params={"custom_llm_provider": "openai"} if body is None else None
|
||||
)
|
||||
reached: Final = _targets(rig.wire.drain())
|
||||
assert response.status_code == 400, response.text
|
||||
assert ProxyError.model_validate_json(response.text).error["message"] == (
|
||||
f"Raw provider {kind} ids cannot be used when require_managed_files is enabled in litellm_settings. "
|
||||
f"Use the LiteLLM managed {kind} id returned when the {kind} was created."
|
||||
), response.text
|
||||
assert reached == []
|
||||
|
||||
|
||||
def test_raw_validation_file_is_refused_when_managed_files_are_required(rig: Rig) -> None:
|
||||
rig.wire.drain()
|
||||
response: Final = rig.proxy.request(
|
||||
"POST",
|
||||
"/v1/fine_tuning/jobs",
|
||||
{"model": MANAGED_MODEL, "training_file": rig.managed_file, "validation_file": "file-raw123"},
|
||||
key=rig.owner,
|
||||
)
|
||||
reached: Final = _targets(rig.wire.drain())
|
||||
assert response.status_code == 400, response.text
|
||||
assert ProxyError.model_validate_json(response.text).error["message"] == (
|
||||
"Raw provider file ids cannot be used when require_managed_files is enabled in litellm_settings. "
|
||||
"Use the LiteLLM managed file id returned when the file was created."
|
||||
), response.text
|
||||
assert reached == []
|
||||
|
||||
|
||||
def test_intruder_cannot_use_the_owners_managed_validation_file(rig: Rig) -> None:
|
||||
with rig.sdk(rig.intruder) as client:
|
||||
uploaded: Final = client.files.with_raw_response.create(
|
||||
file=("train.jsonl", TRAINING_LINE, "application/jsonl"),
|
||||
purpose="fine-tune",
|
||||
extra_body={"target_model_names": MANAGED_MODEL},
|
||||
)
|
||||
assert uploaded.status_code == 200, uploaded.text
|
||||
intruder_file: Final = uploaded.parse().id
|
||||
assert _targets(rig.wire.drain()) == [("POST", "/v1/files", f"Bearer {DEPLOYMENT_KEY}")]
|
||||
response: Final = rig.proxy.request(
|
||||
"POST",
|
||||
"/v1/fine_tuning/jobs",
|
||||
{"model": MANAGED_MODEL, "training_file": intruder_file, "validation_file": rig.managed_file},
|
||||
key=rig.intruder,
|
||||
)
|
||||
_refused(
|
||||
response,
|
||||
rig.wire,
|
||||
f"User {rig.intruder_user} does not have access to the file {rig.managed_file}",
|
||||
)
|
||||
|
||||
|
||||
def test_keys_granted_the_retrieve_route_get_the_owner_check_and_the_decoded_job_id(rig: Rig) -> None:
|
||||
with rig.proxy.scenario() as scenario:
|
||||
owner: Final = scenario.key(user_id=rig.owner_user, allowed_routes=["/v1/fine_tuning/jobs"])
|
||||
intruder: Final = scenario.key(user_id=rig.intruder_user, allowed_routes=["/v1/fine_tuning/jobs"])
|
||||
_refused(
|
||||
rig.proxy.request("GET", f"/v1/fine_tuning/jobs/{rig.managed_job}", key=intruder),
|
||||
rig.wire,
|
||||
"The caller does not have access to this managed fine-tuning job id.",
|
||||
)
|
||||
with rig.sdk(owner) as client:
|
||||
retrieved: Final = client.fine_tuning.jobs.with_raw_response.retrieve(rig.managed_job)
|
||||
assert [(r.method, r.target, r.headers["authorization"], r.body) for r in rig.wire.drain()] == [
|
||||
("GET", f"/v1/fine_tuning/jobs/{PROVIDER_JOB}", f"Bearer {DEPLOYMENT_KEY}", b"")
|
||||
]
|
||||
assert retrieved.status_code == 200, retrieved.text
|
||||
assert Job.model_validate_json(retrieved.text) == _client_job(rig.managed_job, "running", rig.provider_file), (
|
||||
retrieved.text
|
||||
)
|
||||
|
||||
|
||||
def test_keys_granted_the_retrieve_route_get_default_job_ownership_and_decoded_id(default_rig: Rig) -> None:
|
||||
with default_rig.proxy.scenario() as scenario:
|
||||
owner: Final = scenario.key(user_id=default_rig.owner_user, allowed_routes=["/v1/fine_tuning/jobs"])
|
||||
intruder: Final = scenario.key(user_id=default_rig.intruder_user, allowed_routes=["/v1/fine_tuning/jobs"])
|
||||
_refused(
|
||||
default_rig.proxy.request("GET", f"/v1/fine_tuning/jobs/{default_rig.managed_job}", key=intruder),
|
||||
default_rig.wire,
|
||||
f"User {default_rig.intruder_user} does not have access to the object {default_rig.managed_job}",
|
||||
)
|
||||
with default_rig.sdk(owner) as client:
|
||||
retrieved: Final = client.fine_tuning.jobs.with_raw_response.retrieve(default_rig.managed_job)
|
||||
assert [
|
||||
(request.method, request.target, request.headers["authorization"], request.body)
|
||||
for request in default_rig.wire.drain()
|
||||
] == [("GET", f"/v1/fine_tuning/jobs/{DEFAULT_PROVIDER_JOB}", f"Bearer {DEPLOYMENT_KEY}", b"")]
|
||||
assert retrieved.status_code == 200, retrieved.text
|
||||
assert Job.model_validate_json(retrieved.text) == _client_job(
|
||||
default_rig.managed_job, "running", default_rig.provider_file
|
||||
), retrieved.text
|
||||
|
||||
|
||||
def test_owner_retrieve_reaches_the_provider_with_the_decoded_job_id(rig: Rig) -> None:
|
||||
pytest.skip("BUG: a virtual key gets 401 on GET /v1/fine_tuning/jobs/{id}, the route is missing from openai_routes")
|
||||
_refused(
|
||||
rig.proxy.request("GET", f"/v1/fine_tuning/jobs/{rig.managed_job}", key=rig.intruder),
|
||||
rig.wire,
|
||||
"The caller does not have access to this managed fine-tuning job id.",
|
||||
)
|
||||
with rig.sdk(rig.owner) as client:
|
||||
retrieved: Final = client.fine_tuning.jobs.with_raw_response.retrieve(rig.managed_job)
|
||||
assert [(r.method, r.target, r.headers["authorization"], r.body) for r in rig.wire.drain()] == [
|
||||
("GET", f"/v1/fine_tuning/jobs/{PROVIDER_JOB}", f"Bearer {DEPLOYMENT_KEY}", b"")
|
||||
]
|
||||
assert retrieved.status_code == 200, retrieved.text
|
||||
assert Job.model_validate_json(retrieved.text) == _client_job(rig.managed_job, "running", rig.provider_file), (
|
||||
retrieved.text
|
||||
)
|
||||
|
||||
|
||||
def test_owner_cancel_reaches_the_provider_with_the_decoded_job_id(rig: Rig) -> None:
|
||||
with rig.sdk(rig.owner) as client:
|
||||
cancelled: Final = client.fine_tuning.jobs.with_raw_response.cancel(rig.managed_job)
|
||||
assert [(r.method, r.target, r.headers["authorization"], r.body) for r in rig.wire.drain()] == [
|
||||
("POST", f"/v1/fine_tuning/jobs/{PROVIDER_JOB}/cancel", f"Bearer {DEPLOYMENT_KEY}", b"")
|
||||
]
|
||||
assert cancelled.status_code == 200, cancelled.text
|
||||
assert Job.model_validate_json(cancelled.text) == _client_job(rig.managed_job, "cancelled", rig.provider_file), (
|
||||
cancelled.text
|
||||
)
|
||||
Loading…
Add table
Reference in a new issue