`truss predict` bug fixes / unit tests #723

helenlyang · 2023-11-08T23:12:59Z

This PR addresses several issues with truss predict:

~~Unifies logic for getting the deployment version given a model_name or model_id~~
~~Supports using the --published flag with --model (i.e. model ID)~~
- The above two points result in different and potentially unexpected behavior for users. Will leave this for a future PR--it should be bundled with changes to the API calls on the Baseten UI and more transparency in Truss CLI on what version is being used in predict
Calling truss predict --published returns production version instead of the first non-development version
Adds unit test coverage for helper functions in core and for BasetenRemote

Testing

Unit tests

poetry run pytest truss/tests/remote/baseten/test_core.py
poetry run pytest truss/tests/remote/baseten/test_remote.py

Truss CLI

Tested manually on a model with four published deployments and one development deployment:

For the production deployment, I set the predict function of my Model to return "production model" as its output. For the development deployment, I return "development model". The other deployments simply return the model_input.

Model version ID:

@helenlyang ➜ /workspaces/truss/my-test-model (helenyang/bt-8803-fix-truss-predict) $ # production version
@helenlyang ➜ /workspaces/truss/my-test-model (helenyang/bt-8803-fix-truss-predict) $ poetry run truss predict --model-deployment qrj6d03 --data {}
? 🎮 Which remote do you want to connect to? prod
{
  "message": "production model"
}

@helenlyang ➜ /workspaces/truss/my-test-model (helenyang/bt-8803-fix-truss-predict) $ # development version
@helenlyang ➜ /workspaces/truss/my-test-model (helenyang/bt-8803-fix-truss-predict) $ poetry run truss predict --model-deployment qjd502q --data {}
? 🎮 Which remote do you want to connect to? prod
{
  "message": "development model"
}

@helenlyang ➜ /workspaces/truss/my-test-model (helenyang/bt-8803-fix-truss-predict) $ # published, non-production version
@helenlyang ➜ /workspaces/truss/my-test-model (helenyang/bt-8803-fix-truss-predict) $ poetry run truss predict --model-deployment qvvd0eq --data {}
? 🎮 Which remote do you want to connect to? prod
{}

Model name:

@helenlyang ➜ /workspaces/truss/my-test-model (helenyang/bt-8803-fix-truss-predict) $ # without --published; should use dev version
@helenlyang ➜ /workspaces/truss/my-test-model (helenyang/bt-8803-fix-truss-predict) $ poetry run truss predict --data {}
? 🎮 Which remote do you want to connect to? prod
{
  "message": "development model"
}

@helenlyang ➜ /workspaces/truss/my-test-model (helenyang/bt-8803-fix-truss-predict) $ # with --published; should use production version
@helenlyang ➜ /workspaces/truss/my-test-model (helenyang/bt-8803-fix-truss-predict) $ poetry run truss predict --published --data {}
? 🎮 Which remote do you want to connect to? prod
{
  "message": "production model"
}

Model ID:

@helenlyang ➜ /workspaces/truss/my-test-model (helenyang/bt-8803-fix-truss-predict) $ # should use production version
@helenlyang ➜ /workspaces/truss/my-test-model (helenyang/bt-8803-fix-truss-predict) $ poetry run truss predict --model 03ydnk43 --data {}
? 🎮 Which remote do you want to connect to? prod
{
  "message": "production model"
}

Testing on a second model with no production deployment:

@helenlyang ➜ /workspaces/truss/my-test-model (helenyang/bt-8803-fix-truss-predict) $ # if no production version exists, use the development version
@helenlyang ➜ /workspaces/truss/my-test-model (helenyang/bt-8803-fix-truss-predict) $ poetry run truss predict --model yqvvy0jq --data {}
? 🎮 Which remote do you want to connect to? prod
{
  "message": "development model"
}

* adds API for returning versions given model_id * separate functions for getting model and version IDs and finding matching version * service URL only uses model_versions endpoint * check is_primary to find production deployment

linear · 2023-11-08T23:13:04Z

BT-8803 Fix `truss predict`

Problem to Fix

Model name logic is different to the model ID logic
Not respecting the —published flag if you pass a model-id
Errors are very hard to read (separate PR)
Logic is not unit tested

Correct truss predict logic:

If not published, there is only one development deployment, so use that
If it is published, use the production deployment (primary_version is set)
If there is no production deployment, use the latest published deployment

Repro steps:

