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 Pubmatic-specific multi-inference extensions to the Triton HTTP server: a generic POST /v2/multi_infer endpoint that fans out a JSON array of per-model inference requests in parallel, and a POST /v2/predict endpoint that routes impression/campaign data through a MySQL ODBC-backed feature-mapping pipeline before dispatching batched FP32 tensors to LightGBM models.

• New HTTP endpoints (src/multi_infer.cc): HandleMultiInfer and HandlePredict each build a MultiInferAggregator that collects sub-request results and writes a single HTTP reply when all shards complete; a double-buffer pattern in mysql_odbc_connection_pool.cc refreshes campaign-to-model mappings from the DB every 15 minutes without blocking request threads.
• Shared infrastructure (src/http_error_json.h, src/http_server_macros.h): duplicated EVBufferAddErrorJson and RETURN_AND_RESPOND_* macros are extracted into headers consumed by both http_server.cc and sagemaker_server.cc, and the scheduling boilerplate is factored into ScheduleInferAsync.
• Startup/lifecycle (src/main.cc): LoadTritonDmDatabaseConfigAtStartup reads /etc/triton-dmconfig.json, initializes the ODBC connection pool, and (when TRITON_ENABLE_MYSQL_ODBC is set) starts a background refresh thread; a corresponding std::atexit handler tears the pool down on process exit.

Confidence Score: 2/5

Not safe to merge: both new HTTP handlers contain a use-after-free that can be triggered in production whenever Triton rejects a shard's async submission mid-loop, and the startup sequence can abandon the DB refresh thread with its pool freed beneath it.

The partial-scheduling error paths in HandleMultiInfer and HandlePredict free MultiInferShardRequest objects whose InferResponseComplete callbacks are still in-flight in the Triton thread pool, producing dangling-pointer dereferences. Separately, StartTritonModelsRefreshThread runs before the FAIL_IF_ERR call that may exit(1) without joining it, causing the atexit handler to destroy the ODBC pool while the thread is mid-query. Both defects are in the hot path of the new endpoints.

Files Needing Attention: src/multi_infer.cc (both HandleMultiInfer and HandlePredict scheduling loops) and src/main.cc (StartTritonModelsRefreshThread / InitializeReadyModelNames ordering).

Important Files Changed

Filename	Overview
src/multi_infer.cc	New file implementing HandleMultiInfer and HandlePredict HTTP endpoints; contains a use-after-free in both handlers when ScheduleInferAsync fails for a non-first shard slot, because already-in-flight shards are freed before their async callbacks complete.
src/main.cc	Adds ODBC pool startup, refresh thread lifecycle, and DatabaseConfig loading; StartTritonModelsRefreshThread is called before the FAIL_IF_ERR that can exit(1), leaving the thread running when the atexit handler frees the pool.
src/mysql_odbc_connection_pool.cc	New file implementing ODBC connection pool, feature-mapping DB fetch, and double-buffer campaign-to-model lookup; Acquire now correctly uses wait_for with a 30-second timeout; range check on ParseApplicableCampaignIds is present.
src/transform.cc	New file building FP32 feature tensors from imps/camps JSON for BT model inference; ready-model snapshot is populated once at startup and used lock-free thereafter; logic looks correct for the static-model deployment assumption.
src/http_server.cc	Refactors shared scheduling logic into ScheduleInferAsync, adds FillMultiInferSlotTritonRequest / FillImpsTritonRequest helpers, and fixes a decompression evbuffer leak; changes look correct for the single-request (HandleInfer) path.
src/http_server.h	Adds pause_http_request / register_fini_cancel_hook constructor flags to InferRequestClass and new multi-infer handler declarations; the virtual FinalizeResponse signature change is backward-compatible with the default argument.

Sequence Diagram

sequenceDiagram
    participant Client
    participant HTTPServer
    participant HandleMultiInfer
    participant MultiInferAggregator
    participant TritonServer
    participant InferResponseComplete

    Client->>HTTPServer: POST /v2/multi_infer
    HTTPServer->>HandleMultiInfer: dispatch
    HandleMultiInfer->>MultiInferAggregator: create(n slots)
    loop for each slot i
        HandleMultiInfer->>TritonServer: ScheduleInferAsync(shard_i)
        note over HandleMultiInfer: shard_i released only after ALL slots scheduled
    end
    TritonServer-->>InferResponseComplete: callback(slot_i, flags)
    InferResponseComplete->>MultiInferAggregator: OnShardDone(slot_i, shard_json)
    note over MultiInferAggregator: when done_count == n
    MultiInferAggregator->>HTTPServer: evthr_defer(FinishThunk)
    HTTPServer->>Client: "HTTP 200 {responses:[]}"

Loading

_{Reviews (2): Last reviewed commit: "Separated predict and multi_infer endpoi..." | Re-trigger Greptile}

[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.

[greptile-apps]

Use-after-free when shard scheduling fails mid-loop

In both HandleMultiInfer and HandlePredict, once shard slot i is successfully scheduled via ScheduleInferAsync, the corresponding MultiInferShardRequest is owned by shard_holders (a unique_ptr). If scheduling a later slot i+1 fails, the function returns early — destroying shard_holders[0..i] — while the Triton server still holds in-flight async work for those earlier slots. Their InferResponseComplete callback receives userp as a raw MultiInferShardRequest*, so when the callback fires it dereferences a freed object.

CancelAllSubRequests only marks requests for cancellation and does not guarantee callbacks are suppressed; the callbacks will still fire (with an error), using the dangling pointer.

The fix is to release() each shard/release holder immediately after its ScheduleInferAsync succeeds, rather than deferring all releases to a final loop.

[greptile-apps]

Refresh thread left running when FAIL_IF_ERR calls exit(1)

StartTritonModelsRefreshThread is called before FAIL_IF_ERR(InitializeReadyModelNames(...)). If InitializeReadyModelNames returns an error, FAIL_IF_ERR calls exit(1) without calling JoinTritonModelsRefreshThread. The refresh thread is still alive and iterating through UpdateTritonModelsData, which calls into the global ODBC pool. The TritonDmOdbcPoolAtExit handler registered via std::atexit then nulls out and destroys that pool while the thread is mid-use — a use-after-free.

The simplest fix is to swap the order: call InitializeReadyModelNames first (fail fast before any thread is started), then StartTritonModelsRefreshThread.