Skip to content

Commit c2c40d7

Browse files
Fix asyncpg event loop bug with startup wrapper monkey-patch
The first request to llama-stack after a pod start fails with RuntimeError because asyncpg connections created during StackApp init are bound to a temporary event loop. This adds a Python wrapper script that monkey-patches StackApp.__init__ to lazily recreate SQL engines on uvicorn's event loop, matching the upstream fix (PR #5837) that hasn't been built into the container image yet. Remove when the container image includes the upstream fix.
1 parent e2975e3 commit c2c40d7

4 files changed

Lines changed: 63 additions & 1 deletion

File tree

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
#!/usr/bin/env python3
2+
# Workaround for asyncpg event loop bug (llama-stack #5978).
3+
# Monkey-patches StackApp to reset SQL engines after temp event loop init.
4+
# Remove when the container image includes the upstream fix (PR #5837).
5+
6+
import gc
7+
import sys
8+
9+
import llama_stack.core.server.server as server_mod
10+
from llama_stack.core.storage.sqlstore.sqlalchemy_sqlstore import SqlAlchemySqlStoreImpl
11+
12+
_orig_init = server_mod.StackApp.__init__
13+
14+
15+
class SessionMaker:
16+
"""Recreates the SQLAlchemy engine and session on the current event loop."""
17+
18+
def __init__(self, store):
19+
self._store = store
20+
self._maker = None
21+
22+
def __call__(self):
23+
if self._maker is None:
24+
from sqlalchemy.ext.asyncio import async_sessionmaker
25+
26+
self._store._engine = self._store.create_engine()
27+
self._maker = async_sessionmaker(self._store._engine)
28+
self._store.async_session = self._maker
29+
return self._maker()
30+
31+
32+
def _patched_init(self, config, *args, **kwargs):
33+
_orig_init(self, config, *args, **kwargs)
34+
for obj in gc.get_objects():
35+
if isinstance(obj, SqlAlchemySqlStoreImpl):
36+
obj._engine = None
37+
obj.async_session = SessionMaker(obj)
38+
39+
40+
server_mod.StackApp.__init__ = _patched_init
41+
42+
from llama_stack.cli.llama import main # noqa: E402
43+
44+
sys.argv[0] = "llama"
45+
main()

internal/controller/constants.go

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,6 +235,10 @@ const (
235235
// script is stored in the ConfigMap containing vector database init scripts.
236236
VectorDBBuildScriptKey = "vector_database_build.py"
237237

238+
// LlamaStartupWrapperKey is the ConfigMap key for the startup wrapper script
239+
// that monkey-patches the asyncpg event loop bug fix. Remove with PR #5837 backport.
240+
LlamaStartupWrapperKey = "llama_startup_wrapper.py"
241+
238242
// Resource Version Annotation
239243
// These constants define annotation keys used to track the resource versions of specific ConfigMaps.
240244
// By recording the resource version of a ConfigMap in a Deployment, StatefulSet, or similar resource,
@@ -350,6 +354,13 @@ var vectorDatabaseCollectScript string
350354
//go:embed assets/vector_database_build.py
351355
var vectorDatabaseBuildScript string
352356

357+
// llamaStartupWrapperScript is a Python monkey-patch that fixes the asyncpg
358+
// event loop bug (ogx-ai/ogx#5978) by resetting SQL engines after StackApp
359+
// initialization. Remove when the container image includes upstream PR #5837.
360+
//
361+
//go:embed assets/llama_startup_wrapper.py
362+
var llamaStartupWrapperScript string
363+
353364
//go:embed assets/console_nginx.conf.tmpl
354365
var consoleNginxConfigTemplate string
355366

internal/controller/lcore_deployment.go

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,11 +63,16 @@ func buildLCorePodTemplateSpec(h *common_helper.Helper, ctx context.Context, ins
6363
llamaStackMounts := []corev1.VolumeMount{}
6464
llamaStackMounts = append(llamaStackMounts, sharedMounts...)
6565
llamaStackMounts = append(llamaStackMounts, llamaCacheMounts...)
66+
llamaStackMounts = append(llamaStackMounts, corev1.VolumeMount{
67+
Name: VectorDBScriptsVolumeName,
68+
MountPath: VectorDBScriptsMountPath,
69+
ReadOnly: true,
70+
})
6671

6772
llamaStackContainer := corev1.Container{
6873
Name: "llama-stack",
6974
Image: apiv1beta1.OpenStackLightspeedDefaultValues.LCoreImageURL,
70-
Command: []string{"llama", "stack", "run", VectorDBVolumeOGXConfigPath},
75+
Command: []string{"python3", VectorDBScriptsMountPath + "/" + LlamaStartupWrapperKey, "stack", "run", VectorDBVolumeOGXConfigPath},
7176
Ports: []corev1.ContainerPort{{Name: "llama-stack", ContainerPort: LlamaStackContainerPort}},
7277
VolumeMounts: llamaStackMounts,
7378
Env: llamaEnvVars,

internal/controller/lcore_reconciler.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -297,6 +297,7 @@ func reconcileVectorDBScriptsConfigMap(h *common_helper.Helper, ctx context.Cont
297297
cm.Data = map[string]string{
298298
VectorDBCollectScriptKey: vectorDatabaseCollectScript,
299299
VectorDBBuildScriptKey: vectorDatabaseBuildScript,
300+
LlamaStartupWrapperKey: llamaStartupWrapperScript,
300301
}
301302

302303
return controllerutil.SetControllerReference(h.GetBeforeObject(), cm, h.GetScheme())

0 commit comments

Comments
 (0)