Lazy model connection & x-route-info #363

mynhardtburger · 2024-06-06T15:44:44Z

This PR adds support for respecting a TGIS hostname override (x-route-info header/metadata).

This is achieved by lazily registering the model connection with the TGISBackend when .run()/run_stream_out()/run_tokenizer() is called. The x-route-info header/metadata is read from the context arg.

Additionally:

a TextGenerationTGIS._tgis_backend property was added to be consistent with PeftPromptTuningTGIS. The TextGenerationTGIS.tgis_backend property is conditional on the tgis_backend argument to be True, which made accessing TextGenerationTGIS.tgis_backend risky as it might not exist. The additional _tgis_backend property is guaranteed to exist, making using it simpler and less risky.

Depends on caikit-tgis-backend>=0.1.33

Signed-off-by: Mynhardt Burger <[email protected]>

gabe-l-hart

Super close pending changes over in caikit-tgis-backend!

caikit_nlp/modules/text_generation/peft_tgis_remote.py

caikit_nlp/modules/text_generation/text_generation_tgis.py

caikit_nlp/toolkit/text_generation/tgis_utils.py

Signed-off-by: Mynhardt Burger <[email protected]>

gabe-l-hart

A couple of small cleanup suggestions and comment requests

gabe-l-hart · 2024-06-07T16:36:12Z

caikit_nlp/modules/text_generation/peft_tgis_remote.py

@@ -221,6 +233,9 @@ def run(
            self.enable_backend,
            "Backend must be configured and loaded with this module before executing `run` call.",
        )
+        if self._tgis_backend:


Total NIT, but since _register_model_connection_with_context is a member function, you could move the if self._tgis_backend line into the function and save a few duplicate lines in each of the inference methods.

gabe-l-hart · 2024-06-07T16:41:01Z

caikit_nlp/modules/text_generation/peft_tgis_remote.py

+        ok, route_info = get_route_info(context)
+        if ok:


I really want to be able to use := here and make it look more like go, but seems that you can't

gabe-l-hart · 2024-06-07T16:42:04Z

caikit_nlp/modules/text_generation/peft_tgis_remote.py

+                self.base_model_name, {"hostname": route_info}, fill_with_defaults=True
+            )
+        else:
+            self._tgis_backend.register_model_connection(self.base_model_name)


I think we don't need to do this if there's no context because it will be done automatically with get_client

gabe-l-hart · 2024-06-07T16:42:58Z

caikit_nlp/modules/text_generation/text_generation_tgis.py

@@ -240,6 +249,9 @@ def run(
            GeneratedTextResult
                Generated text result produced by TGIS.
        """
+        if self._tgis_backend:


Same comment about burying the if self._tgis_backend

gabe-l-hart · 2024-06-07T16:44:02Z

caikit_nlp/modules/text_generation/text_generation_tgis.py

+                self.model_name, {"hostname": route_info}, fill_with_defaults=True
+            )
+        else:
+            self._tgis_backend.register_model_connection(self.model_name)


Same comment about not needing the else clause

caikit_nlp/toolkit/text_generation/tgis_utils.py

gabe-l-hart · 2024-06-07T16:50:24Z

caikit_nlp/toolkit/text_generation/tgis_utils.py

+def get_route_info(
+    context: Optional[RuntimeServerContextType],
+) -> Tuple[bool, Optional[str]]:
+    """
+    Returns a tuple `(True, x-route-info)` from context if "x-route-info" was found in
+    the headers/metadata.
+
+    Otherwise returns a tuple `(False, None)` if "x-route-info" was not found in the
+    context or if context is None.
+    """


I think we can collapse the return into a simple Optional[str]. Returning None is indicative of "no route info found." This would simplify the logic that uses it and allow us to get another walrus!

Done. Now using walrus operator for good @gabe-l-hart karma.

gabe-l-hart · 2024-06-07T16:52:53Z

tests/fixtures/__init__.py

+
+class TestServicerContext:
+    """
+    A dummy class for mimicking ServicerContext invocation metadata storage.


Interesting. I was looking for a link to suggest as a citation of this interface and I found this instead. Makes me wonder if we should be using grpcio-testing for all of this in our tests!

Here's the citation link: https://github.com/grpc/grpc/blob/master/src/python/grpcio/grpc/__init__.py#L1128

Co-authored-by: Gabe Goodhart <[email protected]> Signed-off-by: Mynhardt Burger <[email protected]>

Signed-off-by: Mynhardt Burger <[email protected]>

gabe-l-hart

One list NIT, but I think this is also good as is. Thanks!

gabe-l-hart · 2024-06-07T17:34:39Z

caikit_nlp/toolkit/text_generation/tgis_utils.py

+            ValueError(f"context is of an unsupported type: {type(context)}"),
+        )
+
+    return None


