Skip to content

refactor(openai): move endpoint hooks to new internal file, add utils file#5865

Merged
Yun-Kim merged 8 commits into
1.xfrom
yunkim/refactor-openai-endpoints
May 19, 2023
Merged

refactor(openai): move endpoint hooks to new internal file, add utils file#5865
Yun-Kim merged 8 commits into
1.xfrom
yunkim/refactor-openai-endpoints

Conversation

@Yun-Kim

@Yun-Kim Yun-Kim commented May 15, 2023

Copy link
Copy Markdown
Contributor

This PR refactors the openAI integration In order to more easily support multiple endpoints other than the existing Completions/ChatCompletions/Embeddings.

This PR makes the following refactors to the OpenAI integration:

  • The _EndpointHook class and its subclasses that we use to trace each endpoint are contained in their own internal ddtrace.contrib.openai._endpoint_hooks.py file.
  • The _EndpointHook base class's handle_request() is no longer an abstract method, and instead now contains two refactored abstract methods _pre_response() and _post_response() which happen before and after the response/error is returned from the openAI API. This allows more granularity and clarity on what the traced endpoint subclasses are implementing, especially since many endpoints will be able to share the same code for _pre_response() or _post_response().
  • All endpoint subclasses must now have a _default_name, which is the default name used for the openai.response.operation tag.
  • All endpoint subclasses will have a _request_tag_attrs, which is a list of request params that will be automatically tagged by _EndpointHook._record_request().
  • The _est_tokens helper function has been moved into a ddtrace.contrib.openai._utils.py file. More helper functions are expected to be contained in this new openai utils file when additional endpoints will be implemented.

No functionality has been changed so no additional testing is required.

Checklist

  • Change(s) are motivated and described in the PR description.
  • Testing strategy is described if automated tests are not included in the PR.
  • Risk is outlined (performance impact, potential for breakage, maintainability, etc).
  • Change is maintainable (easy to change, telemetry, documentation).
  • Library release note guidelines are followed.
  • Documentation is included (in-code, generated user docs, public corp docs).
  • OPTIONAL: PR description includes explicit acknowledgement of the performance implications of the change as reported in the benchmarks PR comment.

Reviewer Checklist

  • Title is accurate.
  • No unnecessary changes are introduced.
  • Description motivates each change.
  • Avoids breaking API changes unless absolutely necessary.
  • Testing strategy adequately addresses listed risk(s).
  • Change is maintainable (easy to change, telemetry, documentation).
  • Release note makes sense to a user of the library.
  • Reviewer has explicitly acknowledged and discussed the performance implications of this PR as reported in the benchmarks PR comment.

@Yun-Kim Yun-Kim added the changelog/no-changelog A changelog entry is not required for this PR. label May 15, 2023
@Yun-Kim Yun-Kim force-pushed the yunkim/refactor-openai-endpoints branch 2 times, most recently from 93722ba to d3a19df Compare May 15, 2023 00:22
@pr-commenter

pr-commenter Bot commented May 15, 2023

Copy link
Copy Markdown

Benchmarks

Comparing candidate commit b4668f8 in PR branch yunkim/refactor-openai-endpoints with baseline commit 826e601 in branch 1.x.

Found 0 performance improvements and 0 performance regressions! Performance is the same for 94 cases.

@Yun-Kim Yun-Kim marked this pull request as ready for review May 17, 2023 19:30
@Yun-Kim Yun-Kim requested a review from a team as a code owner May 17, 2023 19:30
@Yun-Kim Yun-Kim force-pushed the yunkim/refactor-openai-endpoints branch from d3a19df to 296784d Compare May 17, 2023 19:42
emmettbutler
emmettbutler previously approved these changes May 18, 2023
Comment thread ddtrace/contrib/openai/_endpoint_hooks.py
@Yun-Kim Yun-Kim requested a review from emmettbutler May 18, 2023 16:03
@Yun-Kim Yun-Kim enabled auto-merge (squash) May 19, 2023 14:33
@Yun-Kim Yun-Kim merged commit eb16e8e into 1.x May 19, 2023
@Yun-Kim Yun-Kim deleted the yunkim/refactor-openai-endpoints branch May 19, 2023 15:58
Yun-Kim added a commit that referenced this pull request Jun 15, 2023
## Tips for reviewers

**note: 99% of the line changes in this PR are non-code, in the form of
snapshot files, riot lockfiles, test data (audio/image files), and test
cassette files used to mock request/responses with the OpenAI API, which
store binary image/audio file data in `.yaml` files that Github counts
as lines of code.**

