Bug bash improvements to Python DX, better error messages #1346

nnarayen · 2025-01-28T20:45:42Z

🚀 What

This is a first pass at improving the new python DX after our bug bash:

Validation that runs during framework.import_target will now be appropriately scoped to either a model or chainlet target
In error message, we surface the classes that users interact with in public_api rather than our abstract ABCChainlet

Next steps:

Improve error messages related to code examples

💻 How

🔬 Testing

nnarayen · 2025-01-28T20:48:17Z

truss-chains/truss_chains/framework.py

-        if utils.issubclass_safe(sym, definitions.ABCChainlet)
-        and cast(definitions.ABCChainlet, sym).meta_data.is_entrypoint
-    }
+class _ABCImporter(abc.ABC):


I personally think this is another good use case for inheritance, since otherwise we'd end up passing lots of function arguments around to account for differences (no_entrypoint_error, multiple_entrypoints_error, target_cls_type, and any in the future)

However, I also brought all dependencies of _import_target inside this class, like _get_entrypoint_chainlets and _load_module. Since nobody else needs those functions, I personally like encapsulating within the class to better indicate who relies on these. An additional benefit (in the future) could be to more easily extract this class to a standalone file if we want to break up the large framework.py one day

Sounds good.

nnarayen · 2025-01-28T21:02:21Z

truss-chains/truss_chains/framework.py

+    @classmethod
+    def no_entrypoint_error(cls, module_path: pathlib.Path) -> ValueError:
+        return ValueError(
+            f"No Model class in `{module_path}` inherits from {cls.target_cls_type()}."


Previously, the error message would point to ABCChainlet, but users aren't interacting with that class directly. I think that including the public_api class is slightly better here

Good call. It think the reason was the cyclic dependency issue... let's fix both this time.

nnarayen · 2025-01-28T21:02:49Z

truss-chains/truss_chains/framework.py

+        }
+
+    @classmethod
+    def _load_module(cls, module_path: pathlib.Path) -> tuple[types.ModuleType, Loader]:


_load_module and _cleanup_module_imports have no changes, other than the ones needed to make them class methods

I trust you on this, I'm not gonna do a line-by-line review this time ;)

nnarayen · 2025-01-28T21:03:14Z

truss-chains/truss_chains/framework.py

-    try:
+    @classmethod
+    @abc.abstractmethod
+    def target_cls_type(cls) -> Type[definitions.ABCChainlet]:


Having an intended target type solves https://basetenlabs.slack.com/archives/C033TL5SLH0/p1738087812972669?thread_ts=1738086363.048469&cid=C033TL5SLH0

This can be private though (and also others)?

nnarayen · 2025-01-28T21:04:17Z

truss-chains/truss_chains/deployment/code_gen.py

