simonkurtz-msft / python-openai-loadbalancer Goto Github PK

The initial implementation focused on the synchronous flow. It wouldn't take much to expand load_balancer.py to add support for async, I suspect.

I have not yet focused on streaming. As I only change URL and Host header, I don't immediately expect issues with streaming, but it's to be tested.

openai_priority_loadbalancer is not fully typed yet. Prevent use of implicit or explicit any by using proper types.

Consider using numpy for a uniform distribution to selected an available backend. Part of the challenge is that the availableBackends array changes in size, making it difficult to have consistency in obtaining a uniform distribution.

Examples in the code/tests leverage using the managed identity authorization with via token_provider.
If the current solution supports connecting to endpoints with API keys, ti would be good to update documentation.

If the solution does not support connecting to endpoints via API keys at the moment, this is a feature request :)

Presently, the module only accounts for the retry-after header containing a value in seconds. Azure OpenAI Provisioned Throughput (PTU) deployments also pass back the retry-after-ms header containing a value in milliseconds. This value may be preferential to seconds.

From the PTU documentation:

*A 429 response indicates that the allocated PTUs are fully consumed at the time of the call. The response includes the retry-after-ms and retry-after headers that tell you the time to wait before the next call will be accepted. *

As such, there is only a longer delay when not using retry-after-ms when there is simply no PTU left to consume (overall and/or tokens-per-minute?). This may be an acceptable situation that may not surface too often.

One approach may be to convert everything internal to the module to milliseconds. These resources provide the starting point for the enhancement:

• https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/provisioned-get-started#handling-high-utilization
• openai/openai-python#1046

cc @kristapratico

Currently, the code only targets Azure OpenAI. It would be great to also cover OpenAI scenarios.

We keep track of the LastRetryAfter value in load_balanced.py. When no backends are available, and the last attempted backend possibly has a high retryAfter value, we would wait more than we need to as other backends may have become available again.

What I believe we should do then is return the nearest upcoming retry_after datetime property in the backends collection. The spec allows for either a datetime or seconds. I expect the httpx client to honor both.

Presently, the backends are model agnostic. That means that every model being used by the implementer of this code must reside on every Azure OpenAI instance that is defined in the backend. This could be limiting because it would require a lowest common denominator. Take these backends, for example:

• Backend 1 supports model A
• Backend 2 supports models A & B
• Backend 3 supports model A
• Backend 4 supports model B
• Backend 5 supports models A & B

Today, the backend pool can only use backends 2 and 5.

If the backend list could take model into consideration, the following would apply per model:

• Model A: backends 1, 2, 3, and 5
• Model B: backends 2, 4, and 5

I am interested to hear whether there is value in being able to specify backends per model or whether this is a potential solution in search of a problem.

The repo presently does not have any unit tests, but they are needed to mock handling of situations with (Azure) OpenAI backends.

Add python coverage as well.