I did truss push to build my truss. The build failed
I made a change, and did truss push again. The build succeeded, the model successfully deployed
I did truss predict , and got “Model is unhealthy, it is not ready to make predictions”

helenlyang · 2023-11-11T03:12:37Z

truss/remote/baseten/api.py

@@ -157,23 +157,22 @@ def models(self):
    def get_model(self, model_name):


The changes to get_model should be safe--the current callsites are https://github.com/basetenlabs/truss/blob/8901a0121dda36b822fbabe62f8a172ba72c6c6e/truss/remote/baseten/core.py#L50C1-L61 and are updated in this PR

helenlyang · 2023-11-11T03:15:27Z

truss/remote/baseten/api.py

@@ -184,18 +183,19 @@ def get_model_by_id(self, model_id: str):
            model(id: "{model_id}") {{
                name
                id
-                primary_version{{


get_model_by_id changes should also be safe, since this was the only callsite:

truss/truss/remote/baseten/remote.py

Lines 104 to 106 in 8901a01

model = self._api.get_model_by_id(model_identifier.value)

model_id = model["model"]["id"]

model_version_id = model["model"]["primary_version"]["id"]

helenlyang · 2023-11-11T03:17:48Z

truss/remote/baseten/core.py

    for version in versions:
-        if version["is_draft"] is True:
+        if version["is_primary"] and not version["is_draft"]:


I don't think it's possible for is_primary to be true for development models, but thought it'd be safer to check is_draft in case the is_primary definition ever changes

You're right -- it's not possible for is_primary to be true for development models. I think it's a little cleaner to leave this as an is_primary check just to have a cleaner definition of what "production" means (ie: is_primary = True), which is easier to understand. But I dont' feel strongly

truss/remote/baseten/remote.py

squidarth

this is very complex code, so awesome work working through it! Left some code comments, but also in the PR description, it would be awesome if you could drop a few CLI invocations that are worth testing again

squidarth · 2023-11-14T16:01:26Z

truss/remote/baseten/core.py

+    for version in versions:
+        if version["is_draft"] is True:
+            return version
+    raise ValueError("No development version found")


something to consider here -- what if instead of using exceptions to represent the cases where there's no version, we returned None (so this function return type becomes Optional[dict])? And then the callers can decide how they want to handle that case (vs. having to catch exceptions)

I think that's nicer--changed these functions except for get_dev_version_info, which has some callsites in truss watch / patch logic. I'll leave updating that function and its callsites for a future PR to avoid increasing the scope of this one

squidarth · 2023-11-14T16:09:16Z

truss/cli/cli.py

-        raise click.UsageError(
-            "Cannot use --published with --model or --model-version."
-        )
+    if published and model_version_id:


now that we're using the /models endpoint if the person passes --model, we're ignoring the published flag again. So I think we should keep this error message the same.

SG, reverted this change

squidarth · 2023-11-14T16:16:31Z

truss/remote/baseten/core.py

    for version in versions:
-        if version["is_draft"] is True:
+        if version["is_primary"] and not version["is_draft"]:


You're right -- it's not possible for is_primary to be true for development models. I think it's a little cleaner to leave this as an is_primary check just to have a cleaner definition of what "production" means (ie: is_primary = True), which is easier to understand. But I dont' feel strongly

squidarth · 2023-11-14T16:17:34Z

truss/remote/baseten/remote.py

+        # Return the production deployment version.
+        try:
+            return get_prod_version_info_from_versions(model_versions)
+        except ValueError:


per comment in the other file, let's change get_prod_version_info_from_versions to return None in the case where there's no prod version, so we don't ahve to use exceptions for control flow

truss/remote/baseten/remote.py

squidarth · 2023-11-14T16:24:44Z

truss/tests/remote/baseten/test_remote.py

+_TEST_REMOTE_URL = "http://test_remote.com"
+
+
+def test_get_service_by_version_id():


should we also have test for a no model w/ version exists?

Good call, this made me realize that we should probably raise a UsageError rather than propagating the Baseten ApiError if the model version doesn't exist. Added a try / except there and a unit test (let me know if there's a better way to mock errors)

…test

This reverts commit 2d7bc49.

helenlyang

Thanks for the review! Added CLI commands to the description.

Following up on our conversation about model_id vs model_name inconsistency--I updated the PR description with our decision to stick with the current behavior. However, I do think we should consider unifying these code paths in the future.

I think it might make model version resolution less of a black box if we logged which model version is being called in truss push, e.g. f"Calling {'production' if published else 'development'} version {model_version_id} of model {model_id}". That could make changing truss predict by model_id to match model_name behavior less surprising to users--but could also be useful regardless

helenlyang · 2023-11-14T22:35:20Z

truss/remote/baseten/core.py

+    for version in versions:
+        if version["is_draft"] is True:
+            return version
+    raise ValueError("No development version found")


I think that's nicer--changed these functions except for get_dev_version_info, which has some callsites in truss watch / patch logic. I'll leave updating that function and its callsites for a future PR to avoid increasing the scope of this one

helenlyang · 2023-11-14T22:39:17Z

truss/remote/baseten/core.py

    for version in versions:
-        if version["is_draft"] is True:
+        if version["is_primary"] and not version["is_draft"]:


helenlyang · 2023-11-14T22:56:16Z

truss/tests/remote/baseten/test_remote.py

+_TEST_REMOTE_URL = "http://test_remote.com"
+
+
+def test_get_service_by_version_id():


Good call, this made me realize that we should probably raise a UsageError rather than propagating the Baseten ApiError if the model version doesn't exist. Added a try / except there and a unit test (let me know if there's a better way to mock errors)

helenlyang · 2023-11-14T23:36:53Z

truss/cli/cli.py

-        raise click.UsageError(
-            "Cannot use --published with --model or --model-version."
-        )
+    if published and model_version_id:


SG, reverted this change

helenlyang · 2023-11-15T00:54:01Z

truss/remote/baseten/remote.py

+            # TODO(helen): make this consistent with getting the service via
+            # model_name and respect --published in service_url_path.
+            model_version = BasetenRemote._get_matching_version(
+                model_versions, published
+            )
+            model_version_id = model_version["id"]


This implementation is now weirdly in between the original implementation and my attempt to unify the model_id and model_name code paths. It might be more readable to revert this back to the original model_id code path and leave these changes for if / when we unify model_id and model_name paths

Yeah, I think for now it might better to revert back for now. It's one less call at least.

squidarth

Overall, LGTM! Awesome work here. Just left a couple more comments, feel free to resolve & ship

squidarth · 2023-11-15T21:40:51Z

truss/remote/baseten/remote.py

+            # TODO(helen): make this consistent with getting the service via
+            # model_name and respect --published in service_url_path.
+            model_version = BasetenRemote._get_matching_version(
+                model_versions, published
+            )
+            model_version_id = model_version["id"]


Yeah, I think for now it might better to revert back for now. It's one less call at least.

squidarth · 2023-11-15T21:44:40Z

truss/remote/baseten/core.py

+    return (query_result["id"], query_result["versions"])
+
+
+def get_dev_version_info_from_versions(versions: List[dict]) -> Optional[dict]:


Let's make this private (ie: def _get_dev_version_from_versions()...)

I left this as public so we can get the dev version without requiring another GraphQL query (api.get_model is called inside get_dev_version_info). I added docstrings to try to make this more clear, let me know if that makes sense

I see, sounds good

squidarth · 2023-11-15T21:45:11Z

truss/remote/baseten/core.py

    return (query_result["id"], query_result["versions"])


+def get_model_versions_info_by_id(


A lot of these functions have info which I think doesn't add a lot, I think this could just be def get_model_versions_by_id

squidarth · 2023-11-15T21:57:45Z

truss/remote/baseten/core.py

+    versions = model["model"]["versions"]
+    dev_version = get_dev_version_info_from_versions(versions)
+    if not dev_version:
+        # TODO(helen): return dev_version in all cases rather than raising an error


let's try to do this todo in a follow-up should be fairly straightforward

Follow-up to #723. This updates `core.get_dev_version` to return None if a development version doesn’t exist and updates callsites to raise errors.

helenlyang added 3 commits November 8, 2023 05:37

WIP first cut

2142cac

* adds API for returning versions given model_id * separate functions for getting model and version IDs and finding matching version * service URL only uses model_versions endpoint * check is_primary to find production deployment

update model_version query, add TODO about bug

d8c823b

use model query for model id

1c26b8a

helenlyang added 5 commits November 10, 2023 23:30

update get_model to use model query

fdb65d4

add API for getting dev version given versions

5880caa

use get_dev_version and get_prod_version helpers

d82d667

raise click.UsageError in cli

e2cddc1

add test cases to test_core

649e70f

helenlyang commented Nov 11, 2023

View reviewed changes

helenlyang added 4 commits November 11, 2023 03:31

make remote methods static

13b59af

add tests for get_matching_version

dc1a845

mock requests + test get_service instead

58c2799

support --published with --model in cli

2d7bc49

helenlyang changed the title ~~Helenyang/bt 8803 fix truss predict~~ Fix inconsistencies in truss predict Nov 13, 2023

clean up

0c70f04

helenlyang marked this pull request as ready for review November 13, 2023 23:16

helenlyang requested a review from squidarth November 13, 2023 23:17

helenlyang commented Nov 13, 2023

View reviewed changes

truss/remote/baseten/remote.py Outdated Show resolved Hide resolved

squidarth reviewed Nov 14, 2023

View reviewed changes

helenlyang added 5 commits November 14, 2023 22:37

sid CR: return None for no matching version in core.py

0a0271c

sid CR 2: remove is_draft check

89ed524

except ApiError in model version path + sid CR: add no model version …

652be3c

…test

Revert "support --published with --model in cli"

75f2550

This reverts commit 2d7bc49.

model_name uses model_versions endpoint in service url

b3c31c5

helenlyang commented Nov 15, 2023

View reviewed changes

Merge branch 'main' into helenyang/bt-8803-fix-truss-predict

43678ee

squidarth approved these changes Nov 15, 2023

View reviewed changes

sid CR: rm info from function names

b3f0d44

helenlyang added 2 commits November 15, 2023 23:58

add docstring for get_dev_version / get_dev_version_from_versions

9edb25c

revert model ID code path changes

3deb879

squidarth approved these changes Nov 16, 2023

View reviewed changes

helenlyang changed the title ~~Fix inconsistencies in truss predict~~ truss predict bug fixes / unit tests Nov 16, 2023

helenlyang merged commit b5cccc5 into main Nov 16, 2023

helenlyang deleted the helenyang/bt-8803-fix-truss-predict branch November 16, 2023 18:16

linear bot mentioned this pull request Nov 20, 2023

Move error handling out of get_dev_version #746

Merged

helenlyang added a commit that referenced this pull request Nov 29, 2023

Move error handling out of get_dev_version #746

a729155

Follow-up to #723. This updates `core.get_dev_version` to return None if a development version doesn’t exist and updates callsites to raise errors.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

`truss predict` bug fixes / unit tests #723

`truss predict` bug fixes / unit tests #723

helenlyang commented Nov 8, 2023 •

edited

Loading

linear bot commented Nov 8, 2023

helenlyang Nov 11, 2023

helenlyang Nov 11, 2023

helenlyang Nov 11, 2023

squidarth Nov 14, 2023

helenlyang Nov 14, 2023

squidarth left a comment

squidarth Nov 14, 2023

helenlyang Nov 14, 2023

squidarth Nov 14, 2023

helenlyang Nov 14, 2023

squidarth Nov 14, 2023

squidarth Nov 14, 2023

squidarth Nov 14, 2023

helenlyang Nov 14, 2023

helenlyang left a comment

helenlyang Nov 14, 2023

helenlyang Nov 14, 2023

helenlyang Nov 14, 2023

helenlyang Nov 14, 2023

helenlyang Nov 15, 2023

squidarth Nov 15, 2023

helenlyang Nov 16, 2023

squidarth left a comment

squidarth Nov 15, 2023

squidarth Nov 15, 2023

helenlyang Nov 15, 2023

squidarth Nov 16, 2023

squidarth Nov 15, 2023

helenlyang Nov 15, 2023

squidarth Nov 15, 2023

helenlyang Nov 15, 2023

		@@ -157,23 +157,22 @@ def models(self):
		def get_model(self, model_name):

	model = self._api.get_model_by_id(model_identifier.value)
	model_id = model["model"]["id"]
	model_version_id = model["model"]["primary_version"]["id"]

		_TEST_REMOTE_URL = "http://test_remote.com"


		def test_get_service_by_version_id():

		return (query_result["id"], query_result["versions"])


		def get_dev_version_info_from_versions(versions: List[dict]) -> Optional[dict]:

		return (query_result["id"], query_result["versions"])


		def get_model_versions_info_by_id(

truss predict bug fixes / unit tests #723

truss predict bug fixes / unit tests #723

Conversation

helenlyang commented Nov 8, 2023 • edited Loading

Testing

Unit tests

Truss CLI

linear bot commented Nov 8, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

squidarth left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

helenlyang left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

squidarth left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

`truss predict` bug fixes / unit tests #723

`truss predict` bug fixes / unit tests #723

helenlyang commented Nov 8, 2023 •

edited

Loading