Add retry mechanism and dynamic model loading to SONIC framework by kakwok · Pull Request #27 · fastmachinelearning/cmssw

kakwok · 2026-06-24T21:28:50Z

PR description:

Introduces a pluggable retry mechanism for SONIC inference clients and dynamic
model loading/unloading on the Triton fallback server. This enables automatic
recovery when a remote Triton server becomes unavailable during event processing.

Retry mechanism (`SonicCore`)

RetryActionBase: Plugin factory base class for retry strategies, with
retry(), start(), finish() interface
RetrySameServerAction: Retries inference on the same server (configurable
allowedTries)
SonicClientBase::finish(): Drives retry loop — on retryable failure
(finish(false)), iterates through registered retry actions; on non-retryable
failure (finish(false, eptr)), propagates exception directly

Retry with server failover (`SonicTriton`)

RetryActionDiffServer: On failure, queries TritonService::getBestServer()
for an alternative healthy remote server, calls updateServer() + eval() to
retry on the new server
RetryFallbackServerAction: Last-resort action — when all other retries are
exhausted, lazily starts the fallback server (idempotent), dynamically loads the
model to the fallback server via TritonClient::switchToFallback(), and retries locally. Fires at most
once per inference call.
Server health tracking: TritonService maintains per-server health stats
(liveness, readiness, failure count, queue time) via updateServerHealth(),
used by getBestServer() to select the best candidate
Server selection preference: resolveServerName() prefers remote (non-fallback)
servers when multiple servers provide the same model; getBestServer() excludes
fallback servers from candidate pool

Async retry redesign (holder-based)

evaluate() (Async mode) creates an inner WaitingTaskWithArenaHolder per
inference attempt — the gRPC callback calls doneWaiting() and returns
immediately; finish()/retry()/updateServer() are scheduled as TBB tasks

Dynamic model loading (`TritonService`)

loadModel()/unloadModel() with reference counting for the fallback server
startFallbackServer() is idempotent; only started at preBeginJob for
unassigned models (no remote server available)
Fallback server health entry properly registered in serversHealth_

Configuration

Retry actions configurable via Retry VPSet in client config
(retryType, allowedTries)
customize.py updated with retry options, defaulted to add RetryFallbackServerAction as the last action.

PR validation:

retry_diffServer_fallback.log

…olve cmsTriton conflict in CMSSW_17_0_0_pre2. Co-authored-by: Trevin Lee <trl008@ucsd.edu>

kpedro88 · 2026-06-29T22:45:22Z

+<test name="TestHeterogeneousCoreSonicTritonRetryActionSame" command="cmsRun ${LOCALTOP}/src/HeterogeneousCore/SonicTriton/test/tritonTest_cfg.py --modules TritonGraphProducer  --maxEvents 2 --unittest --device cpu --retryAction same"/>
+<test name="TestHeterogeneousCoreSonicTritonRetryActionDiff" command="cmsRun ${LOCALTOP}/src/HeterogeneousCore/SonicTriton/test/tritonTest_cfg.py --modules TritonGraphProducer  --maxEvents 2 --unittest --device cpu --retryAction diff"/>


do these tests work without a driver script to enable/disable servers?

They will pass but they does not do what they suggest to test.
The retryDiff test will require some fiddling because the server choice is still undeterministic among remote servers.
I can commit a retry same script retry_action_same.sh for now.

done in b9a4f8e

I think we really need the RetryActionDiff test as a unit test if we're going to claim that functionality works. For now, maybe we can come up with some arbitrary but deterministic way to do the server choice. Alternatively, could try to have the test script parse the CMSSW output and figure out which server was chosen, in order to kill it and force the retry...

done in kakwok@c9b55a1
Server is stored with std::set<std::string>, which is sorted lexicographically.
Note that the server name from the tritonTest_cfg command in the test script is used to pick the server, not the initial server names from cmsTriton.

The script tests the following situation:

Having two remote servers server1, server2 and starting the fallback server.

picks server1

script kills server1, client continues with server2

script kills server2, client continues with fallback server.

We can add other test scenarios if desired.

kpedro88 · 2026-06-29T22:46:15Z

We will also have to update the various reco/miniaod algorithm configuration files that currently set allowedTries = cms.untracked.uint32(0)

