Skip to content

Commit

Permalink
reverted trace to span, with span being single function and trace bei…
Browse files Browse the repository at this point in the history 
…ng the collection.
  • Loading branch information
@djl11
djl11 committed Dec 19, 2024
1 parent 1934937 commit 47e9773
Show file tree
Hide file tree
Showing 3 changed files with 57 additions and 59 deletions.
54 changes: 26 additions & 28 deletions tests/test_evals/test_logging.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -979,7 +979,7 @@ def some_func(st):
entries = unify.get_logs(project="my_project")[0].entries

assert entries["trace"]["inputs"] == {"st": 0.5}
assert entries["trace"]["trace_name"] == "some_func"
assert entries["trace"]["span_name"] == "some_func"
assert (
entries["trace"]["code"].replace(" ", "").replace("\n", "")
== """
Expand All @@ -997,12 +997,11 @@ def some_func(st):
"",
)
)
assert len(entries["trace"]["child_traces"]) == 2
assert entries["trace"]["child_traces"][0]["trace_name"] == "inner_fn"
assert len(entries["trace"]["child_traces"][0]["child_traces"]) == 1
assert len(entries["trace"]["child_spans"]) == 2
assert entries["trace"]["child_spans"][0]["span_name"] == "inner_fn"
assert len(entries["trace"]["child_spans"][0]["child_spans"]) == 1
assert (
entries["trace"]["child_traces"][0]["child_traces"][0]["trace_name"]
== "deeper_fn"
entries["trace"]["child_spans"][0]["child_spans"][0]["span_name"] == "deeper_fn"
)


Expand All @@ -1027,7 +1026,7 @@ def some_func(a, b, c):
trace = logs[0].entries["trace"]
idx = trace["id"]
assert isinstance(idx, str)
assert trace["trace_name"] == "some_func"
assert trace["span_name"] == "some_func"
exec_time = trace["exec_time"]
assert isinstance(exec_time, float)
assert (
Expand All @@ -1042,7 +1041,7 @@ def some_func(a, b, c):
assert len(logs) == 1
trace = logs[0].entries["trace"]
assert trace["id"] == idx
assert trace["trace_name"] == "some_func"
assert trace["span_name"] == "some_func"
assert trace["exec_time"] == exec_time
assert (
trace["code"].replace(" ", "")
Expand All @@ -1068,7 +1067,7 @@ def some_func(a, b, c, d):
assert len(logs) == 1
entries = logs[0].entries
assert entries["trace"]["inputs"] == {"a": 1, "b": 2, "c": None, "d": 4}
assert entries["trace"]["trace_name"] == "some_func"
assert entries["trace"]["span_name"] == "some_func"
assert (
entries["trace"]["code"].replace(" ", "").replace("\n", "")
== """
Expand All @@ -1083,7 +1082,7 @@ def some_func(a, b, c, d):
"",
)
)
assert len(entries["trace"]["child_traces"]) == 0
assert len(entries["trace"]["child_spans"]) == 0

