Debug why streamining hangs

truss/templates/server/common/truss_server.py

@@ -135,8 +135,10 @@ async def predict(
         model: ModelWrapper = self._safe_lookup_model(model_name)
 
         self.check_healthy(model)
-        # TODO: warn (or skip tracing), if no parent info was found?
-        trace_ctx = otel_propagate.extract(request.headers)
+
+        print(f"############\n{request.headers}")
+
+        trace_ctx = otel_propagate.extract(request.headers) or None
         # This is the top-level span in the truss-server, so we set the context here.
         # Nested spans "inherit" context automatically.
         with self._tracer.start_as_current_span(

truss/templates/server/model_wrapper.py

@@ -203,7 +203,7 @@ def try_load(self):
             retry(
                 self._model.load,
                 NUM_LOAD_RETRIES,
-                self._logger.warn,
+                self._logger.warning,
                 "Failed to load model.",
                 gap_seconds=1.0,
             )
@@ -301,6 +301,7 @@ async def write_response_to_queue(
                 )
             finally:
                 await queue.put(None)
+                print("@@@@@@@@ done write_response_to_queue")
 
     async def _gather_generator(self, response: Any, span: trace.Span) -> str:
         # In the case of gathering, it might make more sense to apply the postprocess
@@ -345,6 +346,7 @@ async def _stream_with_background_task(
             self.write_response_to_queue(response_queue, async_generator, span)
         )
         # Defer the release of the semaphore until the write_response_to_queue task.
+        print("@@@@@@@@ add delayed release")
         gen_task.add_done_callback(lambda _: release_and_end)
 
         # The gap between responses in a stream must be < streaming_read_timeout