NIT: this is both implicit and unreachable since else will raise. I'm guessing it's here because your IDE doesn't know that?

Done. Removed.

Signed-off-by: Mynhardt Burger <[email protected]>

mynhardtburger added 4 commits June 6, 2024 10:56

add get_route_info

e278915

Signed-off-by: Mynhardt Burger <[email protected]>

lazily create model_connection and _client

e473c33

Signed-off-by: Mynhardt Burger <[email protected]>

lazy load model_connection and tgis client for peft

5578b7a

Signed-off-by: Mynhardt Burger <[email protected]>

remove commented out code

71371bd

Signed-off-by: Mynhardt Burger <[email protected]>

gabe-l-hart reviewed Jun 6, 2024

View reviewed changes

mynhardtburger added 4 commits June 7, 2024 09:52

Address review comments

21192c0

Signed-off-by: Mynhardt Burger <[email protected]>

Expand test_get_route_info

1c23527

Signed-off-by: Mynhardt Burger <[email protected]>

Lazily create generation client

a4e8539

Signed-off-by: Mynhardt Burger <[email protected]>

Update minimum caikit-tgis-backend version

be72a79

Signed-off-by: Mynhardt Burger <[email protected]>

mynhardtburger changed the title ~~Lazy model connection~~ Lazy model connection & x-route-info Jun 7, 2024

Add debug logs

d5893d9

Signed-off-by: Mynhardt Burger <[email protected]>

mynhardtburger marked this pull request as ready for review June 7, 2024 16:28

mynhardtburger requested review from alex-jw-brooks, gkumbhat, evaline-ju, tharapalanivel and Ssukriti as code owners June 7, 2024 16:28

mynhardtburger force-pushed the lazy-model-connection branch from 031f5cc to d5893d9 Compare June 7, 2024 16:33

mynhardtburger added 2 commits June 7, 2024 12:44

Linting

9848943

Signed-off-by: Mynhardt Burger <[email protected]>

linting

5df0533

Signed-off-by: Mynhardt Burger <[email protected]>

gabe-l-hart requested changes Jun 7, 2024

View reviewed changes

mynhardtburger and others added 2 commits June 7, 2024 13:13

Update caikit_nlp/toolkit/text_generation/tgis_utils.py

6bae66a

Co-authored-by: Gabe Goodhart <[email protected]> Signed-off-by: Mynhardt Burger <[email protected]>

review comments

46ae073

Signed-off-by: Mynhardt Burger <[email protected]>

gabe-l-hart approved these changes Jun 7, 2024

View reviewed changes

remove unreachable code

527b455

Signed-off-by: Mynhardt Burger <[email protected]>

gabe-l-hart merged commit 3b2f8fd into caikit:main Jun 7, 2024
5 checks passed

gabe-l-hart mentioned this pull request Jun 18, 2024

Module backend context update caikit/caikit#735

Merged

3 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Lazy model connection & x-route-info #363

Lazy model connection & x-route-info #363

mynhardtburger commented Jun 6, 2024 •

edited

Loading

gabe-l-hart left a comment

gabe-l-hart left a comment

gabe-l-hart Jun 7, 2024

gabe-l-hart Jun 7, 2024

gabe-l-hart Jun 7, 2024

gabe-l-hart Jun 7, 2024

gabe-l-hart Jun 7, 2024

mynhardtburger Jun 7, 2024

gabe-l-hart Jun 7, 2024

mynhardtburger Jun 7, 2024

gabe-l-hart Jun 7, 2024

gabe-l-hart Jun 7, 2024

gabe-l-hart left a comment

gabe-l-hart Jun 7, 2024

mynhardtburger Jun 7, 2024

Lazy model connection & x-route-info #363

Lazy model connection & x-route-info #363

Conversation

mynhardtburger commented Jun 6, 2024 • edited Loading

gabe-l-hart left a comment

Choose a reason for hiding this comment

gabe-l-hart 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gabe-l-hart left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mynhardtburger commented Jun 6, 2024 •

edited

Loading