@unify.traced(prune_empty=True)
def some_func(a, b, c, d):
Expand All @@ -1094,7 +1093,7 @@ def some_func(a, b, c, d):
assert len(logs) == 2
entries = logs[1].entries
assert entries["trace"]["inputs"] == {"a": 1, "b": 2, "d": 4}
assert entries["trace"]["trace_name"] == "some_func"
assert entries["trace"]["span_name"] == "some_func"
assert (
entries["trace"]["code"].replace(" ", "").replace("\n", "")
== """
Expand All @@ -1109,7 +1108,7 @@ def some_func(a, b, c, d):
"",
)
)
assert "child_traces" not in entries["trace"]
assert "child_spans" not in entries["trace"]


def test_traced_within_log_context():
Expand Down Expand Up @@ -1145,13 +1144,12 @@ def some_func(st):
assert entries["a"] == "a"
assert entries["b"] == "b"
assert entries["trace"]["inputs"] == {"st": 0.5}
assert entries["trace"]["trace_name"] == "some_func"
assert len(entries["trace"]["child_traces"]) == 2
assert entries["trace"]["child_traces"][0]["trace_name"] == "inner_fn"
assert len(entries["trace"]["child_traces"][0]["child_traces"]) == 1
assert entries["trace"]["span_name"] == "some_func"
assert len(entries["trace"]["child_spans"]) == 2
assert entries["trace"]["child_spans"][0]["span_name"] == "inner_fn"
assert len(entries["trace"]["child_spans"][0]["child_spans"]) == 1
assert (
entries["trace"]["child_traces"][0]["child_traces"][0]["trace_name"]
== "deeper_fn"
entries["trace"]["child_spans"][0]["child_spans"][0]["span_name"] == "deeper_fn"
)


Expand Down Expand Up @@ -1195,11 +1193,11 @@ def some_func(st):
for i, log in enumerate(logs):
trace = log.entries["trace"]
assert trace["inputs"] == {"st": i / 100}
assert trace["trace_name"] == "some_func"
assert len(trace["child_traces"]) == 2
assert trace["child_traces"][0]["trace_name"] == "inner_fn"
assert len(trace["child_traces"][0]["child_traces"]) == 1
assert trace["child_traces"][0]["child_traces"][0]["trace_name"] == "deeper_fn"
assert trace["span_name"] == "some_func"
assert len(trace["child_spans"]) == 2
assert trace["child_spans"][0]["span_name"] == "inner_fn"
assert len(trace["child_spans"][0]["child_spans"]) == 1
assert trace["child_spans"][0]["child_spans"][0]["span_name"] == "deeper_fn"


@pytest.mark.asyncio
Expand Down Expand Up @@ -1235,11 +1233,11 @@ async def some_func(st):
for i, log in enumerate(logs):
trace = log.entries["trace"]
assert trace["inputs"] == {"st": i / 100}
assert trace["trace_name"] == "some_func"
assert len(trace["child_traces"]) == 2
assert trace["child_traces"][0]["trace_name"] == "inner_fn"
assert len(trace["child_traces"][0]["child_traces"]) == 1
assert trace["child_traces"][0]["child_traces"][0]["trace_name"] == "deeper_fn"
assert trace["span_name"] == "some_func"
assert len(trace["child_spans"]) == 2
assert trace["child_spans"][0]["span_name"] == "inner_fn"
assert len(trace["child_spans"][0]["child_spans"]) == 1
assert trace["child_spans"][0]["child_spans"][0]["span_name"] == "deeper_fn"


if __name__ == "__main__":
Expand Down
60 changes: 30 additions & 30 deletions unify/evals/logging.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -225,104 +225,104 @@ def wrapped(*args, **kwargs):
log_token = None if ACTIVE_LOG.get() else ACTIVE_LOG.set([unify.log()])
t1 = time.perf_counter()
ts = datetime.datetime.utcnow().isoformat()
if not TRACE.get():
if not SPAN.get():
RUNNING_TIME.set(t1)
signature = inspect.signature(fn)
bound_args = signature.bind(*args, **kwargs)
bound_args.apply_defaults()
inputs = bound_args.arguments
new_trace = {
new_span = {
"id": str(uuid.uuid4()),
"parent_trace_id": (None if not TRACE.get() else TRACE.get()["id"]),
"trace_name": fn.__name__,
"parent_span_id": (None if not SPAN.get() else SPAN.get()["id"]),
"span_name": fn.__name__,
"exec_time": None,
"timestamp": ts,
"offset": round(
0.0 if not TRACE.get() else t1 - RUNNING_TIME.get(),
0.0 if not SPAN.get() else t1 - RUNNING_TIME.get(),
2,
),
"code": inspect.getsource(fn),
"inputs": inputs,
"outputs": None,
"errors": None,
"child_traces": [],
"child_spans": [],
}
if inspect.ismethod(fn) and hasattr(fn.__self__, "endpoint"):
new_trace["endpoint"] = fn.__self__.endpoint
token = TRACE.set(new_trace)
new_span["endpoint"] = fn.__self__.endpoint
token = SPAN.set(new_span)
result = None
try:
result = fn(*args, **kwargs)
return result
except Exception as e:
new_trace["errors"] = str(e)
new_span["errors"] = str(e)
raise e
finally:
t2 = time.perf_counter()
exec_time = t2 - t1
TRACE.get()["exec_time"] = round(exec_time, 2)
TRACE.get()["outputs"] = None if result is None else result
SPAN.get()["exec_time"] = round(exec_time, 2)
SPAN.get()["outputs"] = None if result is None else result
# ToDo: ensure there is a global log set upon the first trace,
# and removed on the last
trace = TRACE.get()
trace = SPAN.get()
if prune_empty:
trace = _prune_dict(trace)
unify.add_log_entries(trace=trace, overwrite=True)
if token.old_value is token.MISSING:
TRACE.reset(token)
SPAN.reset(token)
else:
TRACE.reset(token)
TRACE.get()["child_traces"].append(new_trace)
SPAN.reset(token)
SPAN.get()["child_spans"].append(new_span)
if log_token:
ACTIVE_LOG.set([])

async def async_wrapped(*args, **kwargs):
t1 = time.perf_counter()
if not TRACE.get():
if not SPAN.get():
RUNNING_TIME.set(t1)
signature = inspect.signature(fn)
bound_args = signature.bind(*args, **kwargs)
bound_args.apply_defaults()
inputs = bound_args.arguments
new_trace = {
new_span = {
"id": str(uuid.uuid4()),
"parent_trace_id": (None if not TRACE.get() else TRACE.get()["id"]),
"trace_name": fn.__name__,
"parent_span_id": (None if not SPAN.get() else SPAN.get()["id"]),
"span_name": fn.__name__,
"exec_time": None,
"offset": round(
0.0 if not TRACE.get() else t1 - RUNNING_TIME.get(),
0.0 if not SPAN.get() else t1 - RUNNING_TIME.get(),
2,
),
"inputs": inputs,
"outputs": None,
"errors": None,
"child_traces": [],
"child_spans": [],
}
token = TRACE.set(new_trace)
token = SPAN.set(new_span)
# capture the arguments here
result = None
try:
result = await fn(*args, **kwargs)
return result
except Exception as e:
new_trace["errors"] = str(e)
new_span["errors"] = str(e)
raise e
finally:
t2 = time.perf_counter()
exec_time = t2 - t1
TRACE.get()["exec_time"] = round(exec_time, 2)
TRACE.get()["outputs"] = None if result is None else result
SPAN.get()["exec_time"] = round(exec_time, 2)
SPAN.get()["outputs"] = None if result is None else result
if token.old_value is token.MISSING:
if ACTIVE_LOG.get():
trace = TRACE.get()
trace = SPAN.get()
if prune_empty:
trace = _prune_dict(trace)
unify.add_log_entries(trace=trace, overwrite=True)
else:
unify.log(trace=TRACE.get())
TRACE.reset(token)
unify.log(trace=SPAN.get())
SPAN.reset(token)
else:
TRACE.reset(token)
TRACE.get()["child_traces"].append(new_trace)
SPAN.reset(token)
SPAN.get()["child_spans"].append(new_span)

return wrapped if not inspect.iscoroutinefunction(fn) else async_wrapped
2 changes: 1 addition & 1 deletion unify/evals/utils/logging.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@
PARAMS_NEST_LEVEL = ContextVar("params_nest_level", default=0)

# span
TRACE = ContextVar("span", default={})
SPAN = ContextVar("span", default={})
RUNNING_TIME = ContextVar("running_time", default=0.0)


Expand Down

0 comments on commit 47e9773

Please sign in to comment.