Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 9 additions & 4 deletions python/kthena/downloader/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,15 +23,20 @@
logger = setup_logger()


def load_config(config_str: str = None) -> dict:
def load_config(raw_config: str | dict | None = None) -> dict:
"""Build the downloader config from a JSON string (CLI) or a dict (runtime API)."""
config = {}
if config_str:
if isinstance(raw_config, dict):
config = dict(raw_config)
elif raw_config:
try:
config = json.loads(config_str)
config = json.loads(raw_config)
logger.info("Loaded configuration from JSON string.")
except json.JSONDecodeError as e:
except (json.JSONDecodeError, TypeError) as e:
logger.error(f"Error parsing configuration JSON: {e}")
raise ValueError("Invalid configuration JSON format.") from e
if not isinstance(config, dict):
raise ValueError("Configuration must be a JSON object.")

env_config = {
"hf_token": os.getenv("HF_AUTH_TOKEN"),
Expand Down
75 changes: 75 additions & 0 deletions python/kthena/tests/test_runtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import uvicorn
from fastapi.testclient import TestClient

from kthena.downloader.app import load_config
from kthena.runtime import app as runtime_app
from kthena.runtime import memory_kv_manager as mkm
from kthena.runtime.app import (
Expand Down Expand Up @@ -170,6 +171,80 @@ def test_load_lora_adapter_async_download_success(monkeypatch):
assert resp.status_code == 202


@pytest.mark.parametrize("path, extra_body", [
("/v1/download_model", {}),
("/v1/load_lora_adapter", {"lora_name": "lora-sql"}),
])
def test_download_endpoints_pass_object_config_to_downloader(monkeypatch, path, extra_body):
# Arrange: record the config the endpoint hands to the downloader
monkeypatch.delenv("HF_AUTH_TOKEN", raising=False)
app = create_application(_make_args())
state = get_app_state(app)
monkeypatch.setattr("httpx.AsyncClient", _FakeAsyncClient)
state.client = _FakeAsyncClient()
monkeypatch.setattr("httpx.Client", _FakeSyncClient)

received = []
monkeypatch.setattr(
"kthena.runtime.app.download_model",
lambda source, output_dir, config, max_workers: received.append(config),
)

with TestClient(app) as client:
body = {
"source": "hf://org/repo",
"output_dir": "/tmp/model",
"config": {"hf_token": "from-request"},
**extra_body,
}

# Act
resp = client.post(path, json=body)

# Assert
assert resp.status_code == 200
assert received == [{"hf_token": "from-request"}]


@pytest.mark.parametrize("config, expected", [
(None, {}),
("", {}),
({}, {}),
({"hf_token": "from-config"}, {"hf_token": "from-config"}),
('{"hf_token": "from-config"}', {"hf_token": "from-config"}),
])
def test_load_config_accepts_supported_forms(monkeypatch, config, expected):
monkeypatch.delenv("HF_AUTH_TOKEN", raising=False)

assert load_config(config) == expected


@pytest.mark.parametrize("config", [
{"hf_token": "from-config"},
'{"hf_token": "from-config"}',
])
def test_load_config_prefers_explicit_values_over_environment(monkeypatch, config):
monkeypatch.setenv("HF_AUTH_TOKEN", "from-env")

assert load_config(config) == {"hf_token": "from-config"}


def test_load_config_fills_missing_values_without_mutating_input(monkeypatch):
monkeypatch.setenv("HF_AUTH_TOKEN", "from-env")
config = {"endpoint": "http://storage.local"}

result = load_config(config)

assert result == {"endpoint": "http://storage.local", "hf_token": "from-env"}
assert config == {"endpoint": "http://storage.local"}


@pytest.mark.parametrize("config", [123, ["hf_token"], "[1, 2]", "not json"])
def test_load_config_rejects_values_that_are_not_objects(config):
with pytest.raises(ValueError):
load_config(config)


@pytest.mark.asyncio
@pytest.mark.parametrize("mode, body, expected_status", [
("redis", {"router_id": "router-a"}, 409),
Expand Down
Loading