Pubmatic multi inference

[Shantanu1058]

Thanks for submitting a PR to Triton!
Please go the the Preview tab above this description box and select the appropriate sub-template:

• PR description template for Triton Engineers
• PR description template for External Contributors

If you already created the PR, please replace this message with one of

• External contribution template
• Internal contribution template

and fill it out.

[greptile-apps]

Greptile Summary

This PR adds a Pubmatic-specific "multi-inference" feature to Triton Server: a new POST /v2/multi_infer HTTP endpoint that fans out a single request into parallel sub-inference calls across multiple models, with optional MySQL ODBC integration that loads per-campaign LightGBM model mappings from a database and transforms impression/campaign JSON into batched float tensors before inference.

• New /v2/multi_infer endpoint (src/multi_infer.cc): accepts either a generic {requests:[...]} body or a Pubmatic-specific {imps:[...]} body, dispatches sub-requests asynchronously, and aggregates results into a single HTTP response.
• MySQL ODBC connection pool (src/mysql_odbc_connection_pool.cc/.h): a fixed-size pool that pre-connects to a MySQL DSN at startup; a background thread (15-minute interval) refreshes a double-buffered campaign-to-model mapping from the database.
• Impression transform (src/transform.cc): converts impression+campaign JSON into row-major FP32 tensors using per-campaign feature sequences and categorical-to-index mappings fetched from MySQL.

Confidence Score: 2/5

Not ready to merge: the campaign mapping double-buffer has an unguarded data race that causes UB under concurrent load, and OnShardDone fires before the FINAL-flag check which can send a truncated HTTP reply for any streaming/multi-response model.

Two concurrency bugs affect core correctness of the new endpoint. The double-buffer swap in UpdateTritonModelsData clears the old map while request threads may still hold a raw pointer to it — undefined behaviour on every refresh cycle. The multi_infer shard callback calls OnShardDone before checking TRITONSERVER_RESPONSE_COMPLETE_FINAL, so any model producing more than one response fragment causes a premature reply. Additionally, a declared-but-undefined function in the ODBC header will cause a linker error, and the blocking Acquire() has no timeout to handle a stalled database.

src/mysql_odbc_connection_pool.cc (data race + blocking Acquire), src/multi_infer.cc (OnShardDone before FINAL flag), src/mysql_odbc_connection_pool.h (undefined function body)

Security Review

• Credentials in plaintext connection string (src/mysql_odbc_connection_pool.cc, BuildMySqlDriverConnectString): dsn_user_password is embedded verbatim in the ODBC PWD={...} connection string. If this string is ever captured in a crash dump, diagnostic log, or error message, the database password is exposed.
• Credentials in a plain JSON config file: config/database_config.sample.json stores dsnUserName and dsnUserPassword as plaintext strings; no guidance exists for file-mode hardening or secrets-manager integration.

Important Files Changed

Filename	Overview
src/multi_infer.cc	New file implementing POST /v2/multi_infer. Contains a logic bug where OnShardDone is called before the TRITONSERVER_RESPONSE_COMPLETE_FINAL flag is checked, risking premature HTTP replies for streaming/multi-response models.
src/mysql_odbc_connection_pool.cc	New ODBC connection pool and campaign model data layer. Has a data race in UpdateTritonModelsData, an indefinitely-blocking Acquire(), and a narrowing cast without range check in ParseApplicableCampaignIds.
src/mysql_odbc_connection_pool.h	Header for ODBC pool and feature mapping types. FetchLightgbmFeatureMappingMaxUpdateUnixSeconds is declared but never implemented.
src/transform.cc	Transforms impression+campaign JSON into batched float tensors. Logic appears sound; g_ready_model_names is written once at startup and guarded by an atomic flag.
src/database_config.cc	Loads JSON database config from /etc/triton-dmconfig.json with careful field parsing and range checking.
src/http_server.cc	Adds ScheduleInferAsync helper, FinalizeResponse json_only_out path, and ExtractFirstJsonOutput* methods. Routes POST /v2/multi_infer to HandleMultiInfer.
src/main.cc	Loads database config at startup, initializes ODBC pool, starts the 15-minute model refresh thread, and joins it on shutdown.
src/http_error_json.h	Lightweight helper to write {error:...} JSON into an evbuffer, replacing duplicate implementations. Correct and clean.
src/http_server_macros.h	Extracts shared HTTP response macros into a header. Functionally correct, but included inside a namespace block in multi_infer.cc.

Sequence Diagram

