diff --git a/HeterogeneousCore/SonicCore/BuildFile.xml b/HeterogeneousCore/SonicCore/BuildFile.xml index b0d5e2a08b98f..5208c91638f37 100644 --- a/HeterogeneousCore/SonicCore/BuildFile.xml +++ b/HeterogeneousCore/SonicCore/BuildFile.xml @@ -2,6 +2,7 @@ + diff --git a/HeterogeneousCore/SonicCore/interface/RetryActionBase.h b/HeterogeneousCore/SonicCore/interface/RetryActionBase.h new file mode 100644 index 0000000000000..ead6efd785dd2 --- /dev/null +++ b/HeterogeneousCore/SonicCore/interface/RetryActionBase.h @@ -0,0 +1,36 @@ +#ifndef HeterogeneousCore_SonicCore_RetryActionBase +#define HeterogeneousCore_SonicCore_RetryActionBase + +#include "FWCore/PluginManager/interface/PluginFactory.h" +#include "FWCore/ParameterSet/interface/ParameterSet.h" +#include "HeterogeneousCore/SonicCore/interface/SonicClientBase.h" +#include +#include + +// Base class for retry actions +class RetryActionBase { +public: + RetryActionBase(const edm::ParameterSet& conf, SonicClientBase* client); + virtual ~RetryActionBase() = default; + + bool shouldRetry() const { return shouldRetry_; } // Getter for shouldRetry_ + + virtual void retry() = 0; // Pure virtual function for execution logic + virtual void start() = 0; // Pure virtual function for execution logic for initialization + +protected: + void eval(); // interface for calling evaluate in client + void finish(bool success); // interface for calling finish directly in client + +protected: + SonicClientBase* client_; + bool shouldRetry_; // Flag to track if further retries should happen +}; + +// Define the factory for creating retry actions +using RetryActionFactory = + edmplugin::PluginFactory; + +#endif + +#define DEFINE_RETRY_ACTION(type) DEFINE_EDM_PLUGIN(RetryActionFactory, type, #type); diff --git a/HeterogeneousCore/SonicCore/interface/SonicClientBase.h b/HeterogeneousCore/SonicCore/interface/SonicClientBase.h index 47caaae8b2052..45a089701ed12 100644 --- a/HeterogeneousCore/SonicCore/interface/SonicClientBase.h +++ b/HeterogeneousCore/SonicCore/interface/SonicClientBase.h @@ -9,12 +9,15 @@ #include "HeterogeneousCore/SonicCore/interface/SonicDispatcherPseudoAsync.h" #include +#include #include #include #include enum class SonicMode { Sync = 1, Async = 2, PseudoAsync = 3 }; +class RetryActionBase; + class SonicClientBase { public: //constructor @@ -54,14 +57,23 @@ class SonicClientBase { SonicMode mode_; bool verbose_; std::unique_ptr dispatcher_; - unsigned allowedTries_, tries_; + unsigned totalTries_; std::optional holder_; + // Use a unique_ptr with a custom deleter to avoid incomplete type issues + struct RetryDeleter { + void operator()(RetryActionBase* ptr) const; + }; + + using RetryActionPtr = std::unique_ptr; + std::vector retryActions_; + //for logging/debugging std::string debugName_, clientName_, fullDebugName_; friend class SonicDispatcher; friend class SonicDispatcherPseudoAsync; + friend class RetryActionBase; }; #endif diff --git a/HeterogeneousCore/SonicCore/plugins/BuildFile.xml b/HeterogeneousCore/SonicCore/plugins/BuildFile.xml new file mode 100644 index 0000000000000..eaff0919e46bc --- /dev/null +++ b/HeterogeneousCore/SonicCore/plugins/BuildFile.xml @@ -0,0 +1,6 @@ + + + + + + diff --git a/HeterogeneousCore/SonicCore/plugins/RetrySameServerAction.cc b/HeterogeneousCore/SonicCore/plugins/RetrySameServerAction.cc new file mode 100644 index 0000000000000..09663ec8e4813 --- /dev/null +++ b/HeterogeneousCore/SonicCore/plugins/RetrySameServerAction.cc @@ -0,0 +1,28 @@ +#include "HeterogeneousCore/SonicCore/interface/RetryActionBase.h" +#include "HeterogeneousCore/SonicCore/interface/SonicClientBase.h" + +class RetrySameServerAction : public RetryActionBase { +public: + RetrySameServerAction(const edm::ParameterSet& pset, SonicClientBase* client) + : RetryActionBase(pset, client), allowedTries_(pset.getUntrackedParameter("allowedTries", 0)) {} + + void start() override { tries_ = 0; }; + +protected: + void retry() override; + +private: + unsigned allowedTries_, tries_; +}; + +void RetrySameServerAction::retry() { + ++tries_; + //if max retries has not been exceeded, call evaluate again + if (tries_ >= allowedTries_) { + shouldRetry_ = false; // Flip flag when max retries are reached + edm::LogInfo("RetrySameServerAction") << "Max retry attempts reached. No further retries."; + } + eval(); +} + +DEFINE_RETRY_ACTION(RetrySameServerAction) diff --git a/HeterogeneousCore/SonicCore/src/RetryActionBase.cc b/HeterogeneousCore/SonicCore/src/RetryActionBase.cc new file mode 100644 index 0000000000000..52b1fefb92a8a --- /dev/null +++ b/HeterogeneousCore/SonicCore/src/RetryActionBase.cc @@ -0,0 +1,15 @@ +#include "HeterogeneousCore/SonicCore/interface/RetryActionBase.h" + +// Constructor implementation +RetryActionBase::RetryActionBase(const edm::ParameterSet& conf, SonicClientBase* client) + : client_(client), shouldRetry_(true) { + if (client_ == nullptr) { + throw cms::Exception("RetryActionBase") << "client pointer cannot be null"; + } +} + +void RetryActionBase::eval() { client_->evaluate(); } + +void RetryActionBase::finish(bool success) { client_->finish(success); } + +EDM_REGISTER_PLUGINFACTORY(RetryActionFactory, "RetryActionFactory"); diff --git a/HeterogeneousCore/SonicCore/src/SonicClientBase.cc b/HeterogeneousCore/SonicCore/src/SonicClientBase.cc index 745c51f17aaf3..6ed10089e1bd3 100644 --- a/HeterogeneousCore/SonicCore/src/SonicClientBase.cc +++ b/HeterogeneousCore/SonicCore/src/SonicClientBase.cc @@ -1,18 +1,34 @@ #include "HeterogeneousCore/SonicCore/interface/SonicClientBase.h" +#include "HeterogeneousCore/SonicCore/interface/RetryActionBase.h" #include "FWCore/Utilities/interface/Exception.h" #include "FWCore/ParameterSet/interface/allowedValues.h" +// Custom deleter implementation +void SonicClientBase::RetryDeleter::operator()(RetryActionBase* ptr) const { delete ptr; } + SonicClientBase::SonicClientBase(const edm::ParameterSet& params, const std::string& debugName, const std::string& clientName) - : allowedTries_(params.getUntrackedParameter("allowedTries", 0)), - debugName_(debugName), - clientName_(clientName), - fullDebugName_(debugName_) { + : debugName_(debugName), clientName_(clientName), fullDebugName_(debugName_) { if (!clientName_.empty()) fullDebugName_ += ":" + clientName_; + const auto& retryPSetList = params.getParameter>("Retry"); std::string modeName(params.getParameter("mode")); + + for (const auto& retryPSet : retryPSetList) { + const std::string& actionType = retryPSet.getParameter("retryType"); + + auto retryAction = RetryActionFactory::get()->create(actionType, retryPSet, this); + if (retryAction) { + //Convert to RetryActionPtr Type from raw pointer of retryAction + retryActions_.emplace_back(RetryActionPtr(retryAction.release())); + } else { + throw cms::Exception("Configuration") + << "Unknown Retry type " << actionType << " for SonicClient: " << fullDebugName_; + } + } + if (modeName == "Sync") setMode(SonicMode::Sync); else if (modeName == "Async") @@ -40,24 +56,34 @@ void SonicClientBase::start(edm::WaitingTaskWithArenaHolder holder) { holder_ = std::move(holder); } -void SonicClientBase::start() { tries_ = 0; } +void SonicClientBase::start() { + totalTries_ = 0; + // initialize all actions + for (auto& action : retryActions_) { + action->start(); + } +} void SonicClientBase::finish(bool success, std::exception_ptr eptr) { //retries are only allowed if no exception was raised if (!success and !eptr) { - ++tries_; - //if max retries has not been exceeded, call evaluate again - if (tries_ < allowedTries_) { - evaluate(); - //avoid calling doneWaiting() twice - return; - } - //prepare an exception if exceeded - else { - edm::Exception ex(edm::errors::ExternalFailure); - ex << "SonicCallFailed: call failed after max " << tries_ << " tries"; - eptr = make_exception_ptr(ex); + ++totalTries_; + edm::LogInfo("SonicClientBase") << "finish: failed after total tries of " << totalTries_; + for (const auto& action : retryActions_) { + if (action->shouldRetry()) { + edm::LogInfo("SonicClientBase") << "Calling retry()"; + // retry() must trigger eval() or finish() + action->retry(); + // return because another finish() was already called inside client->evaluate() + return; + } } + //prepare an exception if no more retry actions left + edm::LogInfo("SonicClientBase") << "SonicCallFailed: call failed, no retry actions available after " << totalTries_ + << " tries."; + edm::Exception ex(edm::errors::ExternalFailure); + ex << "SonicCallFailed: call failed, no retry actions available after " << totalTries_ << " tries."; + eptr = make_exception_ptr(ex); } if (holder_) { holder_->doneWaiting(eptr); @@ -74,7 +100,20 @@ void SonicClientBase::fillBasePSetDescription(edm::ParameterSetDescription& desc //restrict allowed values desc.ifValue(edm::ParameterDescription("mode", "PseudoAsync", true), edm::allowedValues("Sync", "Async", "PseudoAsync")); - if (allowRetry) - desc.addUntracked("allowedTries", 0); + if (allowRetry) { + // Defines the structure of each entry in the VPSet + edm::ParameterSetDescription retryDesc; + retryDesc.add("retryType", "RetrySameServerAction"); + retryDesc.addUntracked("allowedTries", 0); + + // Define a default retry action + edm::ParameterSet defaultRetry; + defaultRetry.addParameter("retryType", "RetrySameServerAction"); + defaultRetry.addUntrackedParameter("allowedTries", 0); + + // Add the VPSet with the default retry action + desc.addVPSet("Retry", retryDesc, {defaultRetry}); + } + desc.add("sonicClientBase", desc); desc.addUntracked("verbose", false); } diff --git a/HeterogeneousCore/SonicCore/test/DummyClient.h b/HeterogeneousCore/SonicCore/test/DummyClient.h index ccef888ad9f7d..6504843926c0a 100644 --- a/HeterogeneousCore/SonicCore/test/DummyClient.h +++ b/HeterogeneousCore/SonicCore/test/DummyClient.h @@ -36,7 +36,7 @@ class DummyClient : public SonicClient { this->output_ = this->input_ * factor_; //simulate a failure - if (this->tries_ < fails_) + if (this->totalTries_ < fails_) this->finish(false); else this->finish(true); diff --git a/HeterogeneousCore/SonicCore/test/sonicTestAna_cfg.py b/HeterogeneousCore/SonicCore/test/sonicTestAna_cfg.py index 11c23c6cdfcc9..35cf42fa2b5ae 100644 --- a/HeterogeneousCore/SonicCore/test/sonicTestAna_cfg.py +++ b/HeterogeneousCore/SonicCore/test/sonicTestAna_cfg.py @@ -16,16 +16,27 @@ mode = cms.string("Sync"), factor = cms.int32(-1), wait = cms.int32(10), - allowedTries = cms.untracked.uint32(0), fails = cms.uint32(0), + Retry = cms.VPSet( + cms.PSet( + retryType = cms.string('RetrySameServerAction'), + allowedTries = cms.untracked.uint32(0), + ) + ) ), ) process.dummySyncAnaRetry = process.dummySyncAna.clone( Client = dict( wait = 2, - allowedTries = 2, fails = 1, + Retry = cms.VPSet( + cms.PSet( + retryType = cms.string('RetrySameServerAction'), + allowedTries = cms.untracked.uint32(2), + ) + ) + ) ) diff --git a/HeterogeneousCore/SonicCore/test/sonicTest_cfg.py b/HeterogeneousCore/SonicCore/test/sonicTest_cfg.py index 614297d86e3bb..bcbe820030440 100644 --- a/HeterogeneousCore/SonicCore/test/sonicTest_cfg.py +++ b/HeterogeneousCore/SonicCore/test/sonicTest_cfg.py @@ -17,15 +17,19 @@ process.options.numberOfThreads = 2 process.options.numberOfStreams = 0 - process.dummySync = _moduleClass(_moduleName, input = cms.int32(1), Client = cms.PSet( mode = cms.string("Sync"), factor = cms.int32(-1), wait = cms.int32(10), - allowedTries = cms.untracked.uint32(0), fails = cms.uint32(0), + Retry = cms.VPSet( + cms.PSet( + retryType = cms.string('RetrySameServerAction'), + allowedTries = cms.untracked.uint32(0) + ) + ) ), ) @@ -35,8 +39,14 @@ mode = cms.string("PseudoAsync"), factor = cms.int32(2), wait = cms.int32(10), - allowedTries = cms.untracked.uint32(0), fails = cms.uint32(0), + Retry = cms.VPSet( + cms.PSet( + retryType = cms.string('RetrySameServerAction'), + allowedTries = cms.untracked.uint32(0) + ) + ) + ), ) @@ -46,32 +56,53 @@ mode = cms.string("Async"), factor = cms.int32(5), wait = cms.int32(10), - allowedTries = cms.untracked.uint32(0), fails = cms.uint32(0), + Retry = cms.VPSet( + cms.PSet( + retryType = cms.string('RetrySameServerAction'), + allowedTries = cms.untracked.uint32(0) + ) + ) ), ) process.dummySyncRetry = process.dummySync.clone( Client = dict( wait = 2, - allowedTries = 2, fails = 1, + Retry = cms.VPSet( + cms.PSet( + retryType = cms.string('RetrySameServerAction'), + allowedTries = cms.untracked.uint32(2) + ) + ) + ) ) process.dummyPseudoAsyncRetry = process.dummyPseudoAsync.clone( Client = dict( wait = 2, - allowedTries = 2, fails = 1, + Retry = cms.VPSet( + cms.PSet( + retryType = cms.string('RetrySameServerAction'), + allowedTries = cms.untracked.uint32(2) + ) + ) ) ) process.dummyAsyncRetry = process.dummyAsync.clone( Client = dict( wait = 2, - allowedTries = 2, fails = 1, + Retry = cms.VPSet( + cms.PSet( + allowedTries = cms.untracked.uint32(2), + retryType = cms.string('RetrySameServerAction') + ) + ) ) ) diff --git a/HeterogeneousCore/SonicTriton/BuildFile.xml b/HeterogeneousCore/SonicTriton/BuildFile.xml index b93d51e711e87..4af38d69d89e9 100644 --- a/HeterogeneousCore/SonicTriton/BuildFile.xml +++ b/HeterogeneousCore/SonicTriton/BuildFile.xml @@ -10,6 +10,7 @@ + - + diff --git a/HeterogeneousCore/SonicTriton/interface/RetryActionDiffServer.h b/HeterogeneousCore/SonicTriton/interface/RetryActionDiffServer.h new file mode 100644 index 0000000000000..4593201a473b3 --- /dev/null +++ b/HeterogeneousCore/SonicTriton/interface/RetryActionDiffServer.h @@ -0,0 +1,32 @@ +#ifndef HeterogeneousCore_SonicTriton_RetryActionDiffServer_h +#define HeterogeneousCore_SonicTriton_RetryActionDiffServer_h + +#include "HeterogeneousCore/SonicCore/interface/RetryActionBase.h" + +/** + * @class RetryActionDiffServer + * @brief A concrete implementation of RetryActionBase that attempts to retry an inference + * request on a different Triton server. + * + * This class provides a fallback mechanism. If an initial inference request fails + * (e.g., due to server unavailability or a model-specific error), this action will be + * triggered. It queries the central TritonService to select an alternative server (e.g., + * the fallback server when available) and instructs the TritonClient to reconnect to + * that server for the retry attempt. This action is designed for one-time use per + * inference call; after the retry attempt, it disables itself until the next `start()` + * call. + */ + +class RetryActionDiffServer : public RetryActionBase { +public: + RetryActionDiffServer(const edm::ParameterSet& conf, SonicClientBase* client); + ~RetryActionDiffServer() override = default; + + void retry() override; + void start() override; + +private: + unsigned tries_; +}; + +#endif diff --git a/HeterogeneousCore/SonicTriton/interface/RetryFallbackServerAction.h b/HeterogeneousCore/SonicTriton/interface/RetryFallbackServerAction.h new file mode 100644 index 0000000000000..8b72afcda08c8 --- /dev/null +++ b/HeterogeneousCore/SonicTriton/interface/RetryFallbackServerAction.h @@ -0,0 +1,18 @@ +#ifndef HeterogeneousCore_SonicTriton_RetryFallbackServerAction_h +#define HeterogeneousCore_SonicTriton_RetryFallbackServerAction_h + +#include "HeterogeneousCore/SonicCore/interface/RetryActionBase.h" + +class RetryFallbackServerAction : public RetryActionBase { +public: + RetryFallbackServerAction(const edm::ParameterSet& conf, SonicClientBase* client); + ~RetryFallbackServerAction() override = default; + + void retry() override; + void start() override; + +private: + unsigned tries_; +}; + +#endif diff --git a/HeterogeneousCore/SonicTriton/interface/TritonClient.h b/HeterogeneousCore/SonicTriton/interface/TritonClient.h index 44118d43d09f1..74d6649f6c5e8 100644 --- a/HeterogeneousCore/SonicTriton/interface/TritonClient.h +++ b/HeterogeneousCore/SonicTriton/interface/TritonClient.h @@ -1,104 +1,118 @@ -#ifndef HeterogeneousCore_SonicTriton_TritonClient -#define HeterogeneousCore_SonicTriton_TritonClient - -#include "FWCore/ParameterSet/interface/ParameterSet.h" -#include "FWCore/ParameterSet/interface/ParameterSetDescription.h" -#include "FWCore/ServiceRegistry/interface/ServiceToken.h" -#include "HeterogeneousCore/SonicCore/interface/SonicClient.h" -#include "HeterogeneousCore/SonicTriton/interface/TritonData.h" -#include "HeterogeneousCore/SonicTriton/interface/TritonService.h" - -#include -#include -#include -#include -#include - -#include "grpc_client.h" -#include "grpc_service.pb.h" - -enum class TritonBatchMode { Rectangular = 1, Ragged = 2 }; - -class TritonClient : public SonicClient { -public: - struct ServerSideStats { - uint64_t inference_count_; - uint64_t execution_count_; - uint64_t success_count_; - uint64_t cumm_time_ns_; - uint64_t queue_time_ns_; - uint64_t compute_input_time_ns_; - uint64_t compute_infer_time_ns_; - uint64_t compute_output_time_ns_; - }; - - //constructor - TritonClient(const edm::ParameterSet& params, const std::string& debugName); - - //destructor - ~TritonClient() override; - - //accessors - unsigned batchSize() const; - TritonBatchMode batchMode() const { return batchMode_; } - bool verbose() const { return verbose_; } - bool useSharedMemory() const { return useSharedMemory_; } - void setUseSharedMemory(bool useShm) { useSharedMemory_ = useShm; } - bool setBatchSize(unsigned bsize); - void setBatchMode(TritonBatchMode batchMode); - void resetBatchMode(); - void reset() override; - TritonServerType serverType() const { return serverType_; } - bool isLocal() const { return isLocal_; } - const TritonService* service() const; - const TritonService* localService() const; - - //for fillDescriptions - static void fillPSetDescription(edm::ParameterSetDescription& iDesc); - -protected: - //helpers - bool noOuterDim() const { return noOuterDim_; } - unsigned outerDim() const { return outerDim_; } - unsigned nEntries() const; - void getResults(const std::vector>& results); - void evaluate() override; - template - bool handle_exception(F&& call); - - void reportServerSideStats(const ServerSideStats& stats) const; - ServerSideStats summarizeServerStats(const inference::ModelStatistics& start_status, - const inference::ModelStatistics& end_status) const; - - inference::ModelStatistics getServerSideStatus() const; - - //members - unsigned maxOuterDim_; - unsigned outerDim_; - bool noOuterDim_; - unsigned nEntries_; - TritonBatchMode batchMode_; - bool manualBatchMode_; - bool verbose_; - bool useSharedMemory_; - TritonServerType serverType_; - bool isLocal_; - grpc_compression_algorithm compressionAlgo_; - triton::client::Headers headers_; - - std::unique_ptr client_; - //stores timeout, model name and version - std::vector options_; - edm::ServiceToken token_; - -private: - friend TritonInputData; - friend TritonOutputData; - - //private accessors only used by data - auto client() { return client_.get(); } - void addEntry(unsigned entry); - void resizeEntries(unsigned entry); -}; - -#endif +#ifndef HeterogeneousCore_SonicTriton_TritonClient +#define HeterogeneousCore_SonicTriton_TritonClient + +#include "FWCore/ParameterSet/interface/ParameterSet.h" +#include "FWCore/ParameterSet/interface/ParameterSetDescription.h" +#include "FWCore/Concurrency/interface/WaitingTaskWithArenaHolder.h" +#include "FWCore/ServiceRegistry/interface/ServiceToken.h" +#include "HeterogeneousCore/SonicCore/interface/SonicClient.h" +#include "HeterogeneousCore/SonicTriton/interface/TritonData.h" +#include "HeterogeneousCore/SonicTriton/interface/TritonService.h" + +#include +#include +#include +#include +#include +#include + +#include "grpc_client.h" +#include "grpc_service.pb.h" + +enum class TritonBatchMode { Rectangular = 1, Ragged = 2 }; + +class TritonClient : public SonicClient { +public: + struct ServerSideStats { + uint64_t inference_count_; + uint64_t execution_count_; + uint64_t success_count_; + uint64_t cumm_time_ns_; + uint64_t queue_time_ns_; + uint64_t compute_input_time_ns_; + uint64_t compute_infer_time_ns_; + uint64_t compute_output_time_ns_; + }; + + //constructor + TritonClient(const edm::ParameterSet& params, const std::string& debugName); + + //destructor + ~TritonClient() override; + + //accessors + unsigned batchSize() const; + TritonBatchMode batchMode() const { return batchMode_; } + bool verbose() const { return verbose_; } + bool useSharedMemory() const { return useSharedMemory_; } + void setUseSharedMemory(bool useShm) { useSharedMemory_ = useShm; } + bool setBatchSize(unsigned bsize); + void setBatchMode(TritonBatchMode batchMode); + void resetBatchMode(); + void reset() override; + TritonServerType serverType() const { return serverType_; } + bool isLocal() const { return isLocal_; } + const TritonService* service() const; + const TritonService* localService() const; + std::string modelName() const { return options_[0].model_name_; } + std::string serverName() const { return serverName_; } + virtual void connectToServer(const std::string& url); + virtual void updateServer(const std::string& serverName); + virtual void switchToFallback(); + + //for fillDescriptions + static void fillPSetDescription(edm::ParameterSetDescription& iDesc); + +protected: + // Protected default constructor for unit testing (no framework services) + TritonClient(); + + //helpers + bool noOuterDim() const { return noOuterDim_; } + unsigned outerDim() const { return outerDim_; } + unsigned nEntries() const; + void getResults(const std::vector>& results); + void evaluate() override; + template + bool handle_exception(F&& call); + template + bool handle_exception_holder(edm::WaitingTaskWithArenaHolder& fh, F&& call); + + void reportServerSideStats(const ServerSideStats& stats) const; + ServerSideStats summarizeServerStats(const inference::ModelStatistics& start_status, + const inference::ModelStatistics& end_status) const; + + inference::ModelStatistics getServerSideStatus() const; + + //members + unsigned maxOuterDim_; + unsigned outerDim_; + bool noOuterDim_; + unsigned nEntries_; + TritonBatchMode batchMode_; + bool manualBatchMode_; + bool verbose_; + bool useSharedMemory_; + TritonServerType serverType_; + bool isLocal_; + std::string serverName_; + grpc_compression_algorithm compressionAlgo_; + triton::client::Headers headers_; + + std::unique_ptr client_; + //stores timeout, model name and version + std::vector options_; + edm::ServiceToken token_; + std::atomic inferSuccess_{true}; + +private: + friend TritonInputData; + friend TritonOutputData; + + //private accessors only used by data + auto client() { return client_.get(); } + void addEntry(unsigned entry); + void resizeEntries(unsigned entry); +}; + +#endif diff --git a/HeterogeneousCore/SonicTriton/interface/TritonService.h b/HeterogeneousCore/SonicTriton/interface/TritonService.h index 8ac7a915f8d6d..9be006dfab615 100644 --- a/HeterogeneousCore/SonicTriton/interface/TritonService.h +++ b/HeterogeneousCore/SonicTriton/interface/TritonService.h @@ -3,6 +3,7 @@ #include "FWCore/ParameterSet/interface/ParameterSet.h" #include "FWCore/Utilities/interface/GlobalIdentifier.h" +#include "oneapi/tbb/concurrent_hash_map.h" #include #include @@ -11,6 +12,8 @@ #include #include #include +#include +#include #include "grpc_client.h" @@ -90,18 +93,28 @@ class TritonService { static const std::string fallbackAddress; static const std::string siteconfName; }; + //Dynamic quantities of servers + struct ServerHealth { + bool live{false}; + bool ready{false}; + + uint64_t inferenceCount{0}; + uint64_t failureCount{0}; + double avgQueueTimeMs{0.0}; + double avgInferTimeMs{0.0}; + }; struct Model { Model(const std::string& path_ = "") : path(path_) {} - //members std::string path; std::unordered_set servers; std::unordered_set modules; + int refCount{0}; // for dynamic loading on fallback server + bool isLoaded() const { return refCount > 0; } }; struct Module { //currently assumes that a module can only have one associated model Module(const std::string& model_) : model(model_) {} - //members std::string model; }; @@ -111,12 +124,40 @@ class TritonService { //accessors void addModel(const std::string& modelName, const std::string& path); - Server serverInfo(const std::string& model, const std::string& preferred = "") const; + + const std::string* resolveServerName(const std::string& model, const std::string& preferred = "") const; + const std::pair& resolveServer(const std::string& model, + const std::string& preferred = "") const; + std::vector unassignedModels() const; + + // update health stats of all servers + void updateServerHealth(const std::string& modelName = "") const; + + // return the best server for retry, ignore the current server + std::optional getBestServer(const std::string& modelName, const std::string& IgnoreServer = "") const; + + // helper functions to get server statistics? + // - getServerSideStatus() + // - updateServerStatus() + // - loop over servers_ get statistics + // - getBestServer(model) + // - call updateServerStatus() + // - loop over servers_ get their statistics, compute metric, return server name + const std::string& pid() const { return pid_; } void notifyCallStatus(bool status) const; static void fillDescriptions(edm::ConfigurationDescriptions& descriptions); + // Dynamic model loading/unloading - only supported for the fallback server + // The fallback server must be started with explicit model control mode + // (--model-control-mode explicit) for these functions to work + bool loadModel(const std::string& modelName); + bool unloadModel(const std::string& modelName); + // Start the fallback server if enabled and not already running (idempotent) + void startFallbackServer(); + bool fallbackStarted() const { return startedFallback_; } + private: void preallocate(edm::service::SystemBounds const&); void preModuleConstruction(edm::ModuleDescription const&); @@ -128,6 +169,9 @@ class TritonService { //helper template void printFallbackServerLog() const; + // Internal helpers that operate on Model directly (caller holds lock) + bool loadModel(const std::string& modelName, Model& model); + bool unloadModel(const std::string& modelName, Model& model); bool verbose_; FallbackOpts fallbackOpts_; @@ -136,12 +180,18 @@ class TritonService { bool startedFallback_; mutable std::atomic callFails_; std::string pid_; - std::unordered_map unservedModels_; //this represents a many:many:many map std::unordered_map servers_; + //server health needs concurrent-safe edits + tbb::concurrent_hash_map serversHealth_; std::unordered_map models_; std::unordered_map modules_; int numberOfThreads_; + + //Dynamic model loading and unloading (fallback server only) + std::mutex modelLoadMutex_; + // Model names currently loaded on the fallback server + std::unordered_set fallbackLoadedModels_; }; #endif diff --git a/HeterogeneousCore/SonicTriton/python/customize.py b/HeterogeneousCore/SonicTriton/python/customize.py index 9873a1a498dd7..e633d169e919e 100644 --- a/HeterogeneousCore/SonicTriton/python/customize.py +++ b/HeterogeneousCore/SonicTriton/python/customize.py @@ -13,12 +13,12 @@ def getParser(): from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter) parser.add_argument("--maxEvents", default=-1, type=int, help="Number of events to process (-1 for all)") - parser.add_argument("--serverName", default="default", type=str, help="name for server (used internally)") - parser.add_argument("--address", default="", type=str, help="server address") - parser.add_argument("--port", default=8001, type=int, help="server port") + parser.add_argument("--address", nargs=3, action="append", metavar=("NAME", "HOST", "PORT"), + dest="addresses", default=[], + help="Triton server entry: name host port (repeatable, e.g. --address server1 0.0.0.0 8011)") parser.add_argument("--timeout", default=30, type=int, help="timeout for requests") parser.add_argument("--timeoutUnit", default="seconds", type=str, help="unit for timeout") - parser.add_argument("--params", default="", type=str, help="json file containing server address/port") + parser.add_argument("--params", default="", type=str, help="json file containing server address/port(single-server)") parser.add_argument("--threads", default=1, type=int, help="number of threads") parser.add_argument("--streams", default=0, type=int, help="number of streams") parser.add_argument("--verbose", default=False, action="store_true", help="enable all verbose output") @@ -36,12 +36,27 @@ def getParser(): parser.add_argument("--imageName", default="", type=str, help="container image name for fallback server") parser.add_argument("--sandboxDir", default="", type=str, help="apptainer sandbox directory") parser.add_argument("--tempDir", default="", type=str, help="temp directory for fallback server") + parser.add_argument("--retryAction", default="same", type=str, choices=["same","diff"], help="retry policy: same server or different server") return parser def getOptions(parser, verbose=False): options = parser.parse_args() + # Legacy --params support: loads a single server and appends it to options.addresses + if len(options.params) > 0: + with open(options.params, 'r') as pfile: + pdict = json.load(pfile) + name = pdict.get("name", "default") + host = pdict["address"] + port = str(int(pdict["port"])) + options.addresses.append([name, host, port]) + if verbose: + print("server (from params) = {}:{} [{}]".format(host, port, name)) + + if verbose: + for name, host, port in options.addresses: + print("server = {}:{} [{}]".format(host, port, name)) if len(options.params)>0: with open(options.params,'r') as pfile: pdict = json.load(pfile) @@ -76,13 +91,13 @@ def applyOptions(process, options, applyToModules=False): process.TritonService.fallback.device = options.device if len(options.fallbackName)>0: process.TritonService.fallback.instanceBaseName = options.fallbackName - if len(options.address)>0: + for name, host, port in options.addresses: process.TritonService.servers.append( dict( - name = options.serverName, - address = options.address, - port = options.port, - useSsl = options.ssl, + name = name, + address = host, + port = int(port), + useSsl = options.ssl, ) ) @@ -92,12 +107,19 @@ def applyOptions(process, options, applyToModules=False): return process def getClientOptions(options): + action = cms.PSet( + retryType = cms.string('RetrySameServerAction'), + allowedTries = cms.untracked.uint32(options.tries)) + if options.retryAction != 'same': + action.retryType = cms.string('RetryActionDiffServer') + + fallback = cms.PSet(retryType = cms.string('RetryFallbackServerAction')) return dict( compression = cms.untracked.string(options.compression), useSharedMemory = cms.untracked.bool(not options.noShm), timeout = cms.untracked.uint32(options.timeout), timeoutUnit = cms.untracked.string(options.timeoutUnit), - allowedTries = cms.untracked.uint32(options.tries), + Retry = cms.VPSet(action,fallback) ) def applyClientOptions(client, options): diff --git a/HeterogeneousCore/SonicTriton/scripts/cmsTriton b/HeterogeneousCore/SonicTriton/scripts/cmsTriton index 01932254f9f1d..c3d7a9fc01d87 100755 --- a/HeterogeneousCore/SonicTriton/scripts/cmsTriton +++ b/HeterogeneousCore/SonicTriton/scripts/cmsTriton @@ -31,6 +31,7 @@ fi DEVICE=auto THREADCONTROL="" GLOBALCONNECT= +MODELCONTROL="explicit" get_sandbox(){ if [ -z "$SANDBOX" ]; then @@ -55,6 +56,7 @@ usage() { $ECHO "-G \t accept connections globally (default: only accept connections from localhost)" $ECHO "-i [name] \t server image name (default: ${IMAGE})" $ECHO "-I [num] \t number of model instances (default: ${INSTANCES} -> means no local editing of config files)" + $ECHO "-L \t disable explicit model control mode (enabled by default for dynamic model loading)" $ECHO "-M [dir] \t model repository (can be given more than once)" $ECHO "-m [dir] \t specific model directory (can be given more than once)" $ECHO "-n [name] \t name of container instance, also used for default hidden temporary dir (default: ${SERVER})" @@ -80,7 +82,7 @@ if [ -e /run/shm ]; then SHM=/run/shm fi -while getopts "cC:Dd:fg:i:I:M:m:n:P:p:r:s:t:vw:h" opt; do +while getopts "cC:Dd:fg:i:I:LM:m:n:P:p:r:s:t:vw:h" opt; do case "$opt" in c) CLEANUP="" ;; @@ -100,6 +102,8 @@ while getopts "cC:Dd:fg:i:I:M:m:n:P:p:r:s:t:vw:h" opt; do ;; I) INSTANCES="$OPTARG" ;; + L) MODELCONTROL="none" + ;; M) REPOS+=("$OPTARG") ;; m) MODELS+=("$OPTARG") @@ -250,6 +254,9 @@ start_docker(){ # mount all model repositories MOUNTARGS="" REPOARGS="" + if [ -n "$MODELCONTROL" ]; then + REPOARGS="--model-control-mode=${MODELCONTROL}" + fi for REPO in ${REPOS[@]}; do MOUNTARGS="$MOUNTARGS -v$REPO:$REPO" REPOARGS="$REPOARGS --model-repository=${REPO}" @@ -276,6 +283,9 @@ start_podman(){ # mount all model repositories MOUNTARGS="" REPOARGS="" + if [ -n "$MODELCONTROL" ]; then + REPOARGS="--model-control-mode=${MODELCONTROL}" + fi for REPO in ${REPOS[@]}; do MOUNTARGS="$MOUNTARGS --volume $REPO:$REPO" REPOARGS="$REPOARGS --model-repository=${REPO}" @@ -302,6 +312,9 @@ start_apptainer(){ # mount all model repositories MOUNTARGS="" REPOARGS="" + if [ -n "$MODELCONTROL" ]; then + REPOARGS="--model-control-mode=${MODELCONTROL}" + fi for REPO in ${REPOS[@]}; do MOUNTARGS="$MOUNTARGS -B $REPO" REPOARGS="$REPOARGS --model-repository=${REPO}" diff --git a/HeterogeneousCore/SonicTriton/src/RetryActionDiffServer.cc b/HeterogeneousCore/SonicTriton/src/RetryActionDiffServer.cc new file mode 100644 index 0000000000000..8a096f1bfd672 --- /dev/null +++ b/HeterogeneousCore/SonicTriton/src/RetryActionDiffServer.cc @@ -0,0 +1,50 @@ +#include "HeterogeneousCore/SonicTriton/interface/RetryActionDiffServer.h" +#include "HeterogeneousCore/SonicTriton/interface/TritonClient.h" +#include "HeterogeneousCore/SonicTriton/interface/TritonService.h" +#include "FWCore/MessageLogger/interface/MessageLogger.h" +#include "FWCore/ServiceRegistry/interface/Service.h" + +RetryActionDiffServer::RetryActionDiffServer(const edm::ParameterSet& conf, SonicClientBase* client) + : RetryActionBase(conf, client) {} + +void RetryActionDiffServer::start() { + this->shouldRetry_ = true; + tries_ = 0; +} + +void RetryActionDiffServer::retry() { + ++tries_; + if (tries_ >= 1) { + shouldRetry_ = false; // Flip flag when max retries are reached. Allow 1 try for now. + edm::LogInfo("RetryDiffServerAction") << "Max retry attempts reached. No further retries."; + } + try { + auto* tritonClient = static_cast(client_); + edm::LogInfo("RetryActionDiffServer") << "Asking for a different server from TritonService"; + auto ts = tritonClient->service(); + + // First, try to find another remote server + auto bestServerName = ts->getBestServer(tritonClient->modelName(), tritonClient->serverName()); + + if (bestServerName) { + edm::LogInfo("RetryActionDiffServer") << "Got best server from service "; + tritonClient->updateServer(*bestServerName); + edm::LogInfo("RetryActionDiffServer") << "eval() with new server"; + eval(); + return; + } else { + edm::LogWarning("RetryActionDiffServer") + << "No alternative server found for model " << tritonClient->modelName() << ". Now call client->finish()"; + finish(false); + return; + } + } catch (TritonException& e) { + e.convertToWarning(); + } catch (std::exception& e) { + edm::LogError("RetryActionDiffServer") << "Failed to retry with alternative server: " << e.what(); + } catch (...) { + edm::LogError("RetryActionDiffServer: UnknownFailure") << "An unknown exception was thrown"; + } +} + +DEFINE_RETRY_ACTION(RetryActionDiffServer); diff --git a/HeterogeneousCore/SonicTriton/src/RetryFallbackServerAction.cc b/HeterogeneousCore/SonicTriton/src/RetryFallbackServerAction.cc new file mode 100644 index 0000000000000..823d6ce4eca2b --- /dev/null +++ b/HeterogeneousCore/SonicTriton/src/RetryFallbackServerAction.cc @@ -0,0 +1,54 @@ +// RetryFallbackServerAction: last-resort retry action for TritonClient. +// +// When all other retry actions have been exhausted, this action loads the +// client's model onto the fallback (local) Triton server and re-runs +// inference there. It fires at most once per inference call. +// +// Usage add to the Retry VPSet *after* all other retry actions: +// cms.PSet(retryType = cms.string("RetryFallbackServerAction")) +// +// Requirements: +// - TritonService fallback must be enabled in the job configuration. +// - The model must have a modelConfigPath / repository path known to +// TritonService so it can be loaded dynamically. + +#include "HeterogeneousCore/SonicTriton/interface/RetryFallbackServerAction.h" +#include "HeterogeneousCore/SonicTriton/interface/TritonClient.h" +#include "HeterogeneousCore/SonicTriton/interface/TritonService.h" +#include "FWCore/MessageLogger/interface/MessageLogger.h" +#include "FWCore/ParameterSet/interface/ParameterSet.h" +#include "FWCore/Utilities/interface/Exception.h" + +RetryFallbackServerAction::RetryFallbackServerAction(const edm::ParameterSet& conf, SonicClientBase* client) + : RetryActionBase(conf, client) {} + +void RetryFallbackServerAction::start() { + this->shouldRetry_ = true; + tries_ = 0; +} + +void RetryFallbackServerAction::retry() { + // Allow only one fallback attempt per inference call. + shouldRetry_ = false; + + auto* tc = dynamic_cast(client_); + if (!tc) { + // Should never happen in a correctly configured job. + edm::LogWarning("RetryFallbackServerAction") + << "client_ is not a TritonClient — cannot redirect to fallback server"; + finish(false); + return; + } + + CMS_SA_ALLOW try { + // Start the fallback server (idempotent), load the model, and point + // the client's gRPC connection at the fallback URL. + tc->switchToFallback(); + // Re-run the inference on the fallback server. + eval(); + } catch (...) { + // Non-retryable: propagate the exception so the job fails cleanly. + finish(false); + } +} +DEFINE_RETRY_ACTION(RetryFallbackServerAction); diff --git a/HeterogeneousCore/SonicTriton/src/TritonClient.cc b/HeterogeneousCore/SonicTriton/src/TritonClient.cc index 05fb59ce9b0c0..bc223440a68b4 100644 --- a/HeterogeneousCore/SonicTriton/src/TritonClient.cc +++ b/HeterogeneousCore/SonicTriton/src/TritonClient.cc @@ -1,601 +1,705 @@ -#include "FWCore/MessageLogger/interface/MessageLogger.h" -#include "FWCore/ParameterSet/interface/FileInPath.h" -#include "FWCore/ParameterSet/interface/allowedValues.h" -#include "FWCore/ServiceRegistry/interface/Service.h" -#include "FWCore/Utilities/interface/Exception.h" -#include "HeterogeneousCore/SonicTriton/interface/TritonClient.h" -#include "HeterogeneousCore/SonicTriton/interface/TritonException.h" -#include "HeterogeneousCore/SonicTriton/interface/TritonService.h" -#include "HeterogeneousCore/SonicTriton/interface/triton_utils.h" - -#include "grpc_client.h" -#include "grpc_service.pb.h" -#include "model_config.pb.h" - -#include "google/protobuf/text_format.h" -#include "google/protobuf/io/zero_copy_stream_impl.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace tc = triton::client; - -namespace { - grpc_compression_algorithm getCompressionAlgo(const std::string& name) { - if (name.empty() or name.compare("none") == 0) - return grpc_compression_algorithm::GRPC_COMPRESS_NONE; - else if (name.compare("deflate") == 0) - return grpc_compression_algorithm::GRPC_COMPRESS_DEFLATE; - else if (name.compare("gzip") == 0) - return grpc_compression_algorithm::GRPC_COMPRESS_GZIP; - else - throw cms::Exception("GrpcCompression") - << "Unknown compression algorithm requested: " << name << " (choices: none, deflate, gzip)"; - } - - std::vector> convertToShared(const std::vector& tmp) { - std::vector> results; - results.reserve(tmp.size()); - std::transform(tmp.begin(), tmp.end(), std::back_inserter(results), [](tc::InferResult* ptr) { - return std::shared_ptr(ptr); - }); - return results; - } -} // namespace - -//based on https://github.com/triton-inference-server/server/blob/v2.3.0/src/clients/c++/examples/simple_grpc_async_infer_client.cc -//and https://github.com/triton-inference-server/server/blob/v2.3.0/src/clients/c++/perf_client/perf_client.cc - -TritonClient::TritonClient(const edm::ParameterSet& params, const std::string& debugName) - : SonicClient(params, debugName, "TritonClient"), - batchMode_(TritonBatchMode::Rectangular), - manualBatchMode_(false), - verbose_(params.getUntrackedParameter("verbose")), - useSharedMemory_(params.getUntrackedParameter("useSharedMemory")), - compressionAlgo_(getCompressionAlgo(params.getUntrackedParameter("compression"))) { - options_.emplace_back(params.getParameter("modelName")); - //get appropriate server for this model - edm::Service ts; - - // We save the token to be able to notify the service in case of an exception in the evaluate method. - // The evaluate method can be called outside the frameworks TBB threadpool in the case of a retry. In - // this case the context is not setup to access the service registry, we need the service token to - // create the context. - token_ = edm::ServiceRegistry::instance().presentToken(); - - const auto& server = - ts->serverInfo(options_[0].model_name_, params.getUntrackedParameter("preferredServer")); - serverType_ = server.type; - edm::LogInfo("TritonDiscovery") << debugName_ << " assigned server: " << server.url; - //enforce sync mode for fallback CPU server to avoid contention - //todo: could enforce async mode otherwise (unless mode was specified by user?) - if (serverType_ == TritonServerType::LocalCPU) - setMode(SonicMode::Sync); - isLocal_ = serverType_ == TritonServerType::LocalCPU or serverType_ == TritonServerType::LocalGPU; - - //connect to the server - TRITON_THROW_IF_ERROR( - tc::InferenceServerGrpcClient::Create(&client_, server.url, false, server.useSsl, server.sslOptions), - "TritonClient(): unable to create inference context", - localService()); - - //set options - options_[0].model_version_ = params.getParameter("modelVersion"); - options_[0].client_timeout_ = params.getUntrackedParameter("timeout"); - //convert to microseconds - const auto& timeoutUnit = params.getUntrackedParameter("timeoutUnit"); - unsigned conversion = 1; - if (timeoutUnit == "seconds") - conversion = 1e6; - else if (timeoutUnit == "milliseconds") - conversion = 1e3; - else if (timeoutUnit == "microseconds") - conversion = 1; - else - throw cms::Exception("Configuration") << "Unknown timeout unit: " << timeoutUnit; - options_[0].client_timeout_ *= conversion; - - //get fixed parameters from local config - inference::ModelConfig localModelConfig; - { - const std::string localModelConfigPath(params.getParameter("modelConfigPath").fullPath()); - int fileDescriptor = open(localModelConfigPath.c_str(), O_RDONLY); - if (fileDescriptor < 0) - throw TritonException("LocalFailure") - << "TritonClient(): unable to open local model config: " << localModelConfigPath; - google::protobuf::io::FileInputStream localModelConfigInput(fileDescriptor); - localModelConfigInput.SetCloseOnDelete(true); - if (!google::protobuf::TextFormat::Parse(&localModelConfigInput, &localModelConfig)) - throw TritonException("LocalFailure") - << "TritonClient(): unable to parse local model config: " << localModelConfigPath; - } - - //check batch size limitations (after i/o setup) - //triton uses max batch size = 0 to denote a model that does not support native batching (using the outer dimension) - //but for models that do support batching (native or otherwise), a given event may set batch size 0 to indicate no valid input is present - //so set the local max to 1 and keep track of "no outer dim" case - maxOuterDim_ = localModelConfig.max_batch_size(); - noOuterDim_ = maxOuterDim_ == 0; - maxOuterDim_ = std::max(1u, maxOuterDim_); - //propagate batch size - setBatchSize(1); - - //compare model checksums to remote config to enforce versioning - inference::ModelConfigResponse modelConfigResponse; - TRITON_THROW_IF_ERROR(client_->ModelConfig(&modelConfigResponse, options_[0].model_name_, options_[0].model_version_), - "TritonClient(): unable to get model config", - localService()); - inference::ModelConfig remoteModelConfig(modelConfigResponse.config()); - - std::map> checksums; - size_t fileCounter = 0; - for (const auto& modelConfig : {localModelConfig, remoteModelConfig}) { - const auto& agents = modelConfig.model_repository_agents().agents(); - auto agent = std::find_if(agents.begin(), agents.end(), [](auto const& a) { return a.name() == "checksum"; }); - if (agent != agents.end()) { - const auto& params = agent->parameters(); - for (const auto& [key, val] : params) { - // only check the requested version - if (key.compare(0, options_[0].model_version_.size() + 1, options_[0].model_version_ + "/") == 0) - checksums[key][fileCounter] = val; - } - } - ++fileCounter; - } - std::vector incorrect; - for (const auto& [key, val] : checksums) { - if (checksums[key][0] != checksums[key][1]) - incorrect.push_back(key); - } - if (!incorrect.empty()) - throw TritonException("ModelVersioning") << "The following files have incorrect checksums on the remote server: " - << triton_utils::printColl(incorrect, ", "); - - //get model info - inference::ModelMetadataResponse modelMetadata; - TRITON_THROW_IF_ERROR(client_->ModelMetadata(&modelMetadata, options_[0].model_name_, options_[0].model_version_), - "TritonClient(): unable to get model metadata", - localService()); - - //get input and output (which know their sizes) - const auto& nicInputs = modelMetadata.inputs(); - const auto& nicOutputs = modelMetadata.outputs(); - - //report all model errors at once - std::stringstream msg; - std::string msg_str; - - //currently no use case is foreseen for a model with zero inputs or outputs - if (nicInputs.empty()) - msg << "Model on server appears malformed (zero inputs)\n"; - - if (nicOutputs.empty()) - msg << "Model on server appears malformed (zero outputs)\n"; - - //stop if errors - msg_str = msg.str(); - if (!msg_str.empty()) - throw cms::Exception("ModelErrors") << msg_str; - - //setup input map - std::stringstream io_msg; - if (verbose_) - io_msg << "Model inputs: " - << "\n"; - for (const auto& nicInput : nicInputs) { - const auto& iname = nicInput.name(); - auto [curr_itr, success] = input_.emplace(std::piecewise_construct, - std::forward_as_tuple(iname), - std::forward_as_tuple(iname, nicInput, this, ts->pid())); - auto& curr_input = curr_itr->second; - if (verbose_) { - io_msg << " " << iname << " (" << curr_input.dname() << ", " << curr_input.byteSize() - << " b) : " << triton_utils::printColl(curr_input.shape()) << "\n"; - } - } - - //allow selecting only some outputs from server - const auto& v_outputs = params.getUntrackedParameter>("outputs"); - std::unordered_set s_outputs(v_outputs.begin(), v_outputs.end()); - - //setup output map - if (verbose_) - io_msg << "Model outputs: " - << "\n"; - for (const auto& nicOutput : nicOutputs) { - const auto& oname = nicOutput.name(); - if (!s_outputs.empty() and s_outputs.find(oname) == s_outputs.end()) - continue; - auto [curr_itr, success] = output_.emplace(std::piecewise_construct, - std::forward_as_tuple(oname), - std::forward_as_tuple(oname, nicOutput, this, ts->pid())); - auto& curr_output = curr_itr->second; - if (verbose_) { - io_msg << " " << oname << " (" << curr_output.dname() << ", " << curr_output.byteSize() - << " b) : " << triton_utils::printColl(curr_output.shape()) << "\n"; - } - if (!s_outputs.empty()) - s_outputs.erase(oname); - } - - //check if any requested outputs were not available - if (!s_outputs.empty()) - throw cms::Exception("MissingOutput") - << "Some requested outputs were not available on the server: " << triton_utils::printColl(s_outputs); - - //print model info - std::stringstream model_msg; - if (verbose_) { - model_msg << "Model name: " << options_[0].model_name_ << "\n" - << "Model version: " << options_[0].model_version_ << "\n" - << "Model max outer dim: " << (noOuterDim_ ? 0 : maxOuterDim_) << "\n"; - edm::LogInfo(fullDebugName_) << model_msg.str() << io_msg.str(); - } -} - -TritonClient::~TritonClient() { - //by default: members of this class destroyed before members of base class - //in shared memory case, TritonMemResource (member of TritonData) unregisters from client_ in its destructor - //but input/output objects are member of base class, so destroyed after client_ (member of this class) - //therefore, clear the maps here - input_.clear(); - output_.clear(); -} - -void TritonClient::setBatchMode(TritonBatchMode batchMode) { - unsigned oldBatchSize = batchSize(); - batchMode_ = batchMode; - manualBatchMode_ = true; - //this allows calling setBatchSize() and setBatchMode() in either order consistently to change back and forth - //includes handling of change from ragged to rectangular if multiple entries already created - setBatchSize(oldBatchSize); -} - -void TritonClient::resetBatchMode() { - batchMode_ = TritonBatchMode::Rectangular; - manualBatchMode_ = false; -} - -unsigned TritonClient::nEntries() const { return !input_.empty() ? input_.begin()->second.entries_.size() : 0; } - -unsigned TritonClient::batchSize() const { return batchMode_ == TritonBatchMode::Rectangular ? outerDim_ : nEntries(); } - -bool TritonClient::setBatchSize(unsigned bsize) { - if (batchMode_ == TritonBatchMode::Rectangular) { - if (bsize > maxOuterDim_) { - throw TritonException("LocalFailure") - << "Requested batch size " << bsize << " exceeds server-specified max batch size " << maxOuterDim_ << "."; - return false; - } else { - outerDim_ = bsize; - //take min to allow resizing to 0 - resizeEntries(std::min(outerDim_, 1u)); - return true; - } - } else { - resizeEntries(bsize); - outerDim_ = 1; - return true; - } -} - -void TritonClient::resizeEntries(unsigned entry) { - if (entry > nEntries()) - //addEntry(entry) extends the vector to size entry+1 - addEntry(entry - 1); - else if (entry < nEntries()) { - for (auto& element : input_) { - element.second.entries_.resize(entry); - } - for (auto& element : output_) { - element.second.entries_.resize(entry); - } - } -} - -void TritonClient::addEntry(unsigned entry) { - for (auto& element : input_) { - element.second.addEntryImpl(entry); - } - for (auto& element : output_) { - element.second.addEntryImpl(entry); - } - if (entry > 0) { - batchMode_ = TritonBatchMode::Ragged; - outerDim_ = 1; - } -} - -void TritonClient::reset() { - if (!manualBatchMode_) - batchMode_ = TritonBatchMode::Rectangular; - for (auto& element : input_) { - element.second.reset(); - } - for (auto& element : output_) { - element.second.reset(); - } -} - -template -bool TritonClient::handle_exception(F&& call) { - //caught exceptions will be propagated to edm::WaitingTaskWithArenaHolder - CMS_SA_ALLOW try { - call(); - return true; - } - //TritonExceptions are intended/expected to be recoverable, i.e. retries should be allowed - catch (TritonException& e) { - e.convertToWarning(); - finish(false); - return false; - } - //other exceptions are not: execution should stop if they are encountered - catch (...) { - finish(false, std::current_exception()); - return false; - } -} - -const TritonService* TritonClient::service() const { - edm::ServiceRegistry::Operate op(token_); - edm::Service ts; - return &(*ts); -} - -const TritonService* TritonClient::localService() const { return isLocal_ ? service() : nullptr; } - -void TritonClient::getResults(const std::vector>& results) { - for (unsigned i = 0; i < results.size(); ++i) { - const auto& result = results[i]; - for (auto& [oname, output] : output_) { - //set shape here before output becomes const - if (output.variableDims()) { - std::vector tmp_shape; - TRITON_THROW_IF_ERROR(result->Shape(oname, &tmp_shape), - "getResults(): unable to get output shape for " + oname); - if (!noOuterDim_) - tmp_shape.erase(tmp_shape.begin()); - output.setShape(tmp_shape, i); - } - //extend lifetime - output.setResult(result, i); - //compute size after getting all result entries - if (i == results.size() - 1) - output.computeSizes(); - } - } -} - -//default case for sync and pseudo async -void TritonClient::evaluate() { - //undo previous signal from TritonException - if (tries_ > 0) { - // If we are retrying then the evaluate method is called outside the frameworks TBB thread pool. - // So we need to setup the service token for the current thread to access the service registry. - edm::ServiceRegistry::Operate op(token_); - edm::Service ts; - ts->notifyCallStatus(true); - } - - //in case there is nothing to process - if (batchSize() == 0) { - //call getResults on an empty vector - std::vector> empty_results; - getResults(empty_results); - finish(true); - return; - } - - //set up input pointers for triton (generalized for multi-request ragged batching case) - //one vector per request - unsigned nEntriesVal = nEntries(); - std::vector> inputsTriton(nEntriesVal); - for (auto& inputTriton : inputsTriton) { - inputTriton.reserve(input_.size()); - } - for (auto& [iname, input] : input_) { - for (unsigned i = 0; i < nEntriesVal; ++i) { - inputsTriton[i].push_back(input.data(i)); - } - } - - //set up output pointers similarly - std::vector> outputsTriton(nEntriesVal); - for (auto& outputTriton : outputsTriton) { - outputTriton.reserve(output_.size()); - } - for (auto& [oname, output] : output_) { - for (unsigned i = 0; i < nEntriesVal; ++i) { - outputsTriton[i].push_back(output.data(i)); - } - } - - //set up shared memory for output - auto success = handle_exception([&]() { - for (auto& element : output_) { - element.second.prepare(); - } - }); - if (!success) - return; - - // Get the status of the server prior to the request being made. - inference::ModelStatistics start_status; - success = handle_exception([&]() { - if (verbose()) - start_status = getServerSideStatus(); - }); - if (!success) - return; - - if (mode_ == SonicMode::Async) { - //non-blocking call - success = handle_exception([&]() { - TRITON_THROW_IF_ERROR(client_->AsyncInferMulti( - [start_status, this](std::vector resultsTmp) { - //immediately convert to shared_ptr - const auto& results = convertToShared(resultsTmp); - //check results - for (auto ptr : results) { - auto success = handle_exception([&]() { - TRITON_THROW_IF_ERROR( - ptr->RequestStatus(), "evaluate(): unable to get result(s)", localService()); - }); - if (!success) - return; - } - - if (verbose()) { - inference::ModelStatistics end_status; - auto success = handle_exception([&]() { end_status = getServerSideStatus(); }); - if (!success) - return; - - const auto& stats = summarizeServerStats(start_status, end_status); - reportServerSideStats(stats); - } - - //check result - auto success = handle_exception([&]() { getResults(results); }); - if (!success) - return; - - //finish - finish(true); - }, - options_, - inputsTriton, - outputsTriton, - headers_, - compressionAlgo_), - "evaluate(): unable to launch async run", - localService()); - }); - if (!success) - return; - } else { - //blocking call - std::vector resultsTmp; - success = handle_exception([&]() { - TRITON_THROW_IF_ERROR( - client_->InferMulti(&resultsTmp, options_, inputsTriton, outputsTriton, headers_, compressionAlgo_), - "evaluate(): unable to run and/or get result", - localService()); - }); - //immediately convert to shared_ptr - const auto& results = convertToShared(resultsTmp); - if (!success) - return; - - if (verbose()) { - inference::ModelStatistics end_status; - success = handle_exception([&]() { end_status = getServerSideStatus(); }); - if (!success) - return; - - const auto& stats = summarizeServerStats(start_status, end_status); - reportServerSideStats(stats); - } - - success = handle_exception([&]() { getResults(results); }); - if (!success) - return; - - finish(true); - } -} - -void TritonClient::reportServerSideStats(const TritonClient::ServerSideStats& stats) const { - std::stringstream msg; - - // https://github.com/triton-inference-server/server/blob/v2.3.0/src/clients/c++/perf_client/inference_profiler.cc - const uint64_t count = stats.success_count_; - msg << " Inference count: " << stats.inference_count_ << "\n"; - msg << " Execution count: " << stats.execution_count_ << "\n"; - msg << " Successful request count: " << count << "\n"; - - if (count > 0) { - auto get_avg_us = [count](uint64_t tval) { - constexpr uint64_t us_to_ns = 1000; - return tval / us_to_ns / count; - }; - - const uint64_t cumm_avg_us = get_avg_us(stats.cumm_time_ns_); - const uint64_t queue_avg_us = get_avg_us(stats.queue_time_ns_); - const uint64_t compute_input_avg_us = get_avg_us(stats.compute_input_time_ns_); - const uint64_t compute_infer_avg_us = get_avg_us(stats.compute_infer_time_ns_); - const uint64_t compute_output_avg_us = get_avg_us(stats.compute_output_time_ns_); - const uint64_t compute_avg_us = compute_input_avg_us + compute_infer_avg_us + compute_output_avg_us; - const uint64_t overhead = - (cumm_avg_us > queue_avg_us + compute_avg_us) ? (cumm_avg_us - queue_avg_us - compute_avg_us) : 0; - - msg << " Avg request latency: " << cumm_avg_us << " usec" - << "\n" - << " (overhead " << overhead << " usec + " - << "queue " << queue_avg_us << " usec + " - << "compute input " << compute_input_avg_us << " usec + " - << "compute infer " << compute_infer_avg_us << " usec + " - << "compute output " << compute_output_avg_us << " usec)" << std::endl; - } - - if (!debugName_.empty()) - edm::LogInfo(fullDebugName_) << msg.str(); -} - -TritonClient::ServerSideStats TritonClient::summarizeServerStats(const inference::ModelStatistics& start_status, - const inference::ModelStatistics& end_status) const { - TritonClient::ServerSideStats server_stats; - - server_stats.inference_count_ = end_status.inference_count() - start_status.inference_count(); - server_stats.execution_count_ = end_status.execution_count() - start_status.execution_count(); - server_stats.success_count_ = - end_status.inference_stats().success().count() - start_status.inference_stats().success().count(); - server_stats.cumm_time_ns_ = - end_status.inference_stats().success().ns() - start_status.inference_stats().success().ns(); - server_stats.queue_time_ns_ = end_status.inference_stats().queue().ns() - start_status.inference_stats().queue().ns(); - server_stats.compute_input_time_ns_ = - end_status.inference_stats().compute_input().ns() - start_status.inference_stats().compute_input().ns(); - server_stats.compute_infer_time_ns_ = - end_status.inference_stats().compute_infer().ns() - start_status.inference_stats().compute_infer().ns(); - server_stats.compute_output_time_ns_ = - end_status.inference_stats().compute_output().ns() - start_status.inference_stats().compute_output().ns(); - - return server_stats; -} - -inference::ModelStatistics TritonClient::getServerSideStatus() const { - if (verbose_) { - inference::ModelStatisticsResponse resp; - TRITON_THROW_IF_ERROR(client_->ModelInferenceStatistics(&resp, options_[0].model_name_, options_[0].model_version_), - "getServerSideStatus(): unable to get model statistics", - localService()); - return *(resp.model_stats().begin()); - } - return inference::ModelStatistics{}; -} - -//for fillDescriptions -void TritonClient::fillPSetDescription(edm::ParameterSetDescription& iDesc) { - edm::ParameterSetDescription descClient; - fillBasePSetDescription(descClient); - descClient.add("modelName"); - descClient.add("modelVersion", ""); - descClient.add("modelConfigPath"); - //server parameters should not affect the physics results - descClient.addUntracked("preferredServer", ""); - descClient.addUntracked("timeout"); - descClient.ifValue(edm::ParameterDescription("timeoutUnit", "seconds", false), - edm::allowedValues("seconds", "milliseconds", "microseconds")); - descClient.addUntracked("useSharedMemory", true); - descClient.addUntracked("compression", ""); - descClient.addUntracked>("outputs", {}); - iDesc.add("Client", descClient); -} +#include "FWCore/Concurrency/interface/WaitingTask.h" +#include "FWCore/MessageLogger/interface/MessageLogger.h" +#include "FWCore/ParameterSet/interface/FileInPath.h" +#include "FWCore/ParameterSet/interface/allowedValues.h" +#include "FWCore/ServiceRegistry/interface/Service.h" +#include "FWCore/Utilities/interface/Exception.h" +#include "HeterogeneousCore/SonicTriton/interface/TritonClient.h" +#include "HeterogeneousCore/SonicTriton/interface/TritonException.h" +#include "HeterogeneousCore/SonicTriton/interface/TritonService.h" +#include "HeterogeneousCore/SonicTriton/interface/triton_utils.h" + +#include "grpc_client.h" +#include "grpc_service.pb.h" +#include "model_config.pb.h" + +#include "google/protobuf/text_format.h" +#include "google/protobuf/io/zero_copy_stream_impl.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tc = triton::client; + +namespace { + // Minimal ParameterSet to satisfy SonicClientBase requirements during unit tests + edm::ParameterSet makeMinimalSonicParamsForTest() { + edm::ParameterSet params; + params.addParameter("mode", "PseudoAsync"); + + edm::ParameterSet defaultRetry; + defaultRetry.addParameter("retryType", "RetrySameServerAction"); + defaultRetry.addUntrackedParameter("allowedTries", 0u); + std::vector retryVec{defaultRetry}; + params.addParameter>("Retry", retryVec); + + return params; + } + grpc_compression_algorithm getCompressionAlgo(const std::string& name) { + if (name.empty() or name.compare("none") == 0) + return grpc_compression_algorithm::GRPC_COMPRESS_NONE; + else if (name.compare("deflate") == 0) + return grpc_compression_algorithm::GRPC_COMPRESS_DEFLATE; + else if (name.compare("gzip") == 0) + return grpc_compression_algorithm::GRPC_COMPRESS_GZIP; + else + throw cms::Exception("GrpcCompression") + << "Unknown compression algorithm requested: " << name << " (choices: none, deflate, gzip)"; + } + + std::vector> convertToShared(const std::vector& tmp) { + std::vector> results; + results.reserve(tmp.size()); + std::transform(tmp.begin(), tmp.end(), std::back_inserter(results), [](tc::InferResult* ptr) { + return std::shared_ptr(ptr); + }); + return results; + } +} // namespace + +//based on https://github.com/triton-inference-server/server/blob/v2.3.0/src/clients/c++/examples/simple_grpc_async_infer_client.cc +//and https://github.com/triton-inference-server/server/blob/v2.3.0/src/clients/c++/perf_client/perf_client.cc + +TritonClient::TritonClient(const edm::ParameterSet& params, const std::string& debugName) + : SonicClient(params, debugName, "TritonClient"), + batchMode_(TritonBatchMode::Rectangular), + manualBatchMode_(false), + verbose_(params.getUntrackedParameter("verbose")), + useSharedMemory_(params.getUntrackedParameter("useSharedMemory")), + compressionAlgo_(getCompressionAlgo(params.getUntrackedParameter("compression"))) { + options_.emplace_back(params.getParameter("modelName")); + + edm::Service ts; + + // We save the token to be able to notify the service in case of an exception in the evaluate method. + // The evaluate method can be called outside the frameworks TBB threadpool in the case of a retry. In + // this case the context is not setup to access the service registry, we need the service token to + // create the context. + token_ = edm::ServiceRegistry::instance().presentToken(); + + //Connect to server + updateServer(params.getUntrackedParameter("preferredServer")); + + //set options + options_[0].model_version_ = params.getParameter("modelVersion"); + options_[0].client_timeout_ = params.getUntrackedParameter("timeout"); + //convert to microseconds + const auto& timeoutUnit = params.getUntrackedParameter("timeoutUnit"); + unsigned conversion = 1; + if (timeoutUnit == "seconds") + conversion = 1e6; + else if (timeoutUnit == "milliseconds") + conversion = 1e3; + else if (timeoutUnit == "microseconds") + conversion = 1; + else + throw cms::Exception("Configuration") << "Unknown timeout unit: " << timeoutUnit; + options_[0].client_timeout_ *= conversion; + + //get fixed parameters from local config + inference::ModelConfig localModelConfig; + { + const std::string localModelConfigPath(params.getParameter("modelConfigPath").fullPath()); + int fileDescriptor = open(localModelConfigPath.c_str(), O_RDONLY); + if (fileDescriptor < 0) + throw TritonException("LocalFailure") + << "TritonClient(): unable to open local model config: " << localModelConfigPath; + google::protobuf::io::FileInputStream localModelConfigInput(fileDescriptor); + localModelConfigInput.SetCloseOnDelete(true); + if (!google::protobuf::TextFormat::Parse(&localModelConfigInput, &localModelConfig)) + throw TritonException("LocalFailure") + << "TritonClient(): unable to parse local model config: " << localModelConfigPath; + } + + //check batch size limitations (after i/o setup) + //triton uses max batch size = 0 to denote a model that does not support native batching (using the outer dimension) + //but for models that do support batching (native or otherwise), a given event may set batch size 0 to indicate no valid input is present + //so set the local max to 1 and keep track of "no outer dim" case + maxOuterDim_ = localModelConfig.max_batch_size(); + noOuterDim_ = maxOuterDim_ == 0; + maxOuterDim_ = std::max(1u, maxOuterDim_); + //propagate batch size + setBatchSize(1); + + //compare model checksums to remote config to enforce versioning + inference::ModelConfigResponse modelConfigResponse; + TRITON_THROW_IF_ERROR(client_->ModelConfig(&modelConfigResponse, options_[0].model_name_, options_[0].model_version_), + "TritonClient(): unable to get model config", + localService()); + inference::ModelConfig remoteModelConfig(modelConfigResponse.config()); + + std::map> checksums; + size_t fileCounter = 0; + for (const auto& modelConfig : {localModelConfig, remoteModelConfig}) { + const auto& agents = modelConfig.model_repository_agents().agents(); + auto agent = std::find_if(agents.begin(), agents.end(), [](auto const& a) { return a.name() == "checksum"; }); + if (agent != agents.end()) { + const auto& params = agent->parameters(); + for (const auto& [key, val] : params) { + // only check the requested version + if (key.compare(0, options_[0].model_version_.size() + 1, options_[0].model_version_ + "/") == 0) + checksums[key][fileCounter] = val; + } + } + ++fileCounter; + } + std::vector incorrect; + for (const auto& [key, val] : checksums) { + if (checksums[key][0] != checksums[key][1]) + incorrect.push_back(key); + } + if (!incorrect.empty()) + throw TritonException("ModelVersioning") << "The following files have incorrect checksums on the remote server: " + << triton_utils::printColl(incorrect, ", "); + + //get model info + inference::ModelMetadataResponse modelMetadata; + TRITON_THROW_IF_ERROR(client_->ModelMetadata(&modelMetadata, options_[0].model_name_, options_[0].model_version_), + "TritonClient(): unable to get model metadata", + localService()); + + //get input and output (which know their sizes) + const auto& nicInputs = modelMetadata.inputs(); + const auto& nicOutputs = modelMetadata.outputs(); + + //report all model errors at once + std::stringstream msg; + std::string msg_str; + + //currently no use case is foreseen for a model with zero inputs or outputs + if (nicInputs.empty()) + msg << "Model on server appears malformed (zero inputs)\n"; + + if (nicOutputs.empty()) + msg << "Model on server appears malformed (zero outputs)\n"; + + //stop if errors + msg_str = msg.str(); + if (!msg_str.empty()) + throw cms::Exception("ModelErrors") << msg_str; + + //setup input map + std::stringstream io_msg; + if (verbose_) + io_msg << "Model inputs: " + << "\n"; + for (const auto& nicInput : nicInputs) { + const auto& iname = nicInput.name(); + auto [curr_itr, success] = input_.emplace(std::piecewise_construct, + std::forward_as_tuple(iname), + std::forward_as_tuple(iname, nicInput, this, ts->pid())); + auto& curr_input = curr_itr->second; + if (verbose_) { + io_msg << " " << iname << " (" << curr_input.dname() << ", " << curr_input.byteSize() + << " b) : " << triton_utils::printColl(curr_input.shape()) << "\n"; + } + } + + //allow selecting only some outputs from server + const auto& v_outputs = params.getUntrackedParameter>("outputs"); + std::unordered_set s_outputs(v_outputs.begin(), v_outputs.end()); + + //setup output map + if (verbose_) + io_msg << "Model outputs: " + << "\n"; + for (const auto& nicOutput : nicOutputs) { + const auto& oname = nicOutput.name(); + if (!s_outputs.empty() and s_outputs.find(oname) == s_outputs.end()) + continue; + auto [curr_itr, success] = output_.emplace(std::piecewise_construct, + std::forward_as_tuple(oname), + std::forward_as_tuple(oname, nicOutput, this, ts->pid())); + auto& curr_output = curr_itr->second; + if (verbose_) { + io_msg << " " << oname << " (" << curr_output.dname() << ", " << curr_output.byteSize() + << " b) : " << triton_utils::printColl(curr_output.shape()) << "\n"; + } + if (!s_outputs.empty()) + s_outputs.erase(oname); + } + + //check if any requested outputs were not available + if (!s_outputs.empty()) + throw cms::Exception("MissingOutput") + << "Some requested outputs were not available on the server: " << triton_utils::printColl(s_outputs); + + //print model info + std::stringstream model_msg; + if (verbose_) { + model_msg << "Model name: " << options_[0].model_name_ << "\n" + << "Model version: " << options_[0].model_version_ << "\n" + << "Model max outer dim: " << (noOuterDim_ ? 0 : maxOuterDim_) << "\n"; + edm::LogInfo(fullDebugName_) << model_msg.str() << io_msg.str(); + } +} + +TritonClient::~TritonClient() { + //by default: members of this class destroyed before members of base class + //in shared memory case, TritonMemResource (member of TritonData) unregisters from client_ in its destructor + //but input/output objects are member of base class, so destroyed after client_ (member of this class) + //therefore, clear the maps here + input_.clear(); + output_.clear(); +} + +void TritonClient::setBatchMode(TritonBatchMode batchMode) { + unsigned oldBatchSize = batchSize(); + batchMode_ = batchMode; + manualBatchMode_ = true; + //this allows calling setBatchSize() and setBatchMode() in either order consistently to change back and forth + //includes handling of change from ragged to rectangular if multiple entries already created + setBatchSize(oldBatchSize); +} + +void TritonClient::resetBatchMode() { + batchMode_ = TritonBatchMode::Rectangular; + manualBatchMode_ = false; +} + +unsigned TritonClient::nEntries() const { return !input_.empty() ? input_.begin()->second.entries_.size() : 0; } + +unsigned TritonClient::batchSize() const { return batchMode_ == TritonBatchMode::Rectangular ? outerDim_ : nEntries(); } + +bool TritonClient::setBatchSize(unsigned bsize) { + if (batchMode_ == TritonBatchMode::Rectangular) { + if (bsize > maxOuterDim_) { + throw TritonException("LocalFailure") + << "Requested batch size " << bsize << " exceeds server-specified max batch size " << maxOuterDim_ << "."; + return false; + } else { + outerDim_ = bsize; + //take min to allow resizing to 0 + resizeEntries(std::min(outerDim_, 1u)); + return true; + } + } else { + resizeEntries(bsize); + outerDim_ = 1; + return true; + } +} + +void TritonClient::resizeEntries(unsigned entry) { + if (entry > nEntries()) + //addEntry(entry) extends the vector to size entry+1 + addEntry(entry - 1); + else if (entry < nEntries()) { + for (auto& element : input_) { + element.second.entries_.resize(entry); + } + for (auto& element : output_) { + element.second.entries_.resize(entry); + } + } +} + +void TritonClient::addEntry(unsigned entry) { + for (auto& element : input_) { + element.second.addEntryImpl(entry); + } + for (auto& element : output_) { + element.second.addEntryImpl(entry); + } + if (entry > 0) { + batchMode_ = TritonBatchMode::Ragged; + outerDim_ = 1; + } +} + +void TritonClient::reset() { + if (!manualBatchMode_) + batchMode_ = TritonBatchMode::Rectangular; + for (auto& element : input_) { + element.second.reset(); + } + for (auto& element : output_) { + element.second.reset(); + } +} + +template +bool TritonClient::handle_exception(F&& call) { + //caught exceptions will be propagated to edm::WaitingTaskWithArenaHolder + CMS_SA_ALLOW try { + call(); + return true; + } + //TritonExceptions are intended/expected to be recoverable, i.e. retries should be allowed + catch (TritonException& e) { + e.convertToWarning(); + finish(false); + return false; + } + //other exceptions are not: execution should stop if they are encountered + catch (...) { + finish(false, std::current_exception()); + return false; + } +} + +template +bool TritonClient::handle_exception_holder(edm::WaitingTaskWithArenaHolder& fh, F&& call) { + CMS_SA_ALLOW try { + call(); + return true; + } catch (TritonException& e) { + e.convertToWarning(); + inferSuccess_ = false; + fh.doneWaiting(nullptr); // retryable: schedules finish(false) on TBB + return false; + } catch (...) { + fh.doneWaiting(std::current_exception()); // non-retryable + return false; + } +} + +const TritonService* TritonClient::service() const { + edm::ServiceRegistry::Operate op(token_); + edm::Service ts; + return &(*ts); +} + +const TritonService* TritonClient::localService() const { return isLocal_ ? service() : nullptr; } + +void TritonClient::getResults(const std::vector>& results) { + for (unsigned i = 0; i < results.size(); ++i) { + const auto& result = results[i]; + for (auto& [oname, output] : output_) { + //set shape here before output becomes const + if (output.variableDims()) { + std::vector tmp_shape; + TRITON_THROW_IF_ERROR(result->Shape(oname, &tmp_shape), + "getResults(): unable to get output shape for " + oname); + if (!noOuterDim_) + tmp_shape.erase(tmp_shape.begin()); + output.setShape(tmp_shape, i); + } + //extend lifetime + output.setResult(result, i); + //compute size after getting all result entries + if (i == results.size() - 1) + output.computeSizes(); + } + } +} + +//default case for sync and pseudo async +void TritonClient::evaluate() { + //undo previous signal from TritonException + if (totalTries_ > 0) { + // If we are retrying then the evaluate method is called outside the frameworks TBB thread pool. + // So we need to setup the service token for the current thread to access the service registry. + edm::ServiceRegistry::Operate op(token_); + edm::Service ts; + ts->notifyCallStatus(true); + } + + //in case there is nothing to process + if (batchSize() == 0) { + //call getResults on an empty vector + std::vector> empty_results; + getResults(empty_results); + finish(true); + return; + } + + //set up input pointers for triton (generalized for multi-request ragged batching case) + //one vector per request + unsigned nEntriesVal = nEntries(); + std::vector> inputsTriton(nEntriesVal); + for (auto& inputTriton : inputsTriton) { + inputTriton.reserve(input_.size()); + } + for (auto& [iname, input] : input_) { + for (unsigned i = 0; i < nEntriesVal; ++i) { + inputsTriton[i].push_back(input.data(i)); + } + } + + //set up output pointers similarly + std::vector> outputsTriton(nEntriesVal); + for (auto& outputTriton : outputsTriton) { + outputTriton.reserve(output_.size()); + } + for (auto& [oname, output] : output_) { + for (unsigned i = 0; i < nEntriesVal; ++i) { + outputsTriton[i].push_back(output.data(i)); + } + } + + //set up shared memory for output + auto success = handle_exception([&]() { + for (auto& element : output_) { + element.second.prepare(); + } + }); + //edm::LogInfo("TritonClient") << "evaluate() return 1"; + if (!success) + return; + + // Get the status of the server prior to the request being made. + inference::ModelStatistics start_status; + success = handle_exception([&]() { + if (verbose()) + start_status = getServerSideStatus(); + }); + //edm::LogInfo("TritonClient") << "evaluate() return 2"; + if (!success) + return; + + if (mode_ == SonicMode::Async) { + // Reset before each inference attempt: false=retryable/not-yet-succeeded. + // The callback sets this to true on success before calling doneWaiting(nullptr). + // If the launch throws (lambda destroyed inside AsyncInferMulti), the holder destructor + // calls doneWaiting(nullptr) with inferSuccess_=false → finish(false) [retryable] on TBB. + inferSuccess_ = false; + + // Create holder[finish]: when doneWaiting() is called, finish() is scheduled as a TBB task + // so retry logic (finish→retry→updateServer) never runs on the gRPC callback thread. + auto* finishTask = edm::make_waiting_task([this](std::exception_ptr const* excptr) { + if (excptr) + finish(false, *excptr); // non-retryable exception + else + finish(inferSuccess_.load()); // true=success, false=retryable failure + }); + edm::WaitingTaskWithArenaHolder finishHolder(*holder_->group(), finishTask); + + // Launch async inference. + // On TritonException: fh destructor (in destroyed lambda) calls doneWaiting(nullptr) + // → schedules finish(false) [retryable] on TBB. Do NOT call finish() directly here. + CMS_SA_ALLOW try { + TRITON_THROW_IF_ERROR( + client_->AsyncInferMulti( + [start_status, fh = std::move(finishHolder), this](std::vector resultsTmp) mutable { + //immediately convert to shared_ptr + const auto& results = convertToShared(resultsTmp); + //check results + for (auto ptr : results) { + auto ok = handle_exception_holder(fh, [&]() { + TRITON_THROW_IF_ERROR(ptr->RequestStatus(), "evaluate(): unable to get result(s)", localService()); + }); + if (!ok) + return; + } + + if (verbose()) { + inference::ModelStatistics end_status; + auto ok = handle_exception_holder(fh, [&]() { end_status = getServerSideStatus(); }); + if (!ok) + return; + const auto& stats = summarizeServerStats(start_status, end_status); + reportServerSideStats(stats); + } + + auto ok = handle_exception_holder(fh, [&]() { getResults(results); }); + if (!ok) + return; + + inferSuccess_ = true; + fh.doneWaiting(nullptr); // success: schedules finish(true) on TBB + }, + options_, + inputsTriton, + outputsTriton, + headers_, + compressionAlgo_), + "evaluate(): unable to launch async run", + localService()); + } catch (TritonException& e) { + e.convertToWarning(); + // fh destructor already called doneWaiting(nullptr) → finish(false) scheduled on TBB + } + } else { + //edm::LogInfo("TritonClient") << "evaluate() return 4"; + //blocking call + std::vector resultsTmp; + success = handle_exception([&]() { + TRITON_THROW_IF_ERROR( + client_->InferMulti(&resultsTmp, options_, inputsTriton, outputsTriton, headers_, compressionAlgo_), + "evaluate(): unable to run and/or get result", + localService()); + }); + //immediately convert to shared_ptr + const auto& results = convertToShared(resultsTmp); + if (!success) + return; + + if (verbose()) { + inference::ModelStatistics end_status; + success = handle_exception([&]() { end_status = getServerSideStatus(); }); + if (!success) + return; + + const auto& stats = summarizeServerStats(start_status, end_status); + reportServerSideStats(stats); + } + + success = handle_exception([&]() { getResults(results); }); + if (!success) + return; + + finish(true); + } +} + +void TritonClient::reportServerSideStats(const TritonClient::ServerSideStats& stats) const { + std::stringstream msg; + + // https://github.com/triton-inference-server/server/blob/v2.3.0/src/clients/c++/perf_client/inference_profiler.cc + const uint64_t count = stats.success_count_; + msg << " Inference count: " << stats.inference_count_ << "\n"; + msg << " Execution count: " << stats.execution_count_ << "\n"; + msg << " Successful request count: " << count << "\n"; + + if (count > 0) { + auto get_avg_us = [count](uint64_t tval) { + constexpr uint64_t us_to_ns = 1000; + return tval / us_to_ns / count; + }; + + const uint64_t cumm_avg_us = get_avg_us(stats.cumm_time_ns_); + const uint64_t queue_avg_us = get_avg_us(stats.queue_time_ns_); + const uint64_t compute_input_avg_us = get_avg_us(stats.compute_input_time_ns_); + const uint64_t compute_infer_avg_us = get_avg_us(stats.compute_infer_time_ns_); + const uint64_t compute_output_avg_us = get_avg_us(stats.compute_output_time_ns_); + const uint64_t compute_avg_us = compute_input_avg_us + compute_infer_avg_us + compute_output_avg_us; + const uint64_t overhead = + (cumm_avg_us > queue_avg_us + compute_avg_us) ? (cumm_avg_us - queue_avg_us - compute_avg_us) : 0; + + msg << " Avg request latency: " << cumm_avg_us << " usec" + << "\n" + << " (overhead " << overhead << " usec + " + << "queue " << queue_avg_us << " usec + " + << "compute input " << compute_input_avg_us << " usec + " + << "compute infer " << compute_infer_avg_us << " usec + " + << "compute output " << compute_output_avg_us << " usec)" << std::endl; + } + + if (!debugName_.empty()) + edm::LogInfo(fullDebugName_) << msg.str(); +} + +TritonClient::ServerSideStats TritonClient::summarizeServerStats(const inference::ModelStatistics& start_status, + const inference::ModelStatistics& end_status) const { + TritonClient::ServerSideStats server_stats; + + server_stats.inference_count_ = end_status.inference_count() - start_status.inference_count(); + server_stats.execution_count_ = end_status.execution_count() - start_status.execution_count(); + server_stats.success_count_ = + end_status.inference_stats().success().count() - start_status.inference_stats().success().count(); + server_stats.cumm_time_ns_ = + end_status.inference_stats().success().ns() - start_status.inference_stats().success().ns(); + server_stats.queue_time_ns_ = end_status.inference_stats().queue().ns() - start_status.inference_stats().queue().ns(); + server_stats.compute_input_time_ns_ = + end_status.inference_stats().compute_input().ns() - start_status.inference_stats().compute_input().ns(); + server_stats.compute_infer_time_ns_ = + end_status.inference_stats().compute_infer().ns() - start_status.inference_stats().compute_infer().ns(); + server_stats.compute_output_time_ns_ = + end_status.inference_stats().compute_output().ns() - start_status.inference_stats().compute_output().ns(); + + return server_stats; +} + +inference::ModelStatistics TritonClient::getServerSideStatus() const { + if (verbose_) { + inference::ModelStatisticsResponse resp; + TRITON_THROW_IF_ERROR(client_->ModelInferenceStatistics(&resp, options_[0].model_name_, options_[0].model_version_), + "getServerSideStatus(): unable to get model statistics", + localService()); + return *(resp.model_stats().begin()); + } + return inference::ModelStatistics{}; +} + +void TritonClient::updateServer(const std::string& serverName) { + //get appropriate server for this model + auto ts = service(); + + const auto& serverMap = ts->resolveServer(options_[0].model_name_, serverName); + + const auto& server = serverMap.second; + + //update server name + serverName_ = serverMap.first; + + serverType_ = server.type; + edm::LogInfo("TritonDiscovery") << debugName_ << " assigned server: " << server.url; + //enforce sync mode for fallback CPU server to avoid contention + //todo: could enforce async mode otherwise (unless mode was specified by user?) + if (serverType_ == TritonServerType::LocalCPU) + setMode(SonicMode::Sync); + isLocal_ = serverType_ == TritonServerType::LocalCPU or serverType_ == TritonServerType::LocalGPU; + + // updateServer() is always called from a TBB thread (via finish() -> retry() -> updateServer()), + // so it is safe to destroy the old gRPC client immediately. + client_.reset(); + + //connect to the server + TRITON_THROW_IF_ERROR( + tc::InferenceServerGrpcClient::Create(&client_, server.url, false, server.useSsl, server.sslOptions), + "TritonClient(): unable to create inference context", + localService()); +} + +void TritonClient::switchToFallback() { + edm::ServiceRegistry::Operate op(token_); + edm::Service ts; + + // Start the fallback server if it has not been started yet (idempotent). + ts->startFallbackServer(); + if (!ts->fallbackStarted()) + throw TritonException("LocalFailure") + << "TritonClient::switchToFallback: fallback server is not available " + "(check that TritonService fallback.enable = True and a model path is configured)"; + + // Dynamically load this client's model onto the fallback server. + ts->loadModel(options_[0].model_name_); + + // Re-point the gRPC connection at the fallback server. + updateServer(TritonService::Server::fallbackName); +} + +//for fillDescriptions +void TritonClient::fillPSetDescription(edm::ParameterSetDescription& iDesc) { + edm::ParameterSetDescription descClient; + fillBasePSetDescription(descClient); + descClient.add("modelName"); + descClient.add("modelVersion", ""); + descClient.add("modelConfigPath"); + //server parameters should not affect the physics results + descClient.addUntracked("preferredServer", ""); + descClient.addUntracked("timeout"); + descClient.ifValue(edm::ParameterDescription("timeoutUnit", "seconds", false), + edm::allowedValues("seconds", "milliseconds", "microseconds")); + descClient.addUntracked("useSharedMemory", true); + descClient.addUntracked("compression", ""); + descClient.addUntracked>("outputs", {}); + iDesc.add("Client", descClient); +} + +void TritonClient::connectToServer(const std::string& url) { + // Update client state for a generic remote server + serverType_ = TritonServerType::Remote; + isLocal_ = false; + + edm::LogInfo("TritonDiscovery") << debugName_ << " connecting to server: " << url; + + // Use default SSL options + triton::client::SslOptions sslOptions; + bool useSsl = false; // Assuming no SSL for direct URL connection + + // Connect to the server + TRITON_THROW_IF_ERROR(triton::client::InferenceServerGrpcClient::Create(&client_, url, false, useSsl, sslOptions), + "TritonClient::connectToServer(): unable to create inference context"); +} + +//constructor for testing +TritonClient::TritonClient() : SonicClient(makeMinimalSonicParamsForTest(), "TritonClient_test", "TritonClient") {} diff --git a/HeterogeneousCore/SonicTriton/src/TritonService.cc b/HeterogeneousCore/SonicTriton/src/TritonService.cc index 112e70d1d1ea2..b10ffacadcaa5 100644 --- a/HeterogeneousCore/SonicTriton/src/TritonService.cc +++ b/HeterogeneousCore/SonicTriton/src/TritonService.cc @@ -120,11 +120,14 @@ TritonService::TritonService(const edm::ParameterSet& pset, edm::ActivityRegistr << "TritonService: Not allowed to specify more than one server with same name (" << serverName << ")"; } - //loop over all servers: check which models they have + //loop over all servers: check which models they have, populate serverHealth std::string msg; if (verbose_) msg = "List of models for each server:\n"; for (auto& [serverName, server] : servers_) { + //populate serverHealth + serversHealth_.emplace(serverName, ServerHealth{}); + std::unique_ptr client; TRITON_THROW_IF_ERROR( tc::InferenceServerGrpcClient::Create(&client, server.url, false, server.useSsl, server.sslOptions), @@ -138,7 +141,8 @@ TritonService::TritonService(const edm::ParameterSet& pset, edm::ActivityRegistr edm::LogInfo("TritonService") << "Server " << serverName << ": url = " << server.url << ", version = " << serverMetaResponse.version(); else - edm::LogInfo("TritonService") << "unable to get metadata for " + serverName + " (" + server.url + ")"; + edm::LogInfo("TritonService") << "unable to get metadata for " + serverName + " (" + server.url + ")" + << err.Message(); } //if this query fails, it indicates that the server is nonresponsive or saturated @@ -191,64 +195,214 @@ void TritonService::addModel(const std::string& modelName, const std::string& pa if (!allowAddModel_) throw cms::Exception("DisallowedAddModel") << "TritonService: Attempt to call addModel() outside of module constructors"; - //if model is not in the list, then no specified server provides it - auto mit = models_.find(modelName); - if (mit == models_.end()) { - auto& modelInfo(unservedModels_.emplace(modelName, path).first->second); - modelInfo.modules.insert(currentModuleId_); - //only keep track of modules that need unserved models - modules_.emplace(currentModuleId_, modelName); - } + + auto& modelInfo(models_.emplace(modelName, path).first->second); + // Update path if model was previously added (e.g., by server scanning) with empty path + if (modelInfo.path.empty() && !path.empty()) + modelInfo.path = path; + + modelInfo.modules.insert(currentModuleId_); + modules_.emplace(currentModuleId_, modelName); } void TritonService::postModuleConstruction(edm::ModuleDescription const& desc) { allowAddModel_ = false; } void TritonService::preModuleDestruction(edm::ModuleDescription const& desc) { - //remove destructed modules from unserved list - if (unservedModels_.empty()) - return; auto id = desc.id(); auto oit = modules_.find(id); if (oit != modules_.end()) { const auto& moduleInfo(oit->second); - auto mit = unservedModels_.find(moduleInfo.model); - if (mit != unservedModels_.end()) { + auto mit = models_.find(moduleInfo.model); + if (mit != models_.end()) { auto& modelInfo(mit->second); modelInfo.modules.erase(id); - //remove a model if it is no longer needed by any modules - if (modelInfo.modules.empty()) - unservedModels_.erase(mit); } modules_.erase(oit); } } -//second return value is only true if fallback CPU server is being used -TritonService::Server TritonService::serverInfo(const std::string& model, const std::string& preferred) const { +// Returns the name of the server assigned to serve the given model, or nullptr if no server is currently assigned. +// If a preferred server(current server) is specified but unavailable, falls back to any assigned server. +// Callers are responsible for handling the nullptr case. +const std::string* TritonService::resolveServerName(const std::string& model, const std::string& preferred) const { auto mit = models_.find(model); - if (mit == models_.end()) - throw cms::Exception("MissingModel") << "TritonService: There are no servers that provide model " << model; - const auto& modelInfo(mit->second); - const auto& modelServers = modelInfo.servers; + if (mit == models_.end() || mit->second.servers.empty()) + return nullptr; // no server assigned - caller decides what to do + + const auto& modelServers = mit->second.servers; - auto msit = modelServers.end(); if (!preferred.empty()) { - msit = modelServers.find(preferred); - //todo: add a "strict" parameter to stop execution if preferred server isn't found? - if (msit == modelServers.end()) - edm::LogWarning("PreferredServer") << "Preferred server " << preferred << " for model " << model - << " not available, will choose another server"; + auto msit = modelServers.find(preferred); + if (msit != modelServers.end()) + return &(*msit); + edm::LogWarning("PreferredServer") << "Preferred server " << preferred << " for model " << model + << " not available, will choose another server"; } - const auto& serverName(msit == modelServers.end() ? *modelServers.begin() : preferred); + //Prefer remote servers over fallback if available + if (modelServers.size() > 1) { + auto rit = std::find_if(modelServers.begin(), modelServers.end(), [this](const std::string& name) { + auto sit = servers_.find(name); + return sit != servers_.end() && !sit->second.isFallback; + }); + if (rit != modelServers.end()) + return &(*rit); + } + return &(*modelServers.begin()); +} + +// Returns the full server info for the server assigned to serve the given model. +// Throws MissingModel if no server is currently assigned. +// Wraps resolveServerName; use that directly if nullptr should be handled by the caller +const std::pair& TritonService::resolveServer( + const std::string& model, const std::string& preferred) const { + const auto* name = resolveServerName(model, preferred); + if (!name) + throw cms::Exception("MissingModel") << "TritonService: There are no servers that provide model " << model; + return *servers_.find(*name); +} - //todo: use some algorithm to select server rather than just picking arbitrarily - const auto& server(servers_.find(serverName)->second); - return server; +// Returns the list of model names that are not currently assigned to any server. +std::vector TritonService::unassignedModels() const { + std::vector result; + for (const auto& [name, info] : models_) { + if (info.servers.empty()) + result.push_back(name); + } + return result; } -void TritonService::preBeginJob(edm::ProcessContext const&) { - //only need fallback if there are unserved models - if (!fallbackOpts_.enable or unservedModels_.empty()) +void TritonService::updateServerHealth(const std::string& modelName) const { + for (auto& [serverName, server] : servers_) { + edm::LogInfo("TritonService") << "Updating server health for server = " << serverName; + if (server.isFallback) { + edm::LogInfo("TritonService") << serverName << " is skipped because it is a fallback server"; + continue; // fallback is a last resort, not a candidate for getBestServer + } + try { + std::unique_ptr client; + TRITON_THROW_IF_ERROR( + tc::InferenceServerGrpcClient::Create(&client, server.url, false, server.useSsl, server.sslOptions), + "TritonService(): unable to create inference context for " + serverName + " (" + server.url + ")"); + + bool live = false, ready = false; + TRITON_THROW_IF_ERROR(client->IsServerLive(&live), + "TritonService(): unable to query IsServerLive " + serverName + " (" + server.url + ")"); + TRITON_THROW_IF_ERROR(client->IsServerReady(&ready), + "TritonService(): unable to query IsServerReady " + serverName + " (" + server.url + ")"); + + edm::LogInfo("TritonService") << serverName << " : live = " << live << " ready = " << ready; + + inference::ModelStatisticsResponse stats; + if (!modelName.empty()) { + client->ModelInferenceStatistics(&stats, modelName); + } else { + for (const auto& m : server.models) { + client->ModelInferenceStatistics(&stats, m); + } + } + + uint64_t infer_count = 0, queue_count = 0, failures = 0; + double avgQueueTimeMs = 0.0; + double avgInferTimeMs = 0.0; + + for (const auto& mstat : stats.model_stats()) { + if (modelName.empty() || mstat.name() == modelName) { + const auto& infer = mstat.inference_stats(); + + infer_count += infer.compute_infer().count(); + avgInferTimeMs += infer.compute_infer().ns() / 1e3; + queue_count += infer.queue().count(); + avgQueueTimeMs += infer.queue().ns() / 1e3; + failures += infer.fail().count(); + } + } + // Update health map safely with accessor + tbb::concurrent_hash_map::accessor acc; + serversHealth_.find(acc, serverName); + + ServerHealth& health = acc->second; + health.live = live; + health.ready = ready; + health.failureCount = failures; + health.avgQueueTimeMs = (queue_count > 0) ? avgQueueTimeMs / queue_count : 0.0; + health.avgInferTimeMs = (infer_count > 0) ? avgInferTimeMs / infer_count : 0.0; + + } catch (const TritonException& e) { + // mark existing entry unhealthy if present + tbb::concurrent_hash_map::accessor acc; + if (serversHealth_.find(acc, serverName)) { + ServerHealth& health = acc->second; + health.live = false; + health.ready = false; + } + } catch (const std::exception& e) { + // fallback for other exceptions + tbb::concurrent_hash_map::accessor acc; + if (serversHealth_.find(acc, serverName)) { + ServerHealth& health = acc->second; + health.live = false; + health.ready = false; + } + } + } +} + +std::optional TritonService::getBestServer(const std::string& modelName, + const std::string& IgnoreServer) const { + std::optional bestServerName; + ServerHealth bestHealth; + + // get fresh ServerHealth statistics + updateServerHealth(modelName); + edm::LogInfo("TritonService") << "Getting best server"; + + for (auto& [serverName, server] : servers_) { + if (serverName == IgnoreServer) { + edm::LogInfo("TritonService") << serverName << " is ignored"; + continue; // skip ignored server + } + if (server.isFallback) { + edm::LogInfo("TritonService") << serverName << " is skipped because it is a fallback server"; + continue; // fallback is a last resort, not a candidate for getBestServer + } + if (server.models.find(modelName) == server.models.end()) { + edm::LogInfo("TritonService") << serverName << " is skipped because it does not have " << modelName; + continue; // server doesn't have model + } + + tbb::concurrent_hash_map::const_accessor acc; + if (!serversHealth_.find(acc, serverName)) { + edm::LogInfo("TritonService") << serverName << " is skipped because it does not have health info"; + continue; // no health info + } + + const ServerHealth& health = acc->second; + + if (!health.live || !health.ready) { + edm::LogInfo("TritonService") << serverName << " is skipped because is not live or ready"; + continue; // skip unhealthy + } + + // Select server according to rules: + // 1) lowest failureCount + // 2) tie-breaker: lowest avgQueueTimeMs + if (!bestServerName || health.failureCount < bestHealth.failureCount || + (health.failureCount == bestHealth.failureCount && health.avgQueueTimeMs < bestHealth.avgQueueTimeMs)) { + bestServerName = serverName; + bestHealth = health; + } + } + if (verbose_ && bestServerName) { + edm::LogInfo("TritonDiscovery") << "Chosen server for model '" << modelName << "': " << *bestServerName + << " (failures=" << bestHealth.failureCount + << ", avgQueueTime=" << bestHealth.avgQueueTimeMs << " ms)"; + } + return bestServerName; +} + +void TritonService::startFallbackServer() { + // Idempotent: do nothing if already running or disabled + if (!fallbackOpts_.enable || startedFallback_) return; //include fallback server in set @@ -262,10 +416,13 @@ void TritonService::preBeginJob(edm::ProcessContext const&) { std::string msg; if (verbose_) msg = "List of models for fallback server: "; - //all unserved models are provided by fallback server + // Provide all declared models with known paths via the fallback server auto& server(servers_.find(Server::fallbackName)->second); - for (const auto& [modelName, model] : unservedModels_) { - auto& modelInfo(models_.emplace(modelName, model).first->second); + for (const auto& [modelName, model] : models_) { + // Only seed models for which we have a repository path + if (model.path.empty()) + continue; + auto& modelInfo(models_.find(modelName)->second); modelInfo.servers.insert(Server::fallbackName); server.models.insert(modelName); if (verbose_) @@ -288,7 +445,9 @@ void TritonService::preBeginJob(edm::ProcessContext const&) { fallbackOpts_.command += " -r " + std::to_string(fallbackOpts_.retries); if (fallbackOpts_.wait >= 0) fallbackOpts_.command += " -w " + std::to_string(fallbackOpts_.wait); - for (const auto& [modelName, model] : unservedModels_) { + for (const auto& [modelName, model] : models_) { + if (model.path.empty()) + continue; fallbackOpts_.command += " -m " + model.path; } std::string thread_string = " -I " + std::to_string(numberOfThreads_); @@ -297,8 +456,7 @@ void TritonService::preBeginJob(edm::ProcessContext const&) { fallbackOpts_.command += " -i " + fallbackOpts_.imageName; if (!fallbackOpts_.sandboxDir.empty()) fallbackOpts_.command += " -s " + fallbackOpts_.sandboxDir; - //don't need this anymore - unservedModels_.clear(); + // models_ remains for runtime queries; nothing to clear here //get a random temporary directory if none specified if (fallbackOpts_.tempDir.empty()) { @@ -360,6 +518,25 @@ void TritonService::preBeginJob(edm::ProcessContext const&) { << output; } +void TritonService::preBeginJob(edm::ProcessContext const&) { + // Capture unassigned models *before* startFallbackServer() is called. + // startFallbackServer() seeds all known-path models into the fallback server + // set, which would make unassignedModels() return empty afterward. + const auto& unassigned = unassignedModels(); + + // Always start the fallback server so it is ready for on-demand model + // loading during retries, even when every model has a primary server. + startFallbackServer(); + + if (!unassigned.empty() && startedFallback_) { + auto& server(servers_.find(Server::fallbackName)->second); + for (const auto& modelName : unassigned) { + server.models.insert(modelName); + loadModel(modelName); + } + } +} + void TritonService::notifyCallStatus(bool status) const { if (status) --callFails_; @@ -453,3 +630,107 @@ void TritonService::fillDescriptions(edm::ConfigurationDescriptions& description descriptions.addWithDefaultLabel(desc); } + +bool TritonService::loadModel(const std::string& modelName) { + std::lock_guard lock(modelLoadMutex_); + + // Get model from models_ map (should exist from addModel during module construction) + auto mit = models_.find(modelName); + if (mit == models_.end()) { + edm::LogWarning("TritonService") << "loadModel: Model " << modelName << " not found in models_ map"; + return false; + } + + return loadModel(modelName, mit->second); +} + +bool TritonService::loadModel(const std::string& modelName, Model& model) { + // if already loaded, bump refcount + if (model.refCount > 0) { + ++model.refCount; + if (verbose_) + edm::LogInfo("TritonService") << "Model " << modelName << " already loaded, ref count: " << model.refCount; + return true; + } + + if (!startedFallback_) { + throw cms::Exception("TritonService") + << "loadModel: fallback server not started; cannot load model '" << modelName << "'"; + } + + auto sit = servers_.find(Server::fallbackName); + if (sit == servers_.end()) { + throw cms::Exception("TritonService") << "loadModel: fallback server not found"; + } + + std::unique_ptr client; + TRITON_THROW_IF_ERROR(tc::InferenceServerGrpcClient::Create( + &client, sit->second.url, false, sit->second.useSsl, sit->second.sslOptions), + "loadModel: unable to create client for fallback server"); + + TRITON_THROW_IF_ERROR(client->LoadModel(modelName), + "loadModel: failed to load model " + modelName + " on fallback server"); + + // Update state and tracking + model.refCount = 1; + model.servers.insert(Server::fallbackName); + sit->second.models.insert(modelName); + fallbackLoadedModels_.insert(modelName); + + if (verbose_) + edm::LogInfo("TritonService") << "Successfully loaded model " << modelName << " on fallback server"; + return true; +} + +bool TritonService::unloadModel(const std::string& modelName) { + std::lock_guard lock(modelLoadMutex_); + + // Get model from models_ map + auto mit = models_.find(modelName); + if (mit == models_.end()) { + edm::LogWarning("TritonService") << "unloadModel: Model " << modelName << " not found in models_ map"; + return false; + } + + return unloadModel(modelName, mit->second); +} + +bool TritonService::unloadModel(const std::string& modelName, Model& model) { + if (model.refCount == 0) { + edm::LogWarning("TritonService") << "unloadModel: Model " << modelName << " is not loaded"; + return false; + } + + if (model.refCount > 1) { + --model.refCount; + if (verbose_) + edm::LogInfo("TritonService") << "Model " << modelName << " still in use, ref count: " << model.refCount; + return true; + } + + auto sit = servers_.find(Server::fallbackName); + if (sit == servers_.end()) { + edm::LogWarning("TritonService") << "unloadModel: Fallback server not found"; + return false; + } + + if (verbose_) + edm::LogInfo("TritonService") << "Model " << modelName << " ref count is 1, unloading from fallback server"; + + std::unique_ptr client; + TRITON_THROW_IF_ERROR(tc::InferenceServerGrpcClient::Create( + &client, sit->second.url, false, sit->second.useSsl, sit->second.sslOptions), + "unloadModel: unable to create client for fallback server"); + + TRITON_THROW_IF_ERROR(client->UnloadModel(modelName), + "unloadModel: failed to unload model " + modelName + " from fallback server"); + + model.refCount = 0; + model.servers.erase(Server::fallbackName); + sit->second.models.erase(modelName); + fallbackLoadedModels_.erase(modelName); + + if (verbose_) + edm::LogInfo("TritonService") << "Successfully unloaded model " << modelName << " from fallback server"; + return true; +} diff --git a/HeterogeneousCore/SonicTriton/test/BuildFile.xml b/HeterogeneousCore/SonicTriton/test/BuildFile.xml index e4ff7a0bb56f3..ec32b11b88970 100644 --- a/HeterogeneousCore/SonicTriton/test/BuildFile.xml +++ b/HeterogeneousCore/SonicTriton/test/BuildFile.xml @@ -1,11 +1,31 @@ + + - + + + + + + + + + + + + + + + + + + + diff --git a/HeterogeneousCore/SonicTriton/test/DynamicModelLoadingProducer.cc b/HeterogeneousCore/SonicTriton/test/DynamicModelLoadingProducer.cc new file mode 100644 index 0000000000000..eb885ea29dec1 --- /dev/null +++ b/HeterogeneousCore/SonicTriton/test/DynamicModelLoadingProducer.cc @@ -0,0 +1,83 @@ +#include "HeterogeneousCore/SonicTriton/interface/TritonEDProducer.h" +#include "HeterogeneousCore/SonicTriton/interface/TritonService.h" +#include "DataFormats/TestObjects/interface/ToyProducts.h" +#include "FWCore/Framework/interface/MakerMacros.h" +#include "FWCore/MessageLogger/interface/MessageLogger.h" +#include "FWCore/ServiceRegistry/interface/Service.h" + +#include +#include +#include +#include + +// Test module that explicitly exercises dynamic model loading +// This tests the reference counting and thread safety of loadModel/unloadModel +class DynamicModelLoadingProducer : public TritonEDProducer<> { +public: + explicit DynamicModelLoadingProducer(edm::ParameterSet const& cfg) + : TritonEDProducer<>(cfg), + loadUnloadCycles_(cfg.getParameter("loadUnloadCycles")), + testConcurrency_(cfg.getParameter("testConcurrency")) { + putToken_ = produces(); + } + + void acquire(edm::Event const& iEvent, edm::EventSetup const& iSetup, Input& iInput) override { + edm::Service ts; + const std::string& modelName = client_->modelName(); + + // Test dynamic loading and unloading + if (testConcurrency_) { + // Stress test with multiple rapid load/unload cycles + for (int i = 0; i < loadUnloadCycles_; ++i) { + bool loadResult = ts->loadModel(modelName); + edm::LogInfo("DynamicModelLoadingProducer") + << "Load attempt " << i << ": " << (loadResult ? "success" : "failed"); + + // Small delay to allow other threads to interleave + if (i % 5 == 0) { + std::this_thread::yield(); + } + + bool unloadResult = ts->unloadModel(modelName); + edm::LogInfo("DynamicModelLoadingProducer") + << "Unload attempt " << i << ": " << (unloadResult ? "success" : "failed"); + } + } else { + // Simple test: load once, unload once + bool loadResult = ts->loadModel(modelName); + edm::LogInfo("DynamicModelLoadingProducer") << "Single load: " << (loadResult ? "success" : "failed"); + + bool unloadResult = ts->unloadModel(modelName); + edm::LogInfo("DynamicModelLoadingProducer") << "Single unload: " << (unloadResult ? "success" : "failed"); + } + + // Fill dummy input - use actual input from the model (gat_test expects "x" input) + // This is just to satisfy the base class requirements, not for actual inference + auto& input_x = iInput.at("x"); + auto data_x = input_x.allocate(); + // Minimal dummy data + (*data_x)[0] = std::vector{1.0f}; + input_x.setShape(0, 1, 0); + input_x.toServer(data_x); + } + + void produce(edm::Event& iEvent, edm::EventSetup const& iSetup, Output const& iOutput) override { + // Produce dummy output + iEvent.emplace(putToken_, loadUnloadCycles_); + } + + static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) { + edm::ParameterSetDescription desc; + TritonClient::fillPSetDescription(desc); + desc.add("loadUnloadCycles", 1); + desc.add("testConcurrency", false); + descriptions.addWithDefaultLabel(desc); + } + +private: + int loadUnloadCycles_; + bool testConcurrency_; + edm::EDPutTokenT putToken_; +}; + +DEFINE_FWK_MODULE(DynamicModelLoadingProducer); diff --git a/HeterogeneousCore/SonicTriton/test/RefCount.cc b/HeterogeneousCore/SonicTriton/test/RefCount.cc new file mode 100644 index 0000000000000..735b334542e1a --- /dev/null +++ b/HeterogeneousCore/SonicTriton/test/RefCount.cc @@ -0,0 +1,199 @@ +#define CATCH_CONFIG_MAIN +#include "catch2/catch_all.hpp" + +#include +#include +#include + +// Standalone refcount logic test +// This tests the refcount algorithm without requiring the full TritonService infrastructure + +// Simplified model state for testing refcount logic +struct TestModelState { + std::string modelName; + std::string path; + int refCount{0}; + bool isLoaded() const { return refCount > 0; } +}; + +// Mock class that implements the same refcount logic as TritonService +class RefCountManager { +public: + // Tracks actual server load/unload calls + int serverLoadCalls{0}; + int serverUnloadCalls{0}; + + // Simulates loadModel behavior + bool loadModel(const std::string& modelName, const std::string& path = "") { + std::lock_guard lock(mutex_); + + auto& state = models_[modelName]; + if (state.modelName.empty()) + state.modelName = modelName; + if (state.path.empty() && !path.empty()) + state.path = path; + + // If already loaded, just bump refcount (no server call) + if (state.refCount > 0) { + ++state.refCount; + refCounts_[modelName] = state.refCount; + return true; + } + + // Actually "load" on server (simulated) + ++serverLoadCalls; + state.refCount = 1; + refCounts_[modelName] = state.refCount; + return true; + } + + // Simulates unloadModel behavior + bool unloadModel(const std::string& modelName) { + std::lock_guard lock(mutex_); + + auto it = models_.find(modelName); + if (it == models_.end() || it->second.refCount == 0) { + return false; // Not loaded + } + + auto& state = it->second; + + // If refcount > 1, just decrement (no server call) + if (state.refCount > 1) { + --(state.refCount); + refCounts_[modelName] = state.refCount; + return true; + } + + // Actually "unload" from server (simulated) + ++serverUnloadCalls; + refCounts_.erase(modelName); + state.refCount = 0; + return true; + } + + int getRefCount(const std::string& modelName) const { + auto it = refCounts_.find(modelName); + return (it != refCounts_.end()) ? it->second : 0; + } + +private: + std::unordered_map models_; + std::unordered_map refCounts_; + std::mutex mutex_; +}; + +TEST_CASE("RefCount: single load increments to 1", "[RefCount]") { + RefCountManager mgr; + + REQUIRE(mgr.loadModel("model_a", "/path/to/model_a")); + REQUIRE(mgr.getRefCount("model_a") == 1); + REQUIRE(mgr.serverLoadCalls == 1); +} + +TEST_CASE("RefCount: multiple loads increment without server calls", "[RefCount]") { + RefCountManager mgr; + + // First load - should call server + REQUIRE(mgr.loadModel("model_a")); + REQUIRE(mgr.getRefCount("model_a") == 1); + REQUIRE(mgr.serverLoadCalls == 1); + + // Second load - should NOT call server, just increment + REQUIRE(mgr.loadModel("model_a")); + REQUIRE(mgr.getRefCount("model_a") == 2); + REQUIRE(mgr.serverLoadCalls == 1); // Still 1 + + // Third load - should NOT call server, just increment + REQUIRE(mgr.loadModel("model_a")); + REQUIRE(mgr.getRefCount("model_a") == 3); + REQUIRE(mgr.serverLoadCalls == 1); // Still 1 +} + +TEST_CASE("RefCount: unload decrements without server call until zero", "[RefCount]") { + RefCountManager mgr; + + // Load 3 times + mgr.loadModel("model_a"); + mgr.loadModel("model_a"); + mgr.loadModel("model_a"); + REQUIRE(mgr.getRefCount("model_a") == 3); + REQUIRE(mgr.serverLoadCalls == 1); + + // First unload - decrement only, no server call + REQUIRE(mgr.unloadModel("model_a")); + REQUIRE(mgr.getRefCount("model_a") == 2); + REQUIRE(mgr.serverUnloadCalls == 0); + + // Second unload - decrement only, no server call + REQUIRE(mgr.unloadModel("model_a")); + REQUIRE(mgr.getRefCount("model_a") == 1); + REQUIRE(mgr.serverUnloadCalls == 0); + + // Third unload - should call server (refcount reaches 0) + REQUIRE(mgr.unloadModel("model_a")); + REQUIRE(mgr.getRefCount("model_a") == 0); + REQUIRE(mgr.serverUnloadCalls == 1); +} + +TEST_CASE("RefCount: unload on non-loaded model returns false", "[RefCount]") { + RefCountManager mgr; + + // Unload without loading first + REQUIRE_FALSE(mgr.unloadModel("model_a")); + REQUIRE(mgr.serverUnloadCalls == 0); +} + +TEST_CASE("RefCount: reload after full unload triggers new server load", "[RefCount]") { + RefCountManager mgr; + + // Load and fully unload + mgr.loadModel("model_a"); + mgr.unloadModel("model_a"); + REQUIRE(mgr.getRefCount("model_a") == 0); + REQUIRE(mgr.serverLoadCalls == 1); + REQUIRE(mgr.serverUnloadCalls == 1); + + // Reload - should call server again + REQUIRE(mgr.loadModel("model_a")); + REQUIRE(mgr.getRefCount("model_a") == 1); + REQUIRE(mgr.serverLoadCalls == 2); // Now 2 +} + +TEST_CASE("RefCount: multiple models are independent", "[RefCount]") { + RefCountManager mgr; + + // Load two different models + mgr.loadModel("model_a"); + mgr.loadModel("model_b"); + REQUIRE(mgr.getRefCount("model_a") == 1); + REQUIRE(mgr.getRefCount("model_b") == 1); + REQUIRE(mgr.serverLoadCalls == 2); + + // Load model_a again + mgr.loadModel("model_a"); + REQUIRE(mgr.getRefCount("model_a") == 2); + REQUIRE(mgr.getRefCount("model_b") == 1); + REQUIRE(mgr.serverLoadCalls == 2); // No new server call + + // Unload model_b completely + mgr.unloadModel("model_b"); + REQUIRE(mgr.getRefCount("model_a") == 2); + REQUIRE(mgr.getRefCount("model_b") == 0); + REQUIRE(mgr.serverUnloadCalls == 1); + + // model_a still loaded + REQUIRE(mgr.getRefCount("model_a") == 2); +} + +TEST_CASE("RefCount: path is preserved from first load", "[RefCount]") { + RefCountManager mgr; + + // First load with path + mgr.loadModel("model_a", "/path/to/model"); + REQUIRE(mgr.getRefCount("model_a") == 1); + + // Second load without path - should still work + mgr.loadModel("model_a"); + REQUIRE(mgr.getRefCount("model_a") == 2); +} diff --git a/HeterogeneousCore/SonicTriton/test/RetryActionDiffServer.cc b/HeterogeneousCore/SonicTriton/test/RetryActionDiffServer.cc new file mode 100644 index 0000000000000..9952218a82fc3 --- /dev/null +++ b/HeterogeneousCore/SonicTriton/test/RetryActionDiffServer.cc @@ -0,0 +1,71 @@ +#define CATCH_CONFIG_MAIN +#include "catch2/catch_all.hpp" + +#include "HeterogeneousCore/SonicTriton/interface/RetryActionDiffServer.h" +#include "HeterogeneousCore/SonicTriton/interface/TritonClient.h" +#include "HeterogeneousCore/SonicTriton/interface/TritonService.h" +#include "HeterogeneousCore/SonicCore/interface/RetryActionBase.h" + +#include "FWCore/ParameterSet/interface/ParameterSet.h" + +#include + +// Test double for TritonClient to observe updateServer calls without framework/services +class TestTritonClient : public TritonClient { +public: + TestTritonClient() : TritonClient() {} + + void connectToServer(const std::string& url) override { lastConnectedUrl = url; } + + void updateServer(const std::string& serverName) override { lastUpdatedServerName = serverName; } + + const std::string& lastUrl() const { return lastConnectedUrl; } + const std::string& lastServerName() const { return lastUpdatedServerName; } + +protected: + void evaluate() override {} + +private: + std::string lastConnectedUrl; + std::string lastUpdatedServerName; +}; + +TEST_CASE("RetryActionDiffServer switches to fallback via updateServer", "[RetryActionDiffServer]") { + edm::ParameterSet empty; + TestTritonClient client; + + RetryActionDiffServer action(empty, static_cast(&client)); + + // start should arm the action + action.start(); + REQUIRE(action.shouldRetry()); + + // retry should call updateServer with fallback name then disarm + action.retry(); + REQUIRE(client.lastServerName() == TritonService::Server::fallbackName); + + // second retry without re-arming should be a no-op: lastServerName unchanged + std::string afterFirst = client.lastServerName(); + action.retry(); + REQUIRE(client.lastServerName() == afterFirst); +} + +// A client that throws during updateServer to exercise error handling path +class ThrowingTritonClient : public TritonClient { +public: + ThrowingTritonClient() : TritonClient() {} + void updateServer(const std::string&) override { throw TritonException("updateServer failure"); } + +protected: + void evaluate() override {} +}; + +TEST_CASE("RetryActionDiffServer catches exceptions from updateServer", "[RetryActionDiffServer]") { + edm::ParameterSet empty; + ThrowingTritonClient client; + RetryActionDiffServer action(empty, static_cast(&client)); + action.start(); + + // Should not throw despite client throwing internally; action disarms afterward + REQUIRE_NOTHROW(action.retry()); +} diff --git a/HeterogeneousCore/SonicTriton/test/tritonTest_cfg.py b/HeterogeneousCore/SonicTriton/test/tritonTest_cfg.py index 33d6a9c60aad4..7a5e88a4d6d6a 100644 --- a/HeterogeneousCore/SonicTriton/test/tritonTest_cfg.py +++ b/HeterogeneousCore/SonicTriton/test/tritonTest_cfg.py @@ -9,6 +9,7 @@ "TritonGraphFilter": ["gat_test"], "TritonGraphAnalyzer": ["gat_test"], "TritonIdentityProducer": ["ragged_io"], + "DynamicModelLoadingProducer": ["gat_test"], } # other choices @@ -21,6 +22,9 @@ parser.add_argument("--brief", default=False, action="store_true", help="briefer output for graph modules") parser.add_argument("--unittest", default=False, action="store_true", help="unit test mode: reduce input sizes") parser.add_argument("--testother", default=False, action="store_true", help="also test gRPC communication if shared memory enabled, or vice versa") +parser.add_argument("--loadUnloadCycles", default=3, type=int, help="number of load/unload cycles for dynamic model loading test") +parser.add_argument("--testConcurrency", default=False, action="store_true", help="enable concurrent stress test for dynamic model loading") +options = parser.parse_args() options = getOptions(parser, verbose=True) @@ -83,6 +87,10 @@ processModule.edgeMin = cms.uint32(8000) processModule.edgeMax = cms.uint32(15000) processModule.brief = cms.bool(options.brief) + elif module=="DynamicModelLoadingProducer": + # Configure dynamic model loading test (requires explicit model control mode, enabled by default in cmsTriton) + processModule.loadUnloadCycles = cms.int32(options.loadUnloadCycles) + processModule.testConcurrency = cms.bool(options.testConcurrency) process.p += processModule if options.testother: # clone modules to test both gRPC and shared memory