From f43c721579d17f1ceee0313bb463e9a62a384e29 Mon Sep 17 00:00:00 2001
From: aarnphm-ec2-dev <29749331+aarnphm@users.noreply.github.com>
Date: Tue, 5 Sep 2023 01:08:23 +0000
Subject: [PATCH] chore: only add bentomodel branch during generated service
 with OpenLLM

Signed-off-by: aarnphm-ec2-dev <29749331+aarnphm@users.noreply.github.com>
---
 openllm-python/src/openllm/_llm.py | 11 +++++++----
 wheels.sh                          | 27 +++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 4 deletions(-)
 create mode 100644 wheels.sh

diff --git a/openllm-python/src/openllm/_llm.py b/openllm-python/src/openllm/_llm.py
index 7e0ad43f..08c77cdb 100644
--- a/openllm-python/src/openllm/_llm.py
+++ b/openllm-python/src/openllm/_llm.py
@@ -892,10 +892,13 @@ class LLM(LLMInterface[M, T], ReprMixin):
     '''
     models = models if models is not None else []
 
-    try:
-      models.append(self._bentomodel)
-    except bentoml.exceptions.NotFound as err:
-      raise RuntimeError(f'Failed to locate {self._bentomodel}:{err}') from None
+    if os.environ.get('BENTO_PATH') is None:
+      # Hmm we should only add this if it is not in the container environment
+      # BentoML sets BENTO_PATH so we can use this as switch logic here.
+      try:
+        models.append(self._bentomodel)
+      except bentoml.exceptions.NotFound as err:
+        raise RuntimeError(f'Failed to locate {self._bentomodel}:{err}') from None
 
     generate_sig = ModelSignature.from_dict(ModelSignatureDict(batchable=False))
     embeddings_sig = ModelSignature.from_dict(ModelSignatureDict(batchable=True, batch_dim=0))
diff --git a/wheels.sh b/wheels.sh
new file mode 100644
index 00000000..65581f79
--- /dev/null
+++ b/wheels.sh
@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+
+
+set -eo pipefail
+
+GIT_ROOT="$(git rev-parse --show-toplevel)"
+cd "$GIT_ROOT" || exit 1
+
+mkdir -p dist
+
+pushd openllm-client &>/dev/null
+
+python -m build -w && mv dist/* ../dist
+
+popd &>/dev/null
+
+pushd openllm-core &>/dev/null
+
+python -m build -w && mv dist/* ../dist
+
+popd &>/dev/null
+
+pushd openllm-python &>/dev/null
+
+python -m build -w && mv dist/* ../dist
+
+popd &>/dev/null