Merge pull request #23 from omkarjoshi0304/fix/asyncpg-event-loop-workaround

Fix asyncpg event loop bug with startup wrapper patch

Author: openshift-merge-bot[bot]

internal/controller/assets/llama_startup_wrapper.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+# Workaround for asyncpg event loop bug (llama-stack #5978).
+# Monkey-patches StackApp to reset SQL engines after temp event loop init.
+# Remove when the container image includes the upstream fix (PR #5837).
+
+import gc
+import sys
+
+import llama_stack.core.server.server as server_mod
+from llama_stack.core.storage.sqlstore.sqlalchemy_sqlstore import SqlAlchemySqlStoreImpl
+from sqlalchemy.ext.asyncio import async_sessionmaker
+
+_orig_init = server_mod.StackApp.__init__
+
+
+class SessionMaker:
+    """Recreates the SQLAlchemy engine and session on the current event loop."""
+
+    def __init__(self, store):
+        self._store = store
+        self._maker = None
+
+    def __call__(self):
+        if self._maker is None:
+            self._store._engine = self._store.create_engine()
+            self._maker = async_sessionmaker(self._store._engine)
+            self._store.async_session = self._maker
+        return self._maker()
+
+
+def _patched_init(self, config, *args, **kwargs):
+    _orig_init(self, config, *args, **kwargs)
+    for obj in gc.get_objects():
+        if isinstance(obj, SqlAlchemySqlStoreImpl):
+            obj._engine = None
+            obj.async_session = SessionMaker(obj)
+
+
+server_mod.StackApp.__init__ = _patched_init
+
+from llama_stack.cli.llama import main  # noqa: E402
+
+sys.argv[0] = "llama"
+main()

internal/controller/constants.go

@@ -235,6 +235,10 @@ const (
 	// script is stored in the ConfigMap containing vector database init scripts.
 	VectorDBBuildScriptKey = "vector_database_build.py"
 
+	// LlamaStartupWrapperKey is the ConfigMap key for the startup wrapper script
+	// that monkey-patches the asyncpg event loop bug fix. Remove with PR #5837 backport.
+	LlamaStartupWrapperKey = "llama_startup_wrapper.py" // #nosec G101 -- ConfigMap key, not a credential
+
 	// Resource Version Annotation
 	// These constants define annotation keys used to track the resource versions of specific ConfigMaps.
 	// By recording the resource version of a ConfigMap in a Deployment, StatefulSet, or similar resource,
@@ -350,6 +354,13 @@ var vectorDatabaseCollectScript string
 //go:embed assets/vector_database_build.py
 var vectorDatabaseBuildScript string
 
+// llamaStartupWrapperScript is a Python monkey-patch that fixes the asyncpg
+// event loop bug (ogx-ai/ogx#5978) by resetting SQL engines after StackApp
+// initialization. Remove when the container image includes upstream PR #5837.
+//
+//go:embed assets/llama_startup_wrapper.py
+var llamaStartupWrapperScript string
+
 //go:embed assets/console_nginx.conf.tmpl
 var consoleNginxConfigTemplate string
 

internal/controller/lcore_deployment.go

@@ -63,11 +63,16 @@ func buildLCorePodTemplateSpec(h *common_helper.Helper, ctx context.Context, ins
 	llamaStackMounts := []corev1.VolumeMount{}
 	llamaStackMounts = append(llamaStackMounts, sharedMounts...)
 	llamaStackMounts = append(llamaStackMounts, llamaCacheMounts...)
+	llamaStackMounts = append(llamaStackMounts, corev1.VolumeMount{
+		Name:      VectorDBScriptsVolumeName,
+		MountPath: VectorDBScriptsMountPath,
+		ReadOnly:  true,
+	})
 
 	llamaStackContainer := corev1.Container{
 		Name:         "llama-stack",
 		Image:        apiv1beta1.OpenStackLightspeedDefaultValues.LCoreImageURL,
-		Command:      []string{"llama", "stack", "run", VectorDBVolumeOGXConfigPath},
+		Command:      []string{"python3", VectorDBScriptsMountPath + "/" + LlamaStartupWrapperKey, "stack", "run", VectorDBVolumeOGXConfigPath},
 		Ports:        []corev1.ContainerPort{{Name: "llama-stack", ContainerPort: LlamaStackContainerPort}},
 		VolumeMounts: llamaStackMounts,
 		Env:          llamaEnvVars,

internal/controller/lcore_reconciler.go

@@ -297,6 +297,7 @@ func reconcileVectorDBScriptsConfigMap(h *common_helper.Helper, ctx context.Cont
 		cm.Data = map[string]string{
 			VectorDBCollectScriptKey: vectorDatabaseCollectScript,
 			VectorDBBuildScriptKey:   vectorDatabaseBuildScript,
+			LlamaStartupWrapperKey:   llamaStartupWrapperScript,
 		}
 
 		return controllerutil.SetControllerReference(h.GetBeforeObject(), cm, h.GetScheme())

test/kuttl/common/openstack-lightspeed-instance/assert-openstack-lightspeed-instance.yaml

@@ -256,6 +256,9 @@ spec:
               mountPath: /vector-db-discovered-values
             - name: llama-cache
               mountPath: /tmp/llama-stack
+            - name: vector-db-scripts
+              mountPath: /scripts
+              readOnly: true
         - name: lightspeed-service-api
           resources:
             requests:

test/kuttl/tests/update-openstacklightspeed/08-assert-openstacklightspeed-update.yaml

@@ -138,6 +138,9 @@ spec:
               mountPath: /vector-db-discovered-values
             - name: llama-cache
               mountPath: /tmp/llama-stack
+            - name: vector-db-scripts
+              mountPath: /scripts
+              readOnly: true
         - name: lightspeed-service-api
           resources:
             requests: