Repository: SeldonIO/seldon-core (language: python)

Issue:
fix(agent): allow serial order for servers connection per instance
<!--
Thanks for sending a pull request! 
If this is your first time, please read our contributor guidelines:
https://docs.seldon.io/projects/seldon-core/en/latest/developer/contributing.html
-->

**What this PR does / why we need it**:

In some cases during rolling upgrades of servers we can get a situation where the scheduler detects the disconnection of an old version of a server _after_ the new server had connected and in this case we dont want to actually act upon this disconnection.

We do that by enforcing a serial order per instance, meaning that the current server has to disconnect before the new one can connect.

There are 2 components affected: agent server and chainer server.

**Which issue(s) this PR fixes**:
<!--
Automatically closes linked issue when PR is merged.
Usage: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
-->
Fixes # INFRA-1226 (internal)

**Special notes for your reviewer**:

Tested by unit tests


Discussion / hints:

Fix patch:
diff --git a/scheduler/data-flow/src/main/kotlin/io/seldon/dataflow/Cli.kt b/scheduler/data-flow/src/main/kotlin/io/seldon/dataflow/Cli.kt
index 1e5485aaea..94a5e393a5 100644
--- a/scheduler/data-flow/src/main/kotlin/io/seldon/dataflow/Cli.kt
+++ b/scheduler/data-flow/src/main/kotlin/io/seldon/dataflow/Cli.kt
@@ -11,6 +11,7 @@ package io.seldon.dataflow
 
 import com.natpryce.konfig.CommandLineOption
 import com.natpryce.konfig.Configuration
+import com.natpryce.konfig.ConfigurationMap
 import com.natpryce.konfig.ConfigurationProperties
 import com.natpryce.konfig.EnvironmentVariables
 import com.natpryce.konfig.Key
@@ -25,6 +26,8 @@ import io.klogging.Level
 import io.klogging.noCoLogger
 import io.seldon.dataflow.kafka.security.KafkaSaslMechanisms
 import io.seldon.dataflow.kafka.security.KafkaSecurityProtocols
