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

yinggeh · web-flow · commit f3bdcc7db41b · 2025-12-11T11:26:35.000-08:00
diff --git a/src/sagemaker_server.cc b/src/sagemaker_server.cc
@@ -649,6 +649,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 "
@@ -743,7 +744,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
`@@ -649,6 +649,7 @@ SagemakerAPIServer::SageMakerMMEUnloadModel(`
`649`	`649`	`target_model = model_name_hash;`
`650`	`650`	`}`
`651`	`651`
	`652`	`+ std::lock_guard<std::mutex> lock(models_list_mutex_);`
`652`	`653`	`if (sagemaker_models_list_.find(model_name_hash) ==`
`653`	`654`	`sagemaker_models_list_.end()) {`
`654`	`655`	`LOG_VERBOSE(1) << "Model " << target_model << " with model hash "`
`@@ -743,7 +744,6 @@ SagemakerAPIServer::SageMakerMMEUnloadModel(`
`743`	`744`
`744`	`745`	`TRITONSERVER_ErrorDelete(unregister_err);`
`745`	`746`
`746`		`- std::lock_guard<std::mutex> lock(models_list_mutex_);`
`747`	`747`	`sagemaker_models_list_.erase(model_name_hash);`
`748`	`748`	`}`
`749`	`749`