sequenceDiagram
    participant C as HTTP Client
    participant H as HTTPAPIServer
    participant MI as HandleMultiInfer
    participant T as GenerateImpsInferSlots
    participant TS as TRITONSERVER_InferAsync
    participant A as MultiInferAggregator
    participant DB as MySQL (ODBC)
    participant RT as RefreshThread (15min)

    RT->>DB: FetchLightgbmBtModelsMaxUpdateUnixSeconds
    DB-->>RT: timestamp
    RT->>DB: FetchLightgbmBtModelsForDc
    DB-->>RT: campaign-to-model mappings
    RT->>RT: swap double-buffer (g_triton_models_active)

    C->>H: "POST /v2/multi_infer {imps:[...]}"
    H->>MI: HandleMultiInfer(req)
    MI->>T: GenerateImpsInferSlots(doc, server)
    T->>T: read ActiveCampaignToFeatureMappings()
    T-->>MI: ImpsInferSlot[] + ImpRoutingTable
    loop for each model slot
        MI->>TS: TRITONSERVER_ServerInferAsync
        TS-->>A: InferResponseComplete callback
        A->>A: OnShardDone(slot, err, buf)
    end
    A->>A: WriteHttpReply (when all n_ slots done)
    A-->>C: "JSON response {imps:[{camps:[...]}]}"

Loading

_{Reviews (1): Last reviewed commit: "Removed odbc.ini file in Dockerfile" | Re-trigger Greptile}

[greptile-apps]

Data race: writer clears the buffer readers are still using

The double-buffer swap sequence is:

1. Write g_triton_campaign_feature_mappings[inactive]
2. Store inactive as the new active index
3. Clear g_triton_campaign_feature_mappings[idx]

Between steps 2 and 3, a request thread can call ActiveCampaignToFeatureMappings(), load the old idx, and start iterating over g_triton_campaign_feature_mappings[idx] in GenerateImpsInferSlots. Step 3 then calls .clear() on that same map while the request thread holds a raw pointer to it — this is a data race (undefined behaviour). A read-lock (e.g. std::shared_mutex) is needed to prevent the writer from clearing the old buffer while any reader is still inside it.

[greptile-apps]

OnShardDone called before FINAL flag is checked — premature reply for streaming models

OnShardDone (which increments done_count_ and potentially triggers WriteHttpReply) is called unconditionally for every response callback, including non-final ones. The TRITONSERVER_RESPONSE_COMPLETE_FINAL flag is only checked afterward (line 636). For any model that delivers more than one response fragment, done_count_ can reach n_ on the first (non-final) callback, causing the HTTP reply to be sent before all sub-requests have actually completed and before the final cleanup runs.

[greptile-apps]

FetchLightgbmFeatureMappingMaxUpdateUnixSeconds declared but never defined

This function is declared in the header but has no matching definition in mysql_odbc_connection_pool.cc. Any translation unit that includes this header and calls this function will produce a linker error. Either provide an implementation or remove the declaration.

[greptile-apps]

Acquire() blocks indefinitely — no timeout on the condition variable wait

cv_.wait(lk, predicate) waits forever until a free connection is available. If all connections are tied up in a stalled ODBC query (network partition, unresponsive DB), the calling thread hangs forever with no way to recover or return an error. Consider using cv_.wait_for(lk, timeout, predicate) and returning an empty PooledOdbcConnection on timeout.

[greptile-apps]

ParseApplicableCampaignIds: no range check before narrowing cast to int32_t

On 64-bit Linux where long is 64 bits, strtol can return values outside [INT32_MIN, INT32_MAX]. The unchecked cast to int32_t silently truncates them, potentially routing traffic to the wrong campaign model.

Suggested change

	const long v = std::strtol(tok.c_str(), &endptr, 10);
	if (endptr == tok.c_str() || *endptr != '\0') {
	return std::string("invalid campaign id token: ") + tok;
	}
	out->push_back(static_cast<int32_t>(v));
	const long v = std::strtol(tok.c_str(), &endptr, 10);
	if (endptr == tok.c_str() || *endptr != '\0') {
	return std::string("invalid campaign id token: ") + tok;
	}
	if (v < INT32_MIN || v > INT32_MAX) {
	return std::string("campaign id out of int32_t range: ") + tok;
	}
	out->push_back(static_cast<int32_t>(v));

[greptile-apps]

#include of a macro header inside a namespace block

http_server_macros.h is #included after namespace triton { namespace server { has been opened. Preprocessor directives are not namespace-scoped, so this works, but placing an #include inside an open namespace is confusing and violates common coding conventions. Move it to the top of the translation unit before any namespace is opened.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[greptile-apps]

Password embedded verbatim in ODBC connection string

BuildMySqlDriverConnectString embeds dsn_user_password directly as PWD={...}. If this string is ever logged or captured in a crash dump, the plaintext credential is exposed. Audit any future logging paths that might capture this string.