resolves #1881 NLP: add healthcheck for Sagemaker classifiers (#1882)

Pazuzzu · vsct-jburet · Pazuzzu · web-flow · commit a853ebae951b · 2025-05-02T07:43:46.000+02:00
* resolves #1881 NLP: add healthcheck for Sagemaker classifiers * resolves #1881 NLP: address reviews * resolves #1881 NLP: fix case, where no intialized sagemaker client returns OK * resolves #1881 NLP: address client buiding reviews --------- Co-authored-by: Julien Buret <jburet@voyages-sncf.com> Co-authored-by: Pazuzzu <bogaforlife@gmai.com>
diff --git a/nlp/api/service/src/main/kotlin/NlpVerticle.kt b/nlp/api/service/src/main/kotlin/NlpVerticle.kt
@@ -16,6 +16,10 @@
 
 package ai.tock.nlp.api
 
+/**
+ *
+ */
+// Add this import at the top with the other imports
 import ai.tock.nlp.front.client.FrontClient
 import ai.tock.nlp.front.service.UnknownApplicationException
 import ai.tock.nlp.front.shared.codec.ApplicationDump
@@ -28,6 +32,7 @@ import ai.tock.nlp.front.shared.merge.ValuesMergeQuery
 import ai.tock.nlp.front.shared.monitoring.MarkAsUnknownQuery
 import ai.tock.nlp.front.shared.monitoring.ParseRequestLogCountQuery
 import ai.tock.nlp.front.shared.parser.ParseQuery
+import ai.tock.nlp.model.service.NlpClassifierService
 import ai.tock.shared.Executor
 import ai.tock.shared.TOCK_FRONT_DATABASE
 import ai.tock.shared.TOCK_MODEL_DATABASE
@@ -46,9 +51,6 @@ import mu.KotlinLogging
 import org.litote.kmongo.Id
 import java.util.Locale
 
-/**
- *
- */
 class NlpVerticle : WebVerticle() {
 
     private val protectPath = verticleBooleanProperty("tock_nlp_protect_path", false)
@@ -59,6 +61,9 @@ class NlpVerticle : WebVerticle() {
 
     private val executor: Executor by injector.instance()
 
+    // Add this line to access NlpClassifierService
+    private val nlpClassifierService = NlpClassifierService
+
     override val logger: KLogger = KotlinLogging.logger {}
 
     override fun authProvider(): TockAuthProvider? {
@@ -232,7 +237,7 @@ class NlpVerticle : WebVerticle() {
             Pair("duckling_service", { FrontClient.healthcheck() }),
             Pair("tock_front_database", { pingMongoDatabase(TOCK_FRONT_DATABASE) }),
             Pair("tock_model_database", { pingMongoDatabase(TOCK_MODEL_DATABASE) })
-        )
+        ) + NlpClassifierService.healthcheck()
     )
 }
 
diff --git a/nlp/core/service/src/main/kotlin/NlpCoreService.kt b/nlp/core/service/src/main/kotlin/NlpCoreService.kt
@@ -16,6 +16,7 @@
 
 package ai.tock.nlp.core.service
 
+
 import ai.tock.nlp.core.CallContext
 import ai.tock.nlp.core.Entity
 import ai.tock.nlp.core.EntityRecognition
@@ -40,8 +41,8 @@ import ai.tock.nlp.model.ModelNotInitializedException
 import ai.tock.nlp.model.NlpClassifier
 import ai.tock.shared.checkMaxLengthAllowed
 import ai.tock.shared.error
-import ai.tock.shared.normalize
 import ai.tock.shared.injector
+import ai.tock.shared.normalize
 import com.github.salomonbrys.kodein.instance
 import mu.KotlinLogging
 
diff --git a/nlp/model/sagemaker/pom.xml b/nlp/model/sagemaker/pom.xml
@@ -36,6 +36,10 @@
             <groupId>software.amazon.awssdk</groupId>
             <artifactId>sagemakerruntime</artifactId>
         </dependency>
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>sagemaker</artifactId>
+        </dependency>
         <dependency>
             <groupId>software.amazon.awssdk</groupId>
             <artifactId>sts</artifactId>
diff --git a/nlp/model/sagemaker/src/main/kotlin/SagemakerAwsClient.kt b/nlp/model/sagemaker/src/main/kotlin/SagemakerAwsClient.kt
@@ -17,13 +17,18 @@ package ai.tock.nlp.sagemaker
 
 import ai.tock.shared.jackson.mapper
 import com.fasterxml.jackson.module.kotlin.readValue