kakwok · 2026-07-14T20:07:52Z

Configs are updated as well Update allowedTries config

kakwok · 2026-07-15T13:53:15Z

+    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)


@kpedro88 On the topic of configuration, should we add a RetryFallbackServerAction to all configurations as default?

Probably. We may be able to do that by modifying https://github.com/cms-sw/cmssw/blob/b8fdcd5f6d1e240d3966e67d062d3880c2a33eea/HeterogeneousCore/SonicTriton/src/TritonClient.cc#L586

rather than having to make more modifications to a growing list of algorithm-specific config files...

kpedro88 · 2026-07-24T19:53:44Z

+def getOptions(parser, verbose=False):
+    options = parser.parse_args()


duplicated lines?

kpedro88 · 2026-07-24T20:02:33Z

@@ -126,6 +131,6 @@ void MyClient::fillPSetDescription(edm::ParameterSetDescription& iDesc) {

 As indicated, the `fillBasePSetDescription()` function should always be applied to the `descClient` object,
 to ensure that it includes the necessary parameters.
-(Calling `fillBasePSetDescription(descClient, false)` will omit the `allowedTries` parameter, disabling retries.)
+(Calling `fillBasePSetDescription(descClient, false)` will omit the `Retry` parameter, disabling retries.)


I'm not sure this behavior is still desirable. Since Retry is a vector, it could just always exist, but be an empty vector if retries are disabled.

kpedro88 · 2026-07-24T20:06:39Z

+  //enforce sync mode for fallback CPU server to avoid contention
+  if (serverType_ == TritonServerType::LocalCPU)
+    setMode(SonicMode::Sync);
+  if (serverType_ == TritonServerType::Remote)


this should be else (async is also desired for LocalGPU)
however, this then overrides the user setting in all cases, making the mode Python parameter pointless. ideally, we should have some default value and only override to Async if the default is not provided. (Sync should always be used for LocalCPU, even if the user provides a different value.)

kakwok mentioned this pull request Jun 24, 2026

Add retry mechanism and dynamic model loading to SONIC framework #26

Closed

Add retry mechanism and dynamic model loading to SONIC framework. Res…

36ca055

…olve cmsTriton conflict in CMSSW_17_0_0_pre2. Co-authored-by: Trevin Lee <trl008@ucsd.edu>

kakwok force-pushed the SonicRetryDML_CMSSW_17_0_0_pre2 branch from 0ea0d7a to 36ca055 Compare June 29, 2026 14:12

kpedro88 reviewed Jun 29, 2026

View reviewed changes

Martin added 4 commits July 14, 2026 09:21

PR comments edit1

53716fa

Fix json support

f43154b

add retry action test script and clean up tests

b9a4f8e

Update allowedTries config

14a5e29

kakwok commented Jul 15, 2026

View reviewed changes

kpedro88 reviewed Jul 24, 2026

View reviewed changes

Martin added 2 commits July 24, 2026 16:00

Add RetryActionDiff test script. Model needs to be ordered.

c9b55a1

remove commented code

e1f2c2f

		<test name="TestHeterogeneousCoreSonicTritonRetryActionSame" command="cmsRun ${LOCALTOP}/src/HeterogeneousCore/SonicTriton/test/tritonTest_cfg.py --modules TritonGraphProducer --maxEvents 2 --unittest --device cpu --retryAction same"/>
		<test name="TestHeterogeneousCoreSonicTritonRetryActionDiff" command="cmsRun ${LOCALTOP}/src/HeterogeneousCore/SonicTriton/test/tritonTest_cfg.py --modules TritonGraphProducer --maxEvents 2 --unittest --device cpu --retryAction diff"/>

		def getOptions(parser, verbose=False):
		options = parser.parse_args()

Uh oh!

Conversation

kakwok commented Jun 24, 2026

PR description:

Retry mechanism (SonicCore)

Retry with server failover (SonicTriton)

Async retry redesign (holder-based)

Dynamic model loading (TritonService)

Configuration

PR validation:

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kpedro88 commented Jun 29, 2026

Uh oh!

kakwok commented Jul 14, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Retry mechanism (`SonicCore`)

Retry with server failover (`SonicTriton`)

Dynamic model loading (`TritonService`)