Yesterday we demonstrated how switching from rate limiting by requests per minute (RPM) to rate limiting by tokens per minute (TPM) allowed us to smooth our Vertex AI API traffic profile for the Gemini model family, allowing us to maximize our throughput for our large historical backfile jobs, while staying just under the endpoint quotas. (No data is used to train or tune any model). In our original work we used a dynamic daemon cost estimator that monitored the realtime inference token counts for each job and devised a formula, updated every few seconds, that precisely estimated the likely input tokens a given transcript set would consume, based on comparing it to a rolling window assessment of recent jobs of similar language, byte, character and word counts.
The end result was a remarkably steady inference rate, but which was still somewhat jagged and not perfectly smooth. Given that Gemini supports a "countTokens" API call that precisely estimates the number of input tokens a given inference request will consume, we received several queries as to whether switching from our client-side token estimation to using the official Gemini countTokens endpoint would further smooth these results and remove much of the final jaggedness of the graph.
You can see the results above. The left third of the graph represents the inference TPM using our client-side estimator algorithm, while the right third uses the Gemini countTokens API (the middle third is the transition period mixing jobs under the new workflow and LROs using the original estimator completed). The Gemini countTokens API definately tightens the variability band, allowing us to move towards pushing ever-closer to the endpoint quota threshold, but overall is not an enormous improvement over our rolling dynamic estimator. Thus, for large intermixed workflows of similar inputs, client-side rolling cost estimators yield nearly identical results to the countTokens API, with the API being useful only for workflows that require optimizing the variability band to absolutely maximize throughput at all costs.
For those interested, the way we use the cost estimates (either from our dynamic estimator or the countTokens API) is that we divide 60 seconds by our per-model TPM quota max (60/TPM_MAX) to convert to TPS (tokens per second) and then for each inference job to be submitted to Gemini, we multiply it by this number and advance a shared globally-locked nanosecond-resolution timestamp to the new position. Whichever job is ready next waits until this timestamp before executing, advancing it based on the number of submitted tokens and so on. This evenly spaces our submission rate through the course of each minute, rather than bunching up all submissions at the start of each minute.