Skip to content

Commit 53792b0

Browse files
lapylapy
committed
Add llama-swap environment variable support and command string generation
- Introduced new environment variables prefixed with `LLAMA_STUDIO_` for managing paths in GGUF commands. - Updated `_llama_swap_env_list` to ignore `LLAMA_STUDIO_` keys in user environment and incorporate a new `studio_env` parameter. - Created `_build_llama_swap_cmd_string` to generate command strings for llama-server using environment variables instead of direct paths. - Refactored `_build_llama_command` to utilize the new command string generation and environment handling, improving overall configuration management. - Enhanced documentation within the code to clarify the purpose of new functions and environment variable handling.
1 parent 5e9947c commit 53792b0

2 files changed

Lines changed: 107 additions & 20 deletions

File tree

‎backend/llama_swap_config.py‎

Lines changed: 63 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,13 @@
3131

3232
_SWAP_ENV_KEY_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
3333

34+
# llama-swap ``env`` keys for GGUF commands (paths not embedded in ``cd`` / argv).
35+
_STUDIO_ENV_SERVER_CWD = "LLAMA_STUDIO_SERVER_CWD"
36+
_STUDIO_ENV_MODEL_PATH = "LLAMA_STUDIO_MODEL_PATH"
37+
_STUDIO_ENV_HF_REPO = "LLAMA_STUDIO_HF_REPO"
38+
_STUDIO_ENV_MMPROJ_PATH = "LLAMA_STUDIO_MMPROJ_PATH"
39+
_STUDIO_ENV_PREFIX = "LLAMA_STUDIO_"
40+
3441

3542
def clear_supported_flags_cache() -> None:
3643
_supported_flags_cache.clear()
@@ -275,18 +282,28 @@ def _llama_swap_env_list(
275282
*,
276283
resolved_ld_library_path: str,
277284
user_env: Dict[str, str],
285+
studio_env: Optional[Dict[str, str]] = None,
278286
) -> List[str]:
279287
"""
280288
Build llama-swap ``env`` entries (``NAME=value`` strings).
281289
User ``LD_LIBRARY_PATH`` is appended after the resolved CUDA/build path.
290+
Keys prefixed with ``LLAMA_STUDIO_`` in ``user_env`` are ignored (reserved for
291+
generated server cwd / model / projector paths).
292+
``studio_env`` is applied last so generated paths win.
282293
"""
283-
merged = dict(user_env)
294+
merged = {
295+
k: v
296+
for k, v in user_env.items()
297+
if not str(k).startswith(_STUDIO_ENV_PREFIX)
298+
}
284299
user_ld = merged.pop("LD_LIBRARY_PATH", None)
285300
ld = resolved_ld_library_path
286301
if user_ld and str(user_ld).strip():
287302
suffix = str(user_ld).strip()
288303
ld = f"{ld}:{suffix}" if ld else suffix
289304
merged["LD_LIBRARY_PATH"] = ld
305+
if studio_env:
306+
merged.update(studio_env)
290307
return [f"{k}={merged[k]}" for k in sorted(merged.keys())]
291308

292309

