Skip to content

Commit

Permalink
[trace][pfs] Support list line runs with line run id (#2822)
Browse files Browse the repository at this point in the history
# Description

Support query with `line_run_ids` for PFS API: `/v1.0/LineRuns/list`.

# All Promptflow Contribution checklist:
- [x] **The pull request does not introduce [breaking changes].**
- [ ] **CHANGELOG is updated for new features, bug fixes or other
significant changes.**
- [x] **I have read the [contribution guidelines](../CONTRIBUTING.md).**
- [ ] **Create an issue and link to the pull request to get dedicated
review from promptflow team. Learn more: [suggested
workflow](../CONTRIBUTING.md#suggested-workflow).**

## General Guidelines and Best Practices
- [x] Title of the pull request is clear and informative.
- [x] There are a small number of commits, each of which have an
informative message. This means that previously merged commits do not
appear in the history of the PR. For more information on cleaning up the
commits in your PR, [see this
page](https://github.com/Azure/azure-powershell/blob/master/documentation/development-docs/cleaning-up-commits.md).

### Testing Guidelines
- [x] Pull request includes test coverage for the included changes.

---------

Co-authored-by: Kaiyi <[email protected]>
  • Loading branch information
zhengfeiwang and RealKai42 authored Apr 16, 2024
1 parent 095129b commit 270ba81
Show file tree
Hide file tree
Showing 10 changed files with 83 additions and 62 deletions.
2 changes: 1 addition & 1 deletion .pre-commit-config.yaml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
# See https://pre-commit.com for more information
# See https://pre-commit.com/hooks.html for more hooks

exclude: '(^docs/)|flows|scripts|src/promptflow-azure/promptflow/azure/_restclient/|src/promptflow-core/promptflow/core/_connection_provider/_models/|src/promptflow-azure/promptflow/azure/_models/|src/promptflow/tests/test_configs|src/promptflow-tools'
exclude: '(^docs/)|flows|scripts|src/promptflow-azure/promptflow/azure/_restclient/|src/promptflow-core/promptflow/core/_connection_provider/_models/|src/promptflow-azure/promptflow/azure/_models/|src/promptflow/tests/test_configs|src/promptflow-tools|src/promptflow-devkit/promptflow/_sdk/_service/static'

repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
Expand Down
3 changes: 3 additions & 0 deletions src/promptflow-devkit/promptflow/_sdk/_orm/trace.py
Original file line number Diff line number Diff line change
Expand Up @@ -190,6 +190,7 @@ def list(
experiments: typing.Optional[typing.List[str]] = None,
trace_ids: typing.Optional[typing.List[str]] = None,
session_id: typing.Optional[str] = None,
line_run_ids: typing.Optional[typing.List[str]] = None,
) -> typing.List["LineRun"]:
with trace_mgmt_db_session() as session:
query = session.query(LineRun)
Expand All @@ -201,6 +202,8 @@ def list(
query = query.filter(LineRun.experiment.in_(experiments))
elif trace_ids is not None:
query = query.filter(LineRun.trace_id.in_(trace_ids))
elif line_run_ids is not None:
query = query.filter(LineRun.line_run_id.in_(line_run_ids))
elif session_id is not None:
query = query.filter(LineRun.session_id == session_id)
query = query.order_by(LineRun.start_time.desc())
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
list_line_run_parser.add_argument("run", type=str, required=False)
list_line_run_parser.add_argument("experiment", type=str, required=False)
list_line_run_parser.add_argument("trace_ids", type=str, required=False)
list_line_run_parser.add_argument("line_run_ids", type=str, required=False)


# use @dataclass for strong type
Expand All @@ -32,6 +33,7 @@ class ListLineRunParser:
experiments: typing.Optional[typing.List[str]] = None
trace_ids: typing.Optional[typing.List[str]] = None
session_id: typing.Optional[str] = None
line_run_ids: typing.Optional[typing.List[str]] = None

@staticmethod
def _parse_string_list(value: typing.Optional[str]) -> typing.Optional[typing.List[str]]:
Expand All @@ -48,6 +50,7 @@ def from_request() -> "ListLineRunParser":
experiments=ListLineRunParser._parse_string_list(args.experiment),
trace_ids=ListLineRunParser._parse_string_list(args.trace_ids),
session_id=args.session,
line_run_ids=ListLineRunParser._parse_string_list(args.line_run_ids),
)


Expand Down Expand Up @@ -94,5 +97,6 @@ def get(self):
experiments=args.experiments,
trace_ids=args.trace_ids,
session_id=args.session_id,
line_run_ids=args.line_run_ids,
)
return [line_run._to_rest_object() for line_run in line_runs]

Large diffs are not rendered by default.

Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,12 @@
<head>
<meta charset="UTF-8" />
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
<link rel="icon" href="/v1.0/ui/traces/assets/icon-IVYk8x5p.svg" media="(prefers-color-scheme:no-preference)">
<link rel="icon" href="/v1.0/ui/traces/assets/icon_for_dark-3C8HbOu4.svg" media="(prefers-color-scheme:dark)">
<link rel="icon" href="/v1.0/ui/traces/assets/icon-IVYk8x5p.svg" media="(prefers-color-scheme:light)">
<link rel="icon" href="/v1.0/ui/traces/assets/icon-oxB8RKmI.svg" media="(prefers-color-scheme:no-preference)">
<link rel="icon" href="/v1.0/ui/traces/assets/icon_for_dark-7fdboOJT.svg" media="(prefers-color-scheme:dark)">
<link rel="icon" href="/v1.0/ui/traces/assets/icon-oxB8RKmI.svg" media="(prefers-color-scheme:light)">
<title>Trace View</title>
<script type="module" crossorigin src="/v1.0/ui/traces/assets/index-M7U6jpSX.js"></script>

<script type="module" crossorigin src="/v1.0/ui/traces/assets/index-L3TKeGx-.js"></script>
</head>
<body>
<div id="root"></div>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,7 @@ def list_line_runs(
experiments: typing.Optional[typing.Union[str, typing.List[str]]] = None,
trace_ids: typing.Optional[typing.Union[str, typing.List[str]]] = None,
session_id: typing.Optional[str] = None,
line_run_ids: typing.Optional[typing.Union[str, typing.List[str]]] = None,
) -> typing.List[LineRun]:
# ensure runs, experiments, and trace_ids are list of string
if isinstance(runs, str):
Expand All @@ -170,6 +171,8 @@ def list_line_runs(
experiments = [experiments]
if isinstance(trace_ids, str):
trace_ids = [trace_ids]
if isinstance(line_run_ids, str):
line_run_ids = [line_run_ids]

# currently we list parent line runs first, and query children for each
# this will query SQLite for N+1 times (N is the number of parent line runs)
Expand All @@ -180,6 +183,7 @@ def list_line_runs(
experiments=experiments,
trace_ids=trace_ids,
session_id=session_id,
line_run_ids=line_run_ids,
)
line_runs = []
for obj in orm_line_runs:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -192,3 +192,10 @@ def test_list_line_run_with_session_id(self, pfs_op: PFSOperations, mock_collect
persist_a_span(collection=mock_collection, custom_attributes=custom_attributes)
line_runs = pfs_op.list_line_runs(session_id=mock_session_id).json
assert len(line_runs) == 1

def test_list_line_run_with_line_run_id(self, pfs_op: PFSOperations, mock_collection: str) -> None:
mock_line_run_id = str(uuid.uuid4())
custom_attributes = {SpanAttributeFieldName.LINE_RUN_ID: mock_line_run_id}
persist_a_span(collection=mock_collection, custom_attributes=custom_attributes)
line_runs = pfs_op.list_line_runs(line_run_ids=[mock_line_run_id]).json
assert len(line_runs) == 1
3 changes: 3 additions & 0 deletions src/promptflow-devkit/tests/sdk_pfs_test/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -248,6 +248,7 @@ def list_line_runs(
runs: Optional[List[str]] = None,
trace_ids: Optional[List[str]] = None,
session_id: Optional[str] = None,
line_run_ids: Optional[List[str]] = None,
):
query_string = {}
if collection is not None:
Expand All @@ -256,6 +257,8 @@ def list_line_runs(
query_string["run"] = ",".join(runs)
if trace_ids is not None:
query_string["trace_ids"] = ",".join(trace_ids)
if line_run_ids is not None:
query_string["line_run_ids"] = ",".join(line_run_ids)
if session_id is not None:
query_string["session"] = session_id
response = self._client.get(
Expand Down

0 comments on commit 270ba81

Please sign in to comment.