Skip to content

Commit

Permalink
fix memory circuit breaker
Browse files Browse the repository at this point in the history
Signed-off-by: Xun Zhang <[email protected]>
  • Loading branch information
Zhangxunmt committed Jul 11, 2023
1 parent 66750bb commit abdccc0
Show file tree
Hide file tree
Showing 3 changed files with 11 additions and 4 deletions.
1 change: 0 additions & 1 deletion plugin/build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -287,7 +287,6 @@ List<String> jacocoExclusions = [
'org.opensearch.ml.action.training.TrainingITTests',
'org.opensearch.ml.action.prediction.PredictionITTests',
'org.opensearch.ml.cluster.MLSyncUpCron',
'org.opensearch.ml.breaker.MemoryCircuitBreaker',
'org.opensearch.ml.model.MLModelGroupManager',
'org.opensearch.ml.helper.ModelAccessControlHelper',
'org.opensearch.ml.action.models.DeleteModelTransportAction.2'
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ public CircuitBreaker getBreaker(BreakerName name) {
*/
public MLCircuitBreakerService init(Path path) {
// Register memory circuit breaker
registerBreaker(BreakerName.MEMORY, new MemoryCircuitBreaker(this.jvmService));
registerBreaker(BreakerName.MEMORY, new MemoryCircuitBreaker(this.settings, this.clusterService, this.jvmService));
log.info("Registered ML memory breaker.");
registerBreaker(BreakerName.DISK, new DiskCircuitBreaker(path.toString()));
log.info("Registered ML disk breaker.");
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
package org.opensearch.ml.breaker;

import static org.mockito.Mockito.when;
import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_JVM_HEAP_MEM_THRESHOLD;
import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_NATIVE_MEM_THRESHOLD;

import java.nio.file.Path;
Expand Down Expand Up @@ -95,8 +96,15 @@ public void testClearBreakers() {

@Test
public void testInit() {
Settings settings = Settings.builder().put(ML_COMMONS_NATIVE_MEM_THRESHOLD.getKey(), 90).build();
ClusterSettings clusterSettings = new ClusterSettings(settings, new HashSet<>(Arrays.asList(ML_COMMONS_NATIVE_MEM_THRESHOLD)));
Settings settings = Settings
.builder()
.put(ML_COMMONS_NATIVE_MEM_THRESHOLD.getKey(), 90)
.put(ML_COMMONS_JVM_HEAP_MEM_THRESHOLD.getKey(), 95)
.build();
ClusterSettings clusterSettings = new ClusterSettings(
settings,
new HashSet<>(Arrays.asList(ML_COMMONS_NATIVE_MEM_THRESHOLD, ML_COMMONS_JVM_HEAP_MEM_THRESHOLD))
);
when(clusterService.getClusterSettings()).thenReturn(clusterSettings);
mlCircuitBreakerService = new MLCircuitBreakerService(jvmService, osService, settings, clusterService);
Assert.assertNotNull(mlCircuitBreakerService.init(Path.of("/")));
Expand Down

0 comments on commit abdccc0

Please sign in to comment.