@@ -773,7 +773,7 @@ def gen_truss_chainlet(
    gen_root = pathlib.Path(tempfile.gettempdir())
    chainlet_dir = _make_chainlet_dir(chain_name, chainlet_descriptor, gen_root)
    logging.info(
-        f"Code generation for Chainlet `{chainlet_descriptor.name}` "
+        f"Code generation for {chainlet_descriptor.chainlet_cls.entity_type} `{chainlet_descriptor.name}` "


This change switches between Model / Chainlet depending on the context

marius-baseten · 2025-01-29T00:10:54Z

truss-chains/tests/test_framework.py

@@ -2,6 +2,7 @@
 import contextlib
 import logging
 import re
+from pathlib import Path


Please stick to convention to not import symbols (except typing)

Ugh sorry, my editor uses this convention by default. I'll investigate a way to change that

what if you just type pathlitb.Path at the call site - will it add import pathlib automatically? Or will it create from pathlib import Path?

Unfortunately doesn't auto import with the pathlib.Path syntax, but I have a unique dev setup locally - I can probably fix this!

marius-baseten · 2025-01-29T00:13:18Z

truss-chains/tests/import/standalone_without_entrypoint.py

@@ -0,0 +1,7 @@
+class PassthroughModel:


I got confused by this name (and didn't see at first that the base class is missing). What about "NonModelClass" or "NonChainletClass" (and also update the file name)?

marius-baseten · 2025-01-29T00:14:35Z

truss-chains/truss_chains/framework.py

-        if utils.issubclass_safe(sym, definitions.ABCChainlet)
-        and cast(definitions.ABCChainlet, sym).meta_data.is_entrypoint
-    }
+class _ABCImporter(abc.ABC):


Sounds good.

marius-baseten · 2025-01-29T00:17:25Z

truss-chains/truss_chains/framework.py

@@ -35,7 +36,7 @@
 import pydantic
 from typing_extensions import ParamSpec

-from truss_chains import definitions, utils
+from truss_chains import definitions, public_api, utils


Can we avoid cyclic dependencies somehow?

Maybe ChainletBase should not be directly in public API.
It might fit into definitions. But the goal there was to have mostly constants, interfaces and datastructures (i.e. no implementaiton logic) -
which is already not well fulfilled, because it does contain some impl.

Maybe it's time to move ChainletBase and some stuff from definitions either into framework or a new module (depending on what makes more sense).

Which ever you chose, the dependencies should be somewhat straight...

framework requires these classes for validation / error messages, and these classes require framework for initialization logic, so for now I think it actually makes sense for them all to belong in framework? I don't have too many opinions on the future of definitions / public_api, but I think framework needs to be split up into multiple modules / files.

We will need to grab these from framework when choosing what to expose from truss_chains now, but since public_api already depends on framework I don't think this is too bad. However, don't have much experience with python import management, so open to thoughts here!

Ok, then move ChainletBase to framework and add an alias in public_api.

Generally we should follow this (for imports and other things): https://google.github.io/styleguide/pyguide.html

But it has the drawback that it is very much focussed on "internal libs" in a monorepo rather than OSS packages for "external endusers" so there are a few places when it comes exposing symbols where to deviated from that guide (e.g. generally it says you should not create aliases, but I think for a public API it's needed).

Ah I was considering an alias but decided against it because the truss_chains already felt like an alias for things across public_api / definitions / framework (now). Can add though!

Good point, no alias in public_api, just directly to __init__.py then?

Sweet, that's what I have currently! Will merge as is then

marius-baseten · 2025-01-29T00:19:08Z

truss-chains/truss_chains/framework.py

-    try:
+    @classmethod
+    @abc.abstractmethod
+    def target_cls_type(cls) -> Type[definitions.ABCChainlet]:


This can be private though (and also others)?

marius-baseten · 2025-01-29T00:20:44Z

truss-chains/truss_chains/framework.py

+        }
+
+    @classmethod
+    def _load_module(cls, module_path: pathlib.Path) -> tuple[types.ModuleType, Loader]:


I trust you on this, I'm not gonna do a line-by-line review this time ;)

marius-baseten · 2025-01-29T00:21:20Z

truss-chains/truss_chains/framework.py

+    @classmethod
+    def no_entrypoint_error(cls, module_path: pathlib.Path) -> ValueError:
+        return ValueError(
+            f"No Model class in `{module_path}` inherits from {cls.target_cls_type()}."


Good call. It think the reason was the cyclic dependency issue... let's fix both this time.

nnarayen force-pushed the nikhil/python-dx-improvements branch 2 times, most recently from 70e26d1 to 0e73efd Compare January 28, 2025 20:51

nnarayen commented Jan 28, 2025

View reviewed changes

nnarayen force-pushed the nikhil/python-dx-improvements branch from 0e73efd to 7f7d449 Compare January 28, 2025 20:59

nnarayen commented Jan 28, 2025

View reviewed changes

nnarayen requested a review from marius-baseten January 28, 2025 21:03

nnarayen commented Jan 28, 2025

View reviewed changes

nnarayen mentioned this pull request Jan 28, 2025

Better example code for models using chains framework #1347

Merged

marius-baseten reviewed Jan 29, 2025

View reviewed changes

nnarayen added 2 commits January 29, 2025 10:31

Bug bash improvements to Python DX, better error messages

02c6009

Address comments

d88e360

nnarayen force-pushed the nikhil/python-dx-improvements branch from 2568a20 to d88e360 Compare January 29, 2025 15:31

marius-baseten approved these changes Jan 29, 2025

View reviewed changes

nnarayen merged commit 7812b1d into main Jan 29, 2025
5 checks passed

nnarayen deleted the nikhil/python-dx-improvements branch January 29, 2025 20:13

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Bug bash improvements to Python DX, better error messages #1346

Bug bash improvements to Python DX, better error messages #1346

nnarayen commented Jan 28, 2025 •

edited

Loading

nnarayen Jan 28, 2025

marius-baseten Jan 29, 2025

nnarayen Jan 28, 2025

marius-baseten Jan 29, 2025

nnarayen Jan 28, 2025

marius-baseten Jan 29, 2025

nnarayen Jan 28, 2025

marius-baseten Jan 29, 2025

nnarayen Jan 28, 2025

marius-baseten Jan 29, 2025

nnarayen Jan 29, 2025

marius-baseten Jan 29, 2025

nnarayen Jan 29, 2025

marius-baseten Jan 29, 2025

marius-baseten Jan 29, 2025

marius-baseten Jan 29, 2025

marius-baseten Jan 29, 2025

nnarayen Jan 29, 2025 •

edited

Loading

marius-baseten Jan 29, 2025

nnarayen Jan 29, 2025

marius-baseten Jan 29, 2025

nnarayen Jan 29, 2025

marius-baseten Jan 29, 2025

marius-baseten Jan 29, 2025

marius-baseten Jan 29, 2025

Bug bash improvements to Python DX, better error messages #1346

Bug bash improvements to Python DX, better error messages #1346

Conversation

nnarayen commented Jan 28, 2025 • edited Loading

🚀 What

💻 How

🔬 Testing

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

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

nnarayen Jan 29, 2025 • edited Loading

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

nnarayen commented Jan 28, 2025 •

edited

Loading

nnarayen Jan 29, 2025 •

edited

Loading