+import java.net.InetAddress
+import java.util.UUID
 
 object Cli {
     private const val ENV_VAR_PREFIX = "SELDON_"
@@ -34,6 +37,7 @@ object Cli {
     val logLevelApplication = Key("log.level.app", enumType(*Level.values()))
     val logLevelKafka = Key("log.level.kafka", enumType(*Level.values()))
     val namespace = Key("pod.namespace", stringType)
+    val dataflowReplicaId = Key("dataflow.replica.id", stringType)
 
     // Seldon components
     val upstreamHost = Key("upstream.host", stringType)
@@ -75,6 +79,7 @@ object Cli {
             logLevelApplication,
             logLevelKafka,
             namespace,
+            dataflowReplicaId,
             upstreamHost,
             upstreamPort,
             kafkaBootstrapServers,
@@ -105,10 +110,26 @@ object Cli {
 
     fun configWith(rawArgs: Array<String>): Configuration {
         val fromProperties = ConfigurationProperties.fromResource("local.properties")
+        val fromSystem = getSystemConfig()
         val fromEnv = EnvironmentVariables(prefix = ENV_VAR_PREFIX)
         val fromArgs = parseArguments(rawArgs)
 
-        return fromArgs overriding fromEnv overriding fromProperties
+        return fromArgs overriding fromEnv overriding fromSystem overriding fromProperties
+    }
+
+    private fun getSystemConfig(): Configuration {
+        val dataflowIdPair = this.dataflowReplicaId to getNewDataflowId()
+        return ConfigurationMap(dataflowIdPair)
+    }
+
+    fun getNewDataflowId(assignRandomUuid: Boolean = false): String {
+        if (!assignRandomUuid) {
+            try {
+                return InetAddress.getLocalHost().hostName
+            } catch (_: Exception) {
+            }
+        }
+        return "seldon-dataflow-engine-" + UUID.randomUUID().toString()
     }
 
     private fun parseArguments(rawArgs: Array<String>): Configuration {
diff --git a/scheduler/data-flow/src/main/kotlin/io/seldon/dataflow/Main.kt b/scheduler/data-flow/src/main/kotlin/io/seldon/dataflow/Main.kt
index 8d4a899eaa..b064a974f5 100644
--- a/scheduler/data-flow/src/main/kotlin/io/seldon/dataflow/Main.kt
+++ b/scheduler/data-flow/src/main/kotlin/io/seldon/dataflow/Main.kt
@@ -102,9 +102,11 @@ object Main {
                 describeRetries = config[Cli.topicDescribeRetries],
                 describeRetryDelayMillis = config[Cli.topicDescribeRetryDelayMillis],
             )
+        val subscriberId = config[Cli.dataflowReplicaId]
+
         val subscriber =
             PipelineSubscriber(
-                "seldon-dataflow-engine",
+                subscriberId,
                 kafkaProperties,
                 kafkaAdminProperties,
                 kafkaStreamsParams,
diff --git a/scheduler/data-flow/src/main/resources/local.properties b/scheduler/data-flow/src/main/resources/local.properties
index 46a7218380..68b3bd408d 100644
--- a/scheduler/data-flow/src/main/resources/local.properties
+++ b/scheduler/data-flow/src/main/resources/local.properties
@@ -1,5 +1,6 @@
 log.level.app=INFO
 log.level.kafka=WARN
+dataflow.replica.id=seldon-dataflow-engine
 kafka.bootstrap.servers=localhost:9092
 kafka.consumer.prefix=
 kafka.security.protocol=PLAINTEXT


Test patch:
diff --git a/scheduler/data-flow/src/test/kotlin/io/seldon/dataflow/CliTest.kt b/scheduler/data-flow/src/test/kotlin/io/seldon/dataflow/CliTest.kt
index 9011ff3d4c..52a97fa4b5 100644
--- a/scheduler/data-flow/src/test/kotlin/io/seldon/dataflow/CliTest.kt
+++ b/scheduler/data-flow/src/test/kotlin/io/seldon/dataflow/CliTest.kt
@@ -16,9 +16,15 @@ import org.junit.jupiter.params.provider.Arguments
 import org.junit.jupiter.params.provider.Arguments.arguments
 import org.junit.jupiter.params.provider.MethodSource
 import strikt.api.expectCatching
+import strikt.api.expectThat
+import strikt.assertions.hasLength
 import strikt.assertions.isEqualTo
+import strikt.assertions.isNotEqualTo
 import strikt.assertions.isSuccess
+import strikt.assertions.startsWith
+import java.util.UUID
 import java.util.stream.Stream
+import kotlin.test.Test
 
 internal class CliTest {
     @DisplayName("Passing auth mechanism via cli argument")
@@ -36,6 +42,31 @@ internal class CliTest {
             .isEqualTo(expectedMechanism)
     }
 
+    @Test
+    fun `should handle dataflow replica id`() {
+        val cliDefault = Cli.configWith(arrayOf<String>())
+        val testReplicaId = "dataflow-id-1"
+        val cli = Cli.configWith(arrayOf("--dataflow-replica-id", testReplicaId))
+
+        expectThat(cliDefault[Cli.dataflowReplicaId]) {
+            isNotEqualTo("seldon-dataflow-engine")
+        }
+        expectThat(cli[Cli.dataflowReplicaId]) {
+            isEqualTo(testReplicaId)
+        }
+
+        // test random Uuid (v4)
+        val expectedReplicaIdPrefix = "seldon-dataflow-engine-"
+        val uuidStringLength = 36
+        val randomReplicaUuid = Cli.getNewDataflowId(true)
+        expectThat(randomReplicaUuid) {
+            startsWith(expectedReplicaIdPrefix)
+            hasLength(expectedReplicaIdPrefix.length + uuidStringLength)
+        }
+        expectCatching { UUID.fromString(randomReplicaUuid.removePrefix(expectedReplicaIdPrefix)) }
+            .isSuccess()
+    }
+
     companion object {
         @JvmStatic
         private fun saslMechanisms(): Stream<Arguments> {