phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-17 11:02:36 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fixes

Browse source
This commit is contained in:
Dinesh Yeduguru 2024-11-26 15:41:08 -08:00
parent af8a1fe5b3
commit b3e149334a
3 changed files with 22 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

10
llama_stack/distribution/tracing.py
View file

@ -155,17 +155,19 @@ def trace_protocol(cls: Type[T]) -> Type[T]:
traced_methods = {}
for parent in cls_child.__mro__[1:]: # Skip the class itself
for name, method in vars(parent).items():
if inspect.isfunction(method) and method._trace_input:
traced_methods[name] = method._trace_input
if inspect.isfunction(method) and getattr(
method, "_trace_input", None
): # noqa: B009
traced_methods[name] = getattr(method, "_trace_input") # noqa: B009
# Trace child class methods if their name matches a traced parent method
for name, method in vars(cls_child).items():
if inspect.isfunction(method) and not name.startswith("_"):
if name in traced_methods:
# Copy the trace configuration from the parent
method._trace_input = traced_methods[name]
setattr(method, "_trace_input", traced_methods[name]) # noqa: B010
cls_child.__dict__[name] = trace_method(method)
setattr(cls_child, name, trace_method(method)) # noqa: B010
# Set the new __init_subclass__
cls.__init_subclass__ = classmethod(__init_subclass__)