fix: Fix race condition in model loading/unloading in Sagemaker server (#8571)

yinggeh · yinggeh · commit 4f02dbcb0cb5 · 2025-12-11T11:27:52.000-08:00
diff --git a/src/sagemaker_server.cc b/src/sagemaker_server.cc
@@ -646,6 +646,7 @@ SagemakerAPIServer::SageMakerMMEUnloadModel(
     target_model = model_name_hash;
   }
 
+  std::lock_guard<std::mutex> lock(models_list_mutex_);
   if (sagemaker_models_list_.find(model_name_hash) ==
       sagemaker_models_list_.end()) {
     LOG_VERBOSE(1) << "Model " << target_model << " with model hash "
@@ -740,7 +741,6 @@ SagemakerAPIServer::SageMakerMMEUnloadModel(
 
   TRITONSERVER_ErrorDelete(unregister_err);
 
-  std::lock_guard<std::mutex> lock(models_list_mutex_);
   sagemaker_models_list_.erase(model_name_hash);
 }
 

Original file line number	Diff line number	Diff line change
`@@ -646,6 +646,7 @@ SagemakerAPIServer::SageMakerMMEUnloadModel(`
`646`	`646`	`target_model = model_name_hash;`
`647`	`647`	`}`
`648`	`648`
	`649`	`+ std::lock_guard<std::mutex> lock(models_list_mutex_);`
`649`	`650`	`if (sagemaker_models_list_.find(model_name_hash) ==`
`650`	`651`	`sagemaker_models_list_.end()) {`
`651`	`652`	`LOG_VERBOSE(1) << "Model " << target_model << " with model hash "`
`@@ -740,7 +741,6 @@ SagemakerAPIServer::SageMakerMMEUnloadModel(`
`740`	`741`
`741`	`742`	`TRITONSERVER_ErrorDelete(unregister_err);`
`742`	`743`
`743`		`- std::lock_guard<std::mutex> lock(models_list_mutex_);`
`744`	`744`	`sagemaker_models_list_.erase(model_name_hash);`
`745`	`745`	`}`
`746`	`746`