Add `setup_environment` to truss #1188

spal1 · 2024-10-11T04:14:46Z

🚀 What

This PR adds a new setup_environment function that we will call when we notice a change to the environment a truss is running in. This environment refers to the release environment, ex: production or staging. In order to notice these changes, we will need to watch for updates to a particular environment file mounted on the b10_dynamic_config directory.

Going to add support for setup_environment for chainlets in a separate PR: #1192

💻 How

From the truss server on_startup, we call a `setup_polling_for_environment_updates() function that will create the asyncio task that will run the file watcher
Inside the task, we will wait for the ModelWrapper._status to be READY, indicating that the load() was successfully completed, and then will start actually watching the /etc/b10_dynamic_config/environment file
Whenever the modification time of the file changes, we will use the dynamic_config_resolver to access the contents of the /etc/b10_dynamic_config/environment file and call the user-defined setup_environment function with the contents

🔬 Testing

TODO:

call user-defined setup_environment from load_impl - this is to enable setting up the environment before user-defined load is called
Add integration tests to truss/tests/test_model_inference.py
Add unit test for new async functions in dynamic_config_resolver
End to end test with the configmap patching flow in dev

Loom with e2e test: https://www.loom.com/share/cbe8e03a7d97434f99f973eaceb33c59?sid=ca51f099-fd45-4b39-a7bf-366d7d01356b

linear · 2024-10-11T04:14:48Z

BT-12427 Add support for `setup_environment` to truss

truss/templates/server/truss_server.py

truss/templates/server/model_wrapper.py

truss/test_data/file_watcher_truss/model/model.py

truss/templates/server/model_wrapper.py

truss/test_data/file_watcher_truss/model/model.py

truss/templates/server/model_wrapper.py

truss/tests/test_model_inference.py

rcano-baseten · 2024-10-15T15:49:29Z

truss/truss_handle.py

@@ -298,6 +298,12 @@ def _run_docker(gpus: Optional[str] = None):
                            f"src={str(LocalConfigHandler.bptr_data_resolution_dir_path())}",
                            "target=/bptr",
                        ],
+                        [


for my own edification, what does this help accomplish?

we mimic the mounted configmap directory (which is /etc/b10_dynamic_config on model pods) here so we can make changes to /etc/b10_dynamic_config/environment in the integration tests to mimic patches to the environment key in the dynamic_config configmap

truss/templates/server/model_wrapper.py

squidarth · 2024-10-16T03:41:49Z

truss/templates/server/model_wrapper.py

+
+            # Skip polling if no setup_environment implementation provided
+            if not self.model_descriptor.setup_environment:
+                self._logger.info("No model.setup_environment definition provided")


I don't think we need to log here, if the user didn't provide a setup_environment method, I don't think this is the right way to communicate this to them.

On the other hand, if they did provide a method, it might be good to log something indicating that we will poll for changes to the env.

truss/templates/server/model_wrapper.py

truss/templates/shared/dynamic_config_resolver.py

truss/tests/test_model_inference.py

squidarth

Tests look good, would you mind making this a context builder and testing on dev/staging? Ideally we can run truss-examples on it (I can walk you through how to do that)

squidarth · 2024-10-16T22:04:18Z

pyproject.toml

@@ -96,6 +97,7 @@ pytest = "7.2.0"
 pytest-cov = "^3.0.0"
 types-PyYAML = "^6.0.12.12"
 types-setuptools = "^69.0.0.0"
+types-aiofiles = "^24.1.0.20240626"


can we use a less specific version here? (maybe "^24.1.0")

going to see if i can update this as a follow-up, still running into poetry lock --no-update hanging for me

truss/templates/server/model_wrapper.py

squidarth

One last comment, but this looks awesome, amazing work here!!

truss/templates/server/model_wrapper.py

spal1 added the do not merge label Oct 11, 2024

spal1 commented Oct 11, 2024

View reviewed changes

truss/templates/server/truss_server.py Outdated Show resolved Hide resolved

spal1 commented Oct 11, 2024

View reviewed changes

truss/templates/server/truss_server.py Outdated Show resolved Hide resolved

spal1 commented Oct 11, 2024

View reviewed changes

truss/templates/server/model_wrapper.py Show resolved Hide resolved

spal1 commented Oct 11, 2024

View reviewed changes

truss/templates/server/model_wrapper.py Outdated Show resolved Hide resolved

spal1 commented Oct 11, 2024

View reviewed changes

truss/templates/server/model_wrapper.py Outdated Show resolved Hide resolved

spal1 commented Oct 14, 2024

View reviewed changes

truss/templates/server/model_wrapper.py Outdated Show resolved Hide resolved

squidarth reviewed Oct 14, 2024

View reviewed changes

truss/templates/server/model_wrapper.py Outdated Show resolved Hide resolved

truss/templates/server/model_wrapper.py Outdated Show resolved Hide resolved

truss/templates/server/model_wrapper.py Outdated Show resolved Hide resolved

spal1 requested a review from rcano-baseten October 14, 2024 13:59

spal1 commented Oct 14, 2024

View reviewed changes

truss/test_data/file_watcher_truss/model/model.py Outdated Show resolved Hide resolved

squidarth reviewed Oct 14, 2024

View reviewed changes

truss/templates/server/model_wrapper.py Outdated Show resolved Hide resolved

rcano-baseten reviewed Oct 15, 2024

View reviewed changes

truss/tests/test_model_inference.py Outdated Show resolved Hide resolved

rcano-baseten reviewed Oct 15, 2024

View reviewed changes

spal1 force-pushed the samiksha/bt-12427-add-support-for-setup_environment-to-truss branch 2 times, most recently from 59a4aa9 to 3bae49a Compare October 16, 2024 00:41

spal1 changed the title ~~Add setup_environment~~ Add setup_environment to truss Oct 16, 2024

squidarth reviewed Oct 16, 2024

View reviewed changes

spal1 marked this pull request as ready for review October 16, 2024 21:49

spal1 removed the do not merge label Oct 16, 2024

squidarth reviewed Oct 16, 2024

View reviewed changes

spal1 requested review from squidarth and rcano-baseten October 18, 2024 22:16

squidarth approved these changes Oct 18, 2024

View reviewed changes

truss/templates/server/model_wrapper.py Outdated Show resolved Hide resolved

spal1 self-assigned this Oct 18, 2024

spal1 added 4 commits October 22, 2024 03:49

Bump Version for CTX builder

3085800

skip long test

60f42f6

add aiofiles to requirements

bfbb0e2

increase ctx builder version

6e52313

spal1 added 26 commits October 22, 2024 03:49

increase ctx builder

d522bc7

pass event_loop down

fc68cd1

allow event_loop to be optional

26b220b

only thread self._model.load

1e0e637

revert test_server changes

776db72

add integration test for setup_environment

314d7a8

add truss for setup_environment integration test

1435142

revert model wrapper test changes

81b9303

revert load changes and check model wrapper status instead + tests

6d2a790

clean up resolver

30c2877

fix resolver test

869f426

call setup env before load

4369a5b

remove print statements

1a3348e

unit tests + fixes

69abf2a

sid CR

54ac737

fix polling structure

36c35e8

CR

078e574

bump rc version

c86803e

support for clearing out env

7bac078

update ctx builder and revert version of types-aiofiles

77b0289

test for Nones in configmap file

005b97f

better error handling

7464089

update ctx builder version

f5c71f8

test fix + better error msg

24df6d3

move sleep

db748cf

add cleanup to env integration tests

ba0e5fd

spal1 force-pushed the samiksha/bt-12427-add-support-for-setup_environment-to-truss branch from bffc63b to ba0e5fd Compare October 22, 2024 05:15

add a todo for async setup nv + more logging

b6577c2

spal1 merged commit 5906d60 into main Oct 22, 2024
5 checks passed

spal1 deleted the samiksha/bt-12427-add-support-for-setup_environment-to-truss branch October 22, 2024 15:15

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `setup_environment` to truss #1188

Add `setup_environment` to truss #1188

spal1 commented Oct 11, 2024 •

edited

Loading

linear bot commented Oct 11, 2024

rcano-baseten Oct 15, 2024

spal1 Oct 15, 2024 •

edited

Loading

squidarth Oct 16, 2024

squidarth left a comment

squidarth Oct 16, 2024

spal1 Oct 22, 2024

squidarth left a comment

Add setup_environment to truss #1188

Add setup_environment to truss #1188

Conversation

spal1 commented Oct 11, 2024 • edited Loading

🚀 What

💻 How

🔬 Testing

linear bot commented Oct 11, 2024

rcano-baseten Oct 15, 2024

Choose a reason for hiding this comment

spal1 Oct 15, 2024 • edited Loading

Choose a reason for hiding this comment

squidarth Oct 16, 2024

Choose a reason for hiding this comment

squidarth left a comment

Choose a reason for hiding this comment

squidarth Oct 16, 2024

Choose a reason for hiding this comment

spal1 Oct 22, 2024

Choose a reason for hiding this comment

squidarth left a comment

Choose a reason for hiding this comment

Add `setup_environment` to truss #1188

Add `setup_environment` to truss #1188

spal1 commented Oct 11, 2024 •

edited

Loading

spal1 Oct 15, 2024 •

edited

Loading