-
Notifications
You must be signed in to change notification settings - Fork 70
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add setup_environment
to truss
#1188
Open
spal1
wants to merge
26
commits into
main
Choose a base branch
from
samiksha/bt-12427-add-support-for-setup_environment-to-truss
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
Show all changes
26 commits
Select commit
Hold shift + click to select a range
114c0f9
Bump Version for CTX builder
spal1 e3d97d1
skip long test
spal1 b2ea022
add aiofiles to requirements
spal1 8a6715d
increase ctx builder version
spal1 3c0d4cb
use async task instead
spal1 97fd7a8
increase ctx builder
spal1 89ffea5
pass event_loop down
spal1 e54e739
allow event_loop to be optional
spal1 84d5339
only thread self._model.load
spal1 ebea840
revert test_server changes
spal1 b929f6e
add integration test for setup_environment
spal1 8898caf
add truss for setup_environment integration test
spal1 5ae6f34
revert model wrapper test changes
spal1 e0b2c31
revert load changes and check model wrapper status instead + tests
spal1 16a6639
clean up resolver
spal1 e30b5c4
fix resolver test
spal1 336fd95
call setup env before load
spal1 3bae49a
remove print statements
spal1 23aa0e0
unit tests + fixes
spal1 9e0c6ec
sid CR
spal1 235bd53
fix polling structure
spal1 c6013e8
CR
spal1 df9d61b
bump rc version
spal1 fd5d479
support for clearing out env
spal1 1d0e9cb
update ctx builder and revert version of types-aiofiles
spal1 7e283e4
test for Nones in configmap file
spal1 File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
Large diffs are not rendered by default.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -4,6 +4,7 @@ | |
import importlib | ||
import importlib.util | ||
import inspect | ||
import json | ||
import logging | ||
import os | ||
import pathlib | ||
|
@@ -35,7 +36,7 @@ | |
from common.schema import TrussSchema | ||
from opentelemetry import trace | ||
from pydantic import BaseModel | ||
from shared import serialization | ||
from shared import dynamic_config_resolver, serialization | ||
from shared.lazy_data_resolver import LazyDataResolver | ||
from shared.secrets_resolver import SecretsResolver | ||
|
||
|
@@ -53,6 +54,7 @@ | |
EXTENSION_CLASS_NAME = "Extension" | ||
EXTENSION_FILE_NAME = "extension" | ||
TRT_LLM_EXTENSION_NAME = "trt_llm" | ||
POLL_FOR_ENVIRONMENT_UPDATES_TIMEOUT_SECS = 30 | ||
|
||
|
||
@asynccontextmanager | ||
|
@@ -191,6 +193,7 @@ class ModelDescriptor: | |
predict: MethodDescriptor | ||
postprocess: Optional[MethodDescriptor] | ||
truss_schema: Optional[TrussSchema] | ||
setup_environment: Optional[MethodDescriptor] | ||
|
||
@cached_property | ||
def skip_input_parsing(self) -> bool: | ||
|
@@ -243,11 +246,19 @@ def from_model(cls, model) -> "ModelDescriptor": | |
else: | ||
return_annotation = inspect.signature(model.predict).return_annotation | ||
|
||
if hasattr(model, "setup_environment"): | ||
setup_environment = MethodDescriptor.from_method( | ||
model.setup_environment, "setup_environment" | ||
) | ||
else: | ||
setup_environment = None | ||
|
||
return cls( | ||
preprocess=preprocess, | ||
predict=predict, | ||
postprocess=postprocess, | ||
truss_schema=TrussSchema.from_signature(parameters, return_annotation), | ||
setup_environment=setup_environment, | ||
) | ||
|
||
|
||
|
@@ -259,6 +270,8 @@ class ModelWrapper: | |
_logger: logging.Logger | ||
_status: "ModelWrapper.Status" | ||
_predict_semaphore: Semaphore | ||
_poll_for_environment_updates_task: Optional[asyncio.Task] | ||
_environment: Optional[dict] | ||
|
||
class Status(Enum): | ||
NOT_READY = 0 | ||
|
@@ -280,6 +293,8 @@ def __init__(self, config: Dict, tracer: sdk_trace.Tracer): | |
"predict_concurrency", DEFAULT_PREDICT_CONCURRENCY | ||
) | ||
) | ||
self._poll_for_environment_updates_task = None | ||
self._environment = None | ||
|
||
@property | ||
def _model(self) -> Any: | ||
|
@@ -419,6 +434,9 @@ def _load_impl(self): | |
|
||
self._maybe_model_descriptor = ModelDescriptor.from_model(self._model) | ||
|
||
if self._maybe_model_descriptor.setup_environment: | ||
self._initialize_environment_before_load() | ||
|
||
if hasattr(self._model, "load"): | ||
retry( | ||
spal1 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
self._model.load, | ||
|
@@ -428,6 +446,71 @@ def _load_impl(self): | |
gap_seconds=1.0, | ||
) | ||
|
||
def setup_polling_for_environment_updates(self): | ||
self._poll_for_environment_updates_task = asyncio.create_task( | ||
self.poll_for_environment_updates() | ||
) | ||
|
||
def _initialize_environment_before_load(self): | ||
environment_str = dynamic_config_resolver.get_dynamic_config_value_sync( | ||
dynamic_config_resolver.ENVIRONMENT_DYNAMIC_CONFIG_KEY | ||
) | ||
if environment_str: | ||
environment_json = json.loads(environment_str) | ||
self._model.setup_environment(environment_json) | ||
self._environment = environment_json | ||
|
||
async def setup_environment(self, environment: Optional[dict]): | ||
descriptor = self.model_descriptor.setup_environment | ||
if not descriptor: | ||
return | ||
self._logger.info( | ||
f"Executing model.setup_environment with new environment: {environment}" | ||
spal1 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
) | ||
if descriptor.is_async: | ||
return await self._model.setup_environment(environment) | ||
else: | ||
return await to_thread.run_sync(self._model.setup_environment, environment) | ||
|
||
async def poll_for_environment_updates(self) -> None: | ||
last_modified_time = None | ||
environment_config_filename = ( | ||
dynamic_config_resolver.get_dynamic_config_file_path( | ||
dynamic_config_resolver.ENVIRONMENT_DYNAMIC_CONFIG_KEY | ||
) | ||
) | ||
|
||
while True: | ||
# Give control back to the event loop while waiting for environment updates | ||
await asyncio.sleep(POLL_FOR_ENVIRONMENT_UPDATES_TIMEOUT_SECS) | ||
|
||
# Wait for load to finish before checking for environment updates | ||
if not self.ready: | ||
continue | ||
|
||
# Skip polling if no setup_environment implementation provided | ||
if not self.model_descriptor.setup_environment: | ||
break | ||
|
||
if environment_config_filename.exists(): | ||
try: | ||
current_mtime = os.path.getmtime(environment_config_filename) | ||
if not last_modified_time or last_modified_time != current_mtime: | ||
environment_str = await dynamic_config_resolver.get_dynamic_config_value_async( | ||
dynamic_config_resolver.ENVIRONMENT_DYNAMIC_CONFIG_KEY | ||
) | ||
if environment_str: | ||
last_modified_time = current_mtime | ||
environment_json = json.loads(environment_str) | ||
# Avoid rerunning `setup_environment` with the same environment | ||
if self._environment != environment_json: | ||
await self.setup_environment(environment_json) | ||
self._environment = environment_json | ||
except Exception as e: | ||
logging.error( | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think it shows up nicer if you do:
|
||
f"An error occurred while polling for environment updates: {e}" | ||
) | ||
|
||
async def preprocess( | ||
self, | ||
inputs: serialization.InputType, | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -17,3 +17,4 @@ pyyaml==6.0.0 | |
requests==2.31.0 | ||
uvicorn==0.24.0 | ||
uvloop==0.19.0 | ||
aiofiles==24.1.0 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,13 +1,28 @@ | ||
from pathlib import Path | ||
from typing import Optional | ||
|
||
import aiofiles | ||
|
||
DYNAMIC_CONFIG_MOUNT_DIR = "/etc/b10_dynamic_config" | ||
ENVIRONMENT_DYNAMIC_CONFIG_KEY = "environment" | ||
|
||
|
||
def get_dynamic_config_value(key: str) -> Optional[str]: | ||
def get_dynamic_config_value_sync(key: str) -> Optional[str]: | ||
dynamic_config_path = Path(DYNAMIC_CONFIG_MOUNT_DIR) / key | ||
if dynamic_config_path.exists() and dynamic_config_path.is_file(): | ||
if dynamic_config_path.exists(): | ||
with dynamic_config_path.open() as dynamic_config_file: | ||
dynamic_config_value = dynamic_config_file.read() | ||
return dynamic_config_value | ||
return dynamic_config_file.read() | ||
return None | ||
|
||
|
||
def get_dynamic_config_file_path(key: str): | ||
dynamic_config_path = Path(DYNAMIC_CONFIG_MOUNT_DIR) / key | ||
return dynamic_config_path | ||
|
||
|
||
async def get_dynamic_config_value_async(key: str) -> Optional[str]: | ||
dynamic_config_path = get_dynamic_config_file_path(key) | ||
if dynamic_config_path.exists(): | ||
async with aiofiles.open(dynamic_config_path, "r") as dynamic_config_file: | ||
return await dynamic_config_file.read() | ||
return None |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,2 +1,2 @@ | ||
model_name: Test Loaf Failure | ||
model_name: Test Load Failure | ||
python_version: py39 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
can we use a less specific version here? (maybe "^24.1.0")