The major files that reviewers should focus on are:
- `ddtrace/contrib/openai/patch.py` - outlines which endpoint hooks are
added, minor changes to overall patching to minimize coupling with
endpoint hook information.
- `ddtrace/contrib/openai/_endpoint_hooks.py` - implementation of each
endpoint hook, as well as minor refactor to minimize code duplication
between endpoint hooks
- `tests/contrib/openai/test_openai.py`. - shows how each endpoint is
expected to be used.

## Summary
This PR adds tracing support for the remaining endpoints not implemented
by #5488. This includes:
- Model (list, retrieve)
- Edits (create)
- Images (create, edit, create variation)
- Audio (transcribe, translate)
- Files (list, create, delete, retrieve, retrieve contents)
- Fine-tunes (list, create, retrieve, cancel, list events, delete fine
tuned model)
- Moderations (create)

Changes in this PR include:
- Modifying the design of the endpoint hooks to minimize duplicate code,
and adding new endpoint hooks for the above endpoints.
- Ensuring rate-limit related metrics are tagged as numeric metrics
rather than string tags.
- Ensure user's OpenAI API key is set as a tag in multiple use scenarios
(set in env var, set as attribute in `openai`, directly added as a
request param)
- Adding and correcting names of attributes extracted from `openai`
(e.g. `organization`, `api_type`)
- Simplifying and separating patch logic from endpoint-specific tracing
logic (only operationID is propagated from endpoint hook level to patch
level)
- Moving some global integration-level tags to request-level (e.g.
`openai.request.model`, `openai.request.endpoint`, added
`openai.request.method`)


## OpenAI Integration Design

### Overall Patch Implementation

**note**: the overall patch implementation has not changed; this is
simply a refresher on how the overall integration works.

The openai integration features patching via generators - each of the
openAI API's endpoint methods (ex: `openai.Completions.create()`) is
wrapped by `_patched_endpoint(endpoint_hook, ...)`, which itself takes a
corresponding endpoint hook as an argument. A rough description of what
happens when `_patched_endpoint()` is called is as follows:
- `_patched_endpoint(endpoint_hook, ...)` starts a trace and starts a
generator `_traced_endpoint(endpoint_hook, ...)`
- `_traced_endpoint(endpoint_hook, ...)` starts
`endpoint_hook().handle_request(...)`, which performs endpoint-specific
tracing logic, and yields back to `_patched_endpoint()` once the request
has been processed by the endpoint hook.
- `_patched_endpoint()` runs the underlying openai API method, then
yields back the response/error back to `_traced_endpoint()`, which in
turn yields the response/error back to the `endpoint_hook`.
- `endpoint_hook()` performs the endpoint-specific tracing logic for the
response/error, then finishes the trace.

### Endpoint Hook Design

The internal implementation design of the endpoint hooks have been
modified (building off of the refactor from #5865) to minimize code
duplication. This means that each endpoint hook now stores the
following:
- `_request_arg_params` and `_request_kwarg_params`, which are tuples
containing the arg/kwarg signature of the underlying openAI API endpoint
method. These tuples are used in `_EndpointHook._record_request(...)` to
add request arg/kwarg parameters as span tags/metrics.
- `ENDPOINT_NAME`, `REQUEST_TYPE`, `OPERATION_ID` constants, which
reflect the base endpoint name (e.g. `/completions`), http request type
(e.g. `POST`), and operation ID as specified by the openAI API
specifcations (e.g. `createCompletion`). The operationID is used as the
span resource name, and the endpoint/request_type values are added as
span tags.
- Each endpoint hook also features (optional) a `_record_request(...)`
and `_record_response(...)` to add any endpoint-specific span tagging
logic.

## Testing

The testing for this PR involves snapshot testing for each endpoint, as
well as `vcrpy` cassettes used to mock request/responses to OpenAI.

## Checklist

- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/contributing.html#Release-Note-Guidelines)
are followed.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] OPTIONAL: PR description includes explicit acknowledgement of the
performance implications of the change as reported in the benchmarks PR
comment.

## Reviewer Checklist

- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer has explicitly acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment.

---------

Co-authored-by: Munir Abdinur <munir.abdinur@datadoghq.com>
Co-authored-by: Federico Mon <federico.mon@datadoghq.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

changelog/no-changelog A changelog entry is not required for this PR.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants