Skip to content

Commit

Permalink
Convert async _close_models in LlamaCppInferenceEngine
Browse files Browse the repository at this point in the history
Fixes RuntimeWarning about unawaited coroutine by converting
the async cleanup method to synchronous since it's called from
del. Also fixes dictionary iteration and adjusts method
name to follow Python protected method conventions.

Resolves: #109
  • Loading branch information
lukehinds committed Dec 2, 2024
1 parent 0bc3c03 commit fbabe31
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 11 deletions.
20 changes: 10 additions & 10 deletions src/codegate/inference/inference_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,16 @@ def __init__(self):
self.__models = {}

def __del__(self):
self.__close_models()
self._close_models()

def _close_models(self):
"""
Closes all open models and samplers
"""
for _, model in self.__models.items():
if model._sampler:
model._sampler.close()
model.close()

async def __get_model(self, model_path, embedding=False, n_ctx=512, n_gpu_layers=0):
"""
Expand Down Expand Up @@ -70,12 +79,3 @@ async def embed(self, model_path, content):
"""
model = await self.__get_model(model_path=model_path, embedding=True)
return model.embed(content)

async def __close_models(self):
"""
Closes all open models and samplers
"""
for _, model in self.__models:
if model._sampler:
model._sampler.close()
model.close()
2 changes: 1 addition & 1 deletion src/codegate/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,8 @@
from codegate.pipeline.codegate_context_retriever.codegate import CodegateContextRetriever
from codegate.pipeline.codegate_system_prompt.codegate import CodegateSystemPrompt
from codegate.pipeline.extract_snippets.extract_snippets import CodeSnippetExtractor
from codegate.pipeline.secrets.signatures import CodegateSignatures
from codegate.pipeline.secrets.secrets import CodegateSecrets
from codegate.pipeline.secrets.signatures import CodegateSignatures
from codegate.pipeline.version.version import CodegateVersion
from codegate.providers.anthropic.provider import AnthropicProvider
from codegate.providers.llamacpp.provider import LlamaCppProvider
Expand Down

0 comments on commit fbabe31

Please sign in to comment.