@@ -476,6 +493,32 @@ def _resolve_mmproj_path(
476493
return None
477494

478495

496+
def _build_llama_swap_cmd_string(
497+
*,
498+
binary_name: str,
499+
proxy_model_name: str,
500+
hf_repo_arg: Optional[str],
501+
mmproj_path: Optional[str],
502+
structured_argv: List[str],
503+
) -> str:
504+
"""
505+
``bash -c`` body for llama-server: paths come from env (``$LLAMA_STUDIO_*``), not ``cd``.
506+
"""
507+
exec_tok = f'"${{{_STUDIO_ENV_SERVER_CWD}}}"/{binary_name}'
508+
parts: List[str] = [exec_tok]
509+
if hf_repo_arg:
510+
parts.extend(["--hf-repo", f'"${{{_STUDIO_ENV_HF_REPO}}}"'])
511+
else:
512+
parts.extend(["--model", f'"${{{_STUDIO_ENV_MODEL_PATH}}}"'])
513+
parts.extend(["--port", "${PORT}", "--alias", _quote_shell_token(proxy_model_name)])
514+
if mmproj_path:
515+
parts.extend(["--mmproj", f'"${{{_STUDIO_ENV_MMPROJ_PATH}}}"'])
516+
if structured_argv:
517+
parts.append(_shell_join(structured_argv))
518+
inner = " ".join(parts)
519+
return f"bash -c {shlex.quote(inner)}"
520+
521+
479522
def _build_llama_command(
480523
*,
481524
model: Any,
@@ -497,32 +540,36 @@ def _build_llama_command(
497540
_, work_cwd = resolve_llama_server_invocation_paths(llama_server_path)
498541
binary_name = os.path.basename(llama_server_path)
499542
library_path = _resolve_cuda_library_path(work_cwd)
500-
501-
argv: List[str] = [f"./{binary_name}"]
502-
if hf_repo_arg:
503-
argv.extend(["--hf-repo", hf_repo_arg])
504-
else:
505-
argv.extend(["--model", str(model_path)])
506-
507-
argv.extend(["--port", "${PORT}", "--alias", proxy_model_name])
508543
mmproj_path = _resolve_mmproj_path(model, hf_id, hf_repo_arg)
509-
if mmproj_path:
510-
argv.extend(["--mmproj", mmproj_path])
511544

512545
structured_engine = (
513546
engine_for_params
514547
if engine_for_params in ("llama_cpp", "ik_llama")
515548
else infer_engine_id_for_binary(llama_server_path)
516549
)
517-
argv.extend(
518-
_emit_structured_tokens(
519-
config, engine=structured_engine, param_index=param_index
520-
)
550+
structured_argv = _emit_structured_tokens(
551+
config, engine=structured_engine, param_index=param_index
552+
)
553+
554+
studio_env: Dict[str, str] = {_STUDIO_ENV_SERVER_CWD: work_cwd}
555+
if hf_repo_arg:
556+
studio_env[_STUDIO_ENV_HF_REPO] = str(hf_repo_arg)
557+
else:
558+
studio_env[_STUDIO_ENV_MODEL_PATH] = str(model_path)
559+
if mmproj_path:
560+
studio_env[_STUDIO_ENV_MMPROJ_PATH] = str(mmproj_path)
561+
562+
cmd = _build_llama_swap_cmd_string(
563+
binary_name=binary_name,
564+
proxy_model_name=proxy_model_name,
565+
hf_repo_arg=hf_repo_arg,
566+
mmproj_path=mmproj_path,
567+
structured_argv=structured_argv,
521568
)
522-
cmd = _render_bash_command(argv, cwd=work_cwd, env=None)
523569
env_list = _llama_swap_env_list(
524570
resolved_ld_library_path=library_path,
525571
user_env=_normalize_swap_env(config),
572+
studio_env=studio_env,
526573
)
527574
return cmd, env_list
528575

‎backend/tests/test_llama_swap_config.py‎

Lines changed: 44 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -300,6 +300,26 @@ def get_active_engine_version(self, engine):
300300
assert llama_swap_config.is_ik_llama_cpp("/tmp/llama") is False
301301

302302

303+
def test_llama_swap_env_list_drops_user_llama_studio_keys_and_merges_studio_env():
304+
user = llama_swap_config._normalize_swap_env(
305+
{"swap_env": {"LLAMA_STUDIO_MODEL_PATH": "/evil.gguf", "FOO": "bar"}}
306+
)
307+
out = llama_swap_config._llama_swap_env_list(
308+
resolved_ld_library_path="/lib",
309+
user_env=user,
310+
studio_env={
311+
llama_swap_config._STUDIO_ENV_MODEL_PATH: "/real.gguf",
312+
llama_swap_config._STUDIO_ENV_SERVER_CWD: "/w",
313+
},
314+
)
315+
assert out == [
316+
"FOO=bar",
317+
"LD_LIBRARY_PATH=/lib",
318+
"LLAMA_STUDIO_MODEL_PATH=/real.gguf",
319+
"LLAMA_STUDIO_SERVER_CWD=/w",
320+
]
321+
322+
303323
def test_normalize_swap_env_and_llama_swap_env_list():
304324
assert llama_swap_config._normalize_swap_env(None) == {}
305325
assert llama_swap_config._normalize_swap_env({}) == {}
@@ -514,13 +534,21 @@ def test_preview_llama_swap_command_uses_catalog_metadata(monkeypatch, tmp_path)
514534

515535
assert preview["ok"] is True
516536
assert "--model" in preview["cmd"]
537+
assert "${LLAMA_STUDIO_MODEL_PATH}" in preview["cmd"]
538+
assert "${LLAMA_STUDIO_SERVER_CWD}" in preview["cmd"]
539+
assert "cd " not in preview["cmd"]
517540
assert "--alias org-model.q4_k_m" in preview["cmd"]
518541
assert "--jinja" in preview["cmd"]
519542
assert "--temperature 0.7" in preview["cmd"]
520543
assert "--stop" in preview["cmd"]
521544
assert "END HERE" in preview["cmd"]
522545
assert preview["use_model_name"] is None
523-
assert preview["env"] == ["LD_LIBRARY_PATH=/fake/lib"]
546+
env_lines = sorted(preview["env"] or [])
547+
assert env_lines == [
548+
"LD_LIBRARY_PATH=/fake/lib",
549+
f"LLAMA_STUDIO_MODEL_PATH={model_path}",
550+
f"LLAMA_STUDIO_SERVER_CWD={tmp_path}",
551+
]
524552
assert " env " not in preview["cmd"]
525553

526554

@@ -680,7 +708,12 @@ def fake_param_index(engine):
680708

681709
assert set(doc["models"].keys()) == {"org-model.q4_k_m", "org-repo-model"}
682710
assert "--temperature 0.9" in doc["models"]["org-model.q4_k_m"]["cmd"]
683-
assert doc["models"]["org-model.q4_k_m"]["env"] == ["LD_LIBRARY_PATH=/fake/lib"]
711+
gguf_env = sorted(doc["models"]["org-model.q4_k_m"]["env"])
712+
assert gguf_env == [
713+
"LD_LIBRARY_PATH=/fake/lib",
714+
f"LLAMA_STUDIO_MODEL_PATH={str(model_path)}",
715+
f"LLAMA_STUDIO_SERVER_CWD={str(tmp_path)}",
716+
]
684717
assert (
685718
"serve api_server org/repo-model --server-port ${PORT} --tp 2"
686719
in doc["models"]["org-repo-model"]["cmd"]
@@ -779,8 +812,15 @@ def inv_paths(path):
779812
)
780813
doc = json.loads(json.dumps(llama_swap_config.yaml.safe_load(yaml_str)))
781814
cmd = doc["models"]["org-model.q4_k_m"]["cmd"]
782-
assert "./ik-server" in cmd
783-
assert "ik-build" in cmd
815+
ik_build = str(tmp_path / "ik-build")
816+
assert '"${LLAMA_STUDIO_SERVER_CWD}"/ik-server' in cmd
817+
assert "${LLAMA_STUDIO_MODEL_PATH}" in cmd
818+
assert "cd " not in cmd
819+
assert sorted(doc["models"]["org-model.q4_k_m"]["env"]) == [
820+
"LD_LIBRARY_PATH=/fake/lib",
821+
f"LLAMA_STUDIO_MODEL_PATH={str(model_path)}",
822+
f"LLAMA_STUDIO_SERVER_CWD={ik_build}",
823+
]
784824
assert "./llama-server" not in cmd
785825

786826

0 commit comments

Comments
 (0)