Review comments

truss-chains/truss_chains/code_gen.py

@@ -511,9 +511,7 @@ def _gen_truss_chainlet_file(
 # Truss Gen ############################################################################
 
 
-def _make_requirements(
-    image: definitions.DockerImage, use_local_chains_src: bool
-) -> list[str]:
+def _make_requirements(image: definitions.DockerImage) -> list[str]:
     """Merges file- and list-based requirements and adds truss git if not present."""
     pip_requirements: set[str] = set()
     if image.pip_requirements_file:
@@ -583,9 +581,7 @@ def _make_truss_config(
     config.runtime.predict_concurrency = compute.predict_concurrency
     # Image.
     _inplace_fill_base_image(chains_config.docker_image, config)
-    pip_requirements = _make_requirements(
-        chains_config.docker_image, use_local_chains_src
-    )
+    pip_requirements = _make_requirements(chains_config.docker_image)
     # TODO: `pip_requirements` will add server requirements which give version
     #  conflicts. Check if that's still the case after relaxing versions.
     # config.requirements = pip_requirements

truss-chains/truss_chains/definitions.py

@@ -621,7 +621,6 @@ class GenericRemoteException(Exception): ...
 class PushOptions(SafeModelNonSerializable):
     chain_name: str
     only_generate_trusses: bool = False
-    use_local_chains_src: bool = False
 
 
 class PushOptionsBaseten(PushOptions):
@@ -662,3 +661,7 @@ class PushOptionsLocalDocker(PushOptions):
     # is unset. Additionally, if local docker containers make calls to models deployed
     # on baseten, a real API key must be provided (i.e. the default must be overridden).
     baseten_chain_api_key: str = "docker_dummy_key"
+    # If enabled, chains code is copied from the local package into `/app/truss_chains`
+    # in the docker image (which takes precedence over potential pip/site-packages).
+    # This should be used for integration tests or quick local dev loops.
+    use_local_chains_src: bool = False

truss-chains/truss_chains/remote.py

@@ -87,11 +87,9 @@ def _push_service_docker(
     options: definitions.PushOptionsLocalDocker,
     port: int,
 ) -> None:
-    truss_handle = truss_build.load(str(truss_dir))
-    truss_handle.add_secret(
-        definitions.BASETEN_API_SECRET_NAME, options.baseten_chain_api_key
-    )
-    truss_handle.docker_run(
+    th = truss_handle.TrussHandle(truss_dir)
+    th.add_secret(definitions.BASETEN_API_SECRET_NAME, options.baseten_chain_api_key)
+    th.docker_run(
         local_port=port,
         detach=True,
         wait_for_server_ready=True,
@@ -319,6 +317,12 @@ def __init__(
         self._options = options
         self._gen_root = gen_root or pathlib.Path(tempfile.gettempdir())
 
+    @property
+    def _use_local_chains_src(self) -> bool:
+        if isinstance(self._options, definitions.PushOptionsLocalDocker):
+            return self._options.use_local_chains_src
+        return False
+
     def generate_chainlet_artifacts(
         self,
         entrypoint: Type[definitions.ABCChainlet],
@@ -340,7 +344,7 @@ def generate_chainlet_artifacts(
                 self._options.chain_name,
                 chainlet_descriptor,
                 model_name,
-                self._options.use_local_chains_src,
+                self._use_local_chains_src,
             )
             artifact = b10_types.ChainletArtifact(
                 truss_dir=chainlet_dir,

truss/templates/server.Dockerfile.jinja

@@ -79,6 +79,7 @@ COPY ./{{config.data_dir}} /app/data
 COPY ./server /app
 
 {%- if use_local_chains_src %}
+{# This path takes precedence over site-packages. #}
 COPY ./truss_chains /app/truss_chains
 {%- endif %}