+import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider
 import software.amazon.awssdk.core.SdkBytes
+import software.amazon.awssdk.services.sagemaker.SageMakerClient
+import software.amazon.awssdk.services.sagemaker.model.DescribeEndpointRequest
 import software.amazon.awssdk.services.sagemakerruntime.SageMakerRuntimeClient
 import software.amazon.awssdk.services.sagemakerruntime.model.InvokeEndpointRequest
 import java.nio.charset.Charset
 
 class SagemakerAwsClient(private val configuration: SagemakerAwsClientProperties) {
 
+    val name = configuration.name
+
     // for intentions and entities
     data class ParsedRequest(
         val text: String
@@ -56,6 +61,10 @@ class SagemakerAwsClient(private val configuration: SagemakerAwsClientProperties
         .region(configuration.region)
         .build()
 
+    private val sagemakerClient: SageMakerClient = SageMakerClient.builder()
+        .region(configuration.region)
+        .build()
+
     fun parseIntent(request: ParsedRequest) = invokeSageMakerIntentEndpoint(request.text)
 
     fun parseEntities(request: ParsedRequest): ParsedEntitiesResponse = invokeSageMakerEntitiesEndpoint(request.text)
@@ -84,4 +93,12 @@ class SagemakerAwsClient(private val configuration: SagemakerAwsClientProperties
         val entities = mapper.readValue<List<ParsedEntity>>(response.body().asInputStream())
         return ParsedEntitiesResponse(entities)
     }
+
+    fun healthcheck(): Boolean {
+        val endpointRequest = DescribeEndpointRequest.builder()
+            .endpointName(configuration.endpointName)
+            .build()
+        val response = sagemakerClient.describeEndpoint(endpointRequest)
+       return response.endpointStatus() == software.amazon.awssdk.services.sagemaker.model.EndpointStatus.IN_SERVICE
+    }
 }
diff --git a/nlp/model/sagemaker/src/main/kotlin/SagemakerAwsClientProperties.kt b/nlp/model/sagemaker/src/main/kotlin/SagemakerAwsClientProperties.kt
@@ -21,6 +21,7 @@ import software.amazon.awssdk.regions.Region
 internal fun String.unescapeSagemakerName(): String = replace("___", ":")
 
 data class SagemakerAwsClientProperties(
+    val name: String,
     val region: Region,
     val endpointName: String,
     val contentType: String,
diff --git a/nlp/model/sagemaker/src/main/kotlin/SagemakerClientProvider.kt b/nlp/model/sagemaker/src/main/kotlin/SagemakerClientProvider.kt
@@ -22,4 +22,7 @@ internal object SagemakerClientProvider {
 
     fun getClient(conf: SagemakerAwsClientProperties): SagemakerAwsClient =
         clientMap.getOrPut(conf) { SagemakerAwsClient(conf) }
+
+    fun getAllClient(): MutableCollection<SagemakerAwsClient> =
+        clientMap.values
 }
diff --git a/nlp/model/sagemaker/src/main/kotlin/SagemakerClientType.kt b/nlp/model/sagemaker/src/main/kotlin/SagemakerClientType.kt
@@ -0,0 +1,26 @@
+/*
+ * Copyright (C) 2017/2024 e-voyageurs technologies
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package ai.tock.nlp.sagemaker
+
+/**
+ * Enum representing the different types of Sagemaker clients.
+ */
+enum class SagemakerClientType(val clientName: String) {
+    INTENT_CLASSIFICATION("intent-classification"),
+    ENTITY_CLASSIFICATION("entity-classification");
+
+    override fun toString(): String = clientName
+}
diff --git a/nlp/model/sagemaker/src/main/kotlin/SagemakerEngineProvider.kt b/nlp/model/sagemaker/src/main/kotlin/SagemakerEngineProvider.kt
@@ -16,6 +16,7 @@
 
 package ai.tock.nlp.sagemaker
 
+import NlpHealthcheckResult
 import ai.tock.nlp.core.NlpEngineType
 import ai.tock.nlp.model.TokenizerContext
 import ai.tock.nlp.model.service.engine.EntityClassifier
@@ -56,4 +57,24 @@ class SagemakerEngineProvider : NlpEngineProvider {
         // do not tokenize anything at this stage
         override fun tokenize(context: TokenizerContext, text: String): Array<String> = arrayOf(text)
     }
+
+    override fun healthcheck(): () -> NlpHealthcheckResult = {
+        val clients = SagemakerClientProvider.getAllClient()
+        if (clients.isEmpty()) {
+            NlpHealthcheckResult(
+                entityClassifier = false,
+                intentClassifier = false
+            )
+        } else {
+            val grouped = clients.groupBy { it.name }.withDefault { emptyList() }
+            val entityClients = grouped.getValue(SagemakerEntityClassifier.CLIENT_TYPE.clientName)
+            val intentClients = grouped.getValue(SagemakerIntentClassifier.CLIENT_TYPE.clientName)
+
+            NlpHealthcheckResult(
+                entityClassifier = entityClients.isNotEmpty() && entityClients.all { it.healthcheck() },
+                intentClassifier = intentClients.isNotEmpty() && intentClients.all { it.healthcheck() }
+            )
+        }
+    }
 }
+
diff --git a/nlp/model/sagemaker/src/main/kotlin/SagemakerEntityClassifier.kt b/nlp/model/sagemaker/src/main/kotlin/SagemakerEntityClassifier.kt
@@ -26,7 +26,11 @@ import ai.tock.nlp.model.service.engine.NlpEntityClassifier
 import ai.tock.shared.property
 import software.amazon.awssdk.regions.Region
 
+
 internal class SagemakerEntityClassifier(model: EntityModelHolder) : NlpEntityClassifier(model) {
+    companion object {
+        val CLIENT_TYPE = SagemakerClientType.ENTITY_CLASSIFICATION
+    }
 
     override fun classifyEntities(
         context: EntityCallContext,
@@ -35,6 +39,7 @@ internal class SagemakerEntityClassifier(model: EntityModelHolder) : NlpEntityCl
     ): List<EntityRecognition> {
         SagemakerClientProvider.getClient(
             SagemakerAwsClientProperties(
+                CLIENT_TYPE.clientName,
                 Region.of(property("tock_sagemaker_aws_region_name", "eu-west-3")),
                 property("tock_sagemaker_aws_entities_endpoint_name", "default"),
                 property("tock_sagemaker_aws_content_type", "application/json"),
diff --git a/nlp/model/sagemaker/src/main/kotlin/SagemakerIntentClassifier.kt b/nlp/model/sagemaker/src/main/kotlin/SagemakerIntentClassifier.kt
@@ -15,19 +15,24 @@
  */
 package ai.tock.nlp.sagemaker
 
-import ai.tock.nlp.sagemaker.SagemakerAwsClient.ParsedRequest
 import ai.tock.nlp.core.Intent
 import ai.tock.nlp.core.IntentClassification
 import ai.tock.nlp.model.IntentContext
 import ai.tock.nlp.model.service.engine.IntentClassifier
+import ai.tock.nlp.sagemaker.SagemakerAwsClient.ParsedRequest
 import ai.tock.shared.property
 import software.amazon.awssdk.regions.Region
 
+
 internal class SagemakerIntentClassifier(private val conf: SagemakerModelConfiguration) : IntentClassifier {
+    companion object {
+        val CLIENT_TYPE = SagemakerClientType.INTENT_CLASSIFICATION
+    }
 
     override fun classifyIntent(context: IntentContext, text: String, tokens: Array<String>): IntentClassification {
         return SagemakerClientProvider.getClient(
             SagemakerAwsClientProperties(
+                CLIENT_TYPE.clientName,
                 Region.of(property("tock_sagemaker_aws_region_name", "eu-west-3")),
                 property("tock_sagemaker_aws_intent_endpoint_name", "default"),
                 property("tock_sagemaker_aws_content_type", "application/json"),
@@ -36,7 +41,6 @@ internal class SagemakerIntentClassifier(private val conf: SagemakerModelConfigu
         ).parseIntent(ParsedRequest(text))
             .run {
                 object : IntentClassification {
-
                     var probability = 0.0
                     val iterator = intent_ranking.iterator()
 
diff --git a/nlp/model/service/src/main/kotlin/NlpClassifierService.kt b/nlp/model/service/src/main/kotlin/NlpClassifierService.kt
@@ -39,6 +39,8 @@ import ai.tock.nlp.model.service.engine.IntentModelHolder
 import ai.tock.nlp.model.service.engine.NlpEngineRepository
 import ai.tock.nlp.model.service.engine.NlpEngineRepository.getModelBuilder
 import ai.tock.nlp.model.service.engine.NlpEngineRepository.getModelIo
+import ai.tock.nlp.model.service.engine.NlpEngineRepository.getProvider
+import ai.tock.nlp.model.service.engine.NlpEngineRepository.registeredNlpEngineTypes
 import ai.tock.nlp.model.service.engine.NlpModelRepository
 import ai.tock.nlp.model.service.engine.NlpModelRepository.saveEntityModel
 import ai.tock.nlp.model.service.engine.NlpModelRepository.saveIntentModel
@@ -235,4 +237,14 @@ object NlpClassifierService : NlpClassifier {
     ) {
         nlpApplicationConfigurationDAO.saveNewConfiguration(applicationName, engineType, configuration)
     }
+
+    fun healthcheck(): List<Pair<String, () -> Boolean>> {
+        return registeredNlpEngineTypes().flatMap {
+            val healthcheck = getProvider(it).healthcheck()
+            listOf(
+                "${it.name}-entity" to { healthcheck().entityClassifier },
+                "${it.name}-intent" to { healthcheck().intentClassifier },
+            )
+        }
+    }
 }
diff --git a/nlp/model/service/src/main/kotlin/engine/NlpEngineProvider.kt b/nlp/model/service/src/main/kotlin/engine/NlpEngineProvider.kt
@@ -16,6 +16,7 @@
 
 package ai.tock.nlp.model.service.engine
 
+import NlpHealthcheckResult
 import ai.tock.nlp.core.NlpEngineType
 
 /**
@@ -53,4 +54,10 @@ interface NlpEngineProvider {
      * Returns the tokenizer from this [TokenizerModelHolder].
      */
     fun getTokenizer(model: TokenizerModelHolder): Tokenizer
+
+    /**
+     * Check if the NLP engine is healthy.
+     * @return true if the engine is healthy, false otherwise
+     */
+    fun healthcheck(): () -> NlpHealthcheckResult = { NlpHealthcheckResult.ALL_OK }
 }
diff --git a/nlp/model/service/src/main/kotlin/engine/NlpEngineRepository.kt b/nlp/model/service/src/main/kotlin/engine/NlpEngineRepository.kt
@@ -90,6 +90,7 @@ internal object NlpEngineRepository {
         }
     }
 
+
     fun <T : ClassifierContextKey> getModelBuilder(context: ClassifierContext<T>): NlpEngineModelBuilder {
         return getProvider(context.engineType).modelBuilder
     }
diff --git a/nlp/model/service/src/main/kotlin/engine/NlpEntityClassifier.kt b/nlp/model/service/src/main/kotlin/engine/NlpEntityClassifier.kt
@@ -22,4 +22,5 @@ import ai.tock.shared.ThreadSafe
  *
  */
 @ThreadSafe
-abstract class NlpEntityClassifier(val model: EntityModelHolder) : EntityClassifier
+abstract class NlpEntityClassifier(val model: EntityModelHolder) : EntityClassifier {
+}
diff --git a/nlp/model/service/src/main/kotlin/engine/NlpHealthcheckResult.kt b/nlp/model/service/src/main/kotlin/engine/NlpHealthcheckResult.kt
@@ -0,0 +1,21 @@
+/*
+ * Copyright (C) 2017/2021 e-voyageurs technologies
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+data class NlpHealthcheckResult(val entityClassifier: Boolean, val intentClassifier: Boolean) {
+    companion object {
+        val ALL_OK = NlpHealthcheckResult(entityClassifier = true, intentClassifier = true)
+    }
+}
diff --git a/pom.xml b/pom.xml
@@ -619,6 +619,11 @@
                 <artifactId>sagemakerruntime</artifactId>
                 <version>${aws-sagemaker}</version>
             </dependency>
+            <dependency>
+                <groupId>software.amazon.awssdk</groupId>
+                <artifactId>sagemaker</artifactId>
+                <version>${aws-sagemaker}</version>
+            </dependency>
             <dependency>
                 <groupId>software.amazon.awssdk</groupId>
                 <artifactId>sts</artifactId>

Original file line number	Diff line number	Diff line change
`@@ -22,4 +22,7 @@ internal object SagemakerClientProvider {`
`22`	`22`
`23`	`23`	`fun getClient(conf: SagemakerAwsClientProperties): SagemakerAwsClient =`
`24`	`24`	`clientMap.getOrPut(conf) { SagemakerAwsClient(conf) }`
	`25`	`+`
	`26`	`+ fun getAllClient(): MutableCollection<SagemakerAwsClient> =`
	`27`	`+ clientMap.values`
`25`	`28`	`}`