Add runner

ancavar · ancavar · commit 1aa2501aa559 · 2024-12-04T18:53:59.000+03:00
diff --git a/buildSrc/src/main/kotlin/Dependencies.kt b/buildSrc/src/main/kotlin/Dependencies.kt
@@ -241,6 +241,12 @@ object Libs {
         version = Versions.onnxruntime
     )
 
+    val onnxruntime_gpu = dep(
+        group = "com.microsoft.onnxruntime",
+        name = "onnxruntime_gpu",
+        version = Versions.onnxruntime
+    )
+
     val ktor_server_core = dep(
         group = "io.ktor",
         name = "ktor-server-core",
diff --git a/settings.gradle.kts b/settings.gradle.kts
@@ -21,6 +21,7 @@ findProject(":usvm-python:usvm-python-runner")?.name = "usvm-python-runner"
 include("usvm-python:usvm-python-commons")
 findProject(":usvm-python:usvm-python-commons")?.name = "usvm-python-commons"
 include("usvm-ml-gameserver")
+include("usvm-runner")
 
 pluginManagement {
     resolutionStrategy {
diff --git a/usvm-core/build.gradle.kts b/usvm-core/build.gradle.kts
@@ -13,6 +13,7 @@ dependencies {
     testImplementation(Libs.ksmt_yices)
 
     implementation(Libs.onnxruntime)
+    implementation(Libs.onnxruntime_gpu)
 }
 
 publishing {
diff --git a/usvm-core/src/main/kotlin/org/usvm/utils/OnnxModelImpl.kt b/usvm-core/src/main/kotlin/org/usvm/utils/OnnxModelImpl.kt
@@ -4,18 +4,31 @@ import ai.onnxruntime.OnnxTensor
 import ai.onnxruntime.OrtEnvironment
 import ai.onnxruntime.OrtSession
 import org.usvm.StateId
+import org.usvm.logger
 import org.usvm.statistics.BasicBlock
 import org.usvm.statistics.BlockGraph
 import org.usvm.util.OnnxModel
 import java.nio.FloatBuffer
 import java.nio.LongBuffer
-import java.nio.file.Path
+
+enum class Mode {
+    CPU,
+    GPU
+}
 
 class OnnxModelImpl<Block: BasicBlock>(
-    pathToONNX: Path
+    pathToONNX: String,
+    mode: Mode
 ): OnnxModel<Game<Block>> {
     private val env: OrtEnvironment = OrtEnvironment.getEnvironment()
-    private val session: OrtSession = env.createSession(pathToONNX.toString())
+    private val sessionOptions: OrtSession.SessionOptions = OrtSession.SessionOptions().apply {
+        if (mode == Mode.GPU) {
+            addCUDA()
+        } else {
+            logger.info("Using CPU execution provider.")
+        }
+    }
+    private val session: OrtSession = env.createSession(pathToONNX, sessionOptions)
 
     override fun predictState(game: Game<Block>): UInt {
         val stateIds = mutableMapOf<StateId, Int>()
diff --git a/usvm-ml-gameserver/build.gradle.kts b/usvm-ml-gameserver/build.gradle.kts
@@ -32,6 +32,7 @@ tasks.test {
 dependencies {
     implementation(project(":usvm-jvm"))
     implementation(project(":usvm-core"))
+    implementation(project(":usvm-runner"))
     implementation(Libs.jacodb_api_jvm)
     implementation(Libs.jacodb_core)
 
@@ -44,6 +45,5 @@ dependencies {
     implementation(Libs.kotlinx_cli)
     implementation(Libs.slf4j_simple)
 
-    implementation("io.github.rchowell:dotlin:1.0.2")
     testImplementation(kotlin("test"))
 }
diff --git a/usvm-ml-gameserver/src/main/kotlin/org.usvm/JavaMethodRunner.kt b/usvm-ml-gameserver/src/main/kotlin/org.usvm/JavaMethodRunner.kt
diff --git a/usvm-ml-gameserver/src/main/kotlin/org.usvm/OracleImpl.kt b/usvm-ml-gameserver/src/main/kotlin/org.usvm/OracleImpl.kt
@@ -4,7 +4,7 @@ import org.usvm.gameserver.GameState
 import org.usvm.statistics.BasicBlock
 import org.usvm.util.Oracle
 import org.usvm.utils.Game
-import org.usvm.util.createGameState
+import org.usvm.gameserver.serialization.createGameState
 
 class OracleImpl<Block: BasicBlock>(
     val predict: (GameState) -> UInt
diff --git a/usvm-ml-gameserver/src/main/kotlin/org.usvm/gameserver/Explorer.kt b/usvm-ml-gameserver/src/main/kotlin/org.usvm/gameserver/Explorer.kt
@@ -1,9 +1,10 @@
-@file:Suppress("UNUSED_PARAMETER", "UNUSED_VARIABLE")
-
 package org.usvm.gameserver
 
-import org.usvm.JavaMethodRunner
+import org.usvm.runner.JavaMethodRunner
 import org.usvm.OracleImpl
+import org.usvm.PathSelectionStrategy
+import org.usvm.UMachineOptions
+import org.usvm.runner.defaultOptions
 import org.usvm.statistics.BasicBlock
 import kotlin.math.floor
 
@@ -22,10 +23,10 @@ fun randomExplorer(
         step.gameStep.stateId
     }
 
-    val (className, methodName) = extractClassAndMethod(gameMap.nameOfObjectToCover)
-    val runner = JavaMethodRunner(gameMap, OracleImpl<BasicBlock>(predict))
+    val options = cloneDefaultOptions(gameMap, predict)
+    val runner = JavaMethodRunner(options, gameMap.assemblyFullName)
 
-    val (results, percentageCoverage) = runner.cover(className, methodName)
+    val (results, percentageCoverage) = runner.cover(gameMap.nameOfObjectToCover)
     val errors = results.count { it.isExceptional }
     val tests = results.size - errors
 
@@ -36,10 +37,16 @@ fun randomExplorer(
     )
 }
 
-private fun extractClassAndMethod(fullName: String): Pair<String, String> {
-    val parts = fullName.split('.')
-    val className = parts.dropLast(1).joinToString(".")
-    val methodName = parts.last()
-
-    return Pair(className, methodName)
+private fun cloneDefaultOptions(gameMap: GameMap, predict: (GameState) -> UInt): UMachineOptions {
+    val defaultSearcher = when (gameMap.defaultSearcher) {
+        Searcher.BFS -> PathSelectionStrategy.BFS
+        Searcher.DFS -> PathSelectionStrategy.DFS
+    }
+    val stepLimit = (gameMap.stepsToStart + gameMap.stepsToStart).toULong()
+    return defaultOptions.copy(
+        pathSelectionStrategies = listOf(defaultSearcher, PathSelectionStrategy.AI),
+        stepLimit = stepLimit,
+        stepsToStart = gameMap.stepsToStart,
+        oracle = OracleImpl<BasicBlock>(predict)
+    )
 }
diff --git a/usvm-ml-gameserver/src/main/kotlin/org.usvm/gameserver/serialization/Conversion.kt b/usvm-ml-gameserver/src/main/kotlin/org.usvm/gameserver/serialization/Conversion.kt
@@ -1,4 +1,4 @@
-package org.usvm.util
+package org.usvm.gameserver.serialization
 
 import org.usvm.gameserver.GameEdgeLabel
 import org.usvm.gameserver.GameMapEdge
diff --git a/usvm-runner/build.gradle.kts b/usvm-runner/build.gradle.kts
@@ -0,0 +1,32 @@
+import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
+
+plugins {
+    id("usvm.kotlin-conventions")
+    kotlin
+}
+
+repositories {
+    mavenCentral()
+}
+
+tasks.test {
+    useJUnitPlatform()
+}
+
+tasks.withType<KotlinCompile> {
+    kotlinOptions.allWarningsAsErrors = false
+}
+
+dependencies {
+    implementation(project(":usvm-jvm"))
+    implementation(project(":usvm-core"))
+    implementation(Libs.jacodb_api_jvm)
+    implementation(Libs.jacodb_core)
+    implementation(Libs.kotlinx_cli)
+
+    implementation(Libs.slf4j_simple)
+    implementation(Libs.logback)
+
+    implementation("io.github.rchowell:dotlin:1.0.2")
+    testImplementation(kotlin("test"))
+}
diff --git a/usvm-runner/src/main/kotlin/org/usvm/runner/JavaMethodRunner.kt b/usvm-runner/src/main/kotlin/org/usvm/runner/JavaMethodRunner.kt
@@ -0,0 +1,57 @@
+package org.usvm.runner
+
+import org.usvm.CoverageZone
+import org.usvm.PathSelectionStrategy
+import org.usvm.PathSelectorCombinationStrategy
+import org.usvm.UMachineOptions
+import org.usvm.machine.JcMachine
+import org.usvm.machine.state.JcState
+import org.usvm.runner.util.JacoDBContainer
+import org.usvm.runner.util.getJcMethodByName
+import org.usvm.runner.util.loadClasspathFromJar
+import org.usvm.runner.util.mapsKey
+import kotlin.time.Duration
+
+class JavaMethodRunner(val options: UMachineOptions, classpath: String) {
+    val jacodbCpKey: String
+        get() = mapsKey
+
+    val cp by lazy {
+        val classpath = loadClasspathFromJar(classpath)
+        JacoDBContainer(jacodbCpKey, classpath).cp
+    }
+
+    fun cover(methodFullName: String): Pair<List<JcState>, Float> {
+        val (className, methodName) = extractClassAndMethod(methodFullName)
+        return cover(className, methodName)
+    }
+
+    fun cover(className: String, methodName: String): Pair<List<JcState>, Float> {
+        val jcMethod = cp.getJcMethodByName(className, methodName)
+
+        val (states, percentageCoverage) = JcMachine(cp, options, interpreterObserver = null).use { machine ->
+            machine.analyze(jcMethod.method, targets = emptyList())
+        }
+        return states to percentageCoverage
+    }
+
+    private fun extractClassAndMethod(fullName: String): Pair<String, String> {
+        val parts = fullName.split('.')
+        val className = parts.dropLast(1).joinToString(".")
+        val methodName = parts.last()
+
+        return Pair(className, methodName)
+    }
+}
+
+val defaultOptions = UMachineOptions(
+    pathSelectionStrategies = listOf(PathSelectionStrategy.AI),
+    pathSelectorCombinationStrategy = PathSelectorCombinationStrategy.SEQUENTIAL,
+    coverageZone = CoverageZone.METHOD,
+    exceptionsPropagation = true,
+    solverTimeout = Duration.INFINITE,
+    timeout = Duration.INFINITE,
+    typeOperationsTimeout = Duration.INFINITE,
+    useSolverForForks = false,
+    stepLimit = 2000U
+)
diff --git a/usvm-runner/src/main/kotlin/org/usvm/runner/Main.kt b/usvm-runner/src/main/kotlin/org/usvm/runner/Main.kt
@@ -0,0 +1,63 @@
+package org.usvm.runner
+
+import ch.qos.logback.classic.Level
+import kotlinx.cli.ArgParser
+import kotlinx.cli.ArgType
+import kotlinx.cli.default
+import kotlinx.cli.required
+import org.slf4j.LoggerFactory
+import org.usvm.statistics.BasicBlock
+import org.usvm.utils.Mode
+import org.usvm.utils.OnnxModelImpl
+
+
+fun main(args: Array<String>) {
+    val parser = ArgParser("usvm.Runner")
+    val model by parser.option(
+        ArgType.String,
+        fullName = "model",
+        shortName = "m",
+        description = "Path to ONNX model."
+    ).required()
+
+    val path by parser.option(
+        ArgType.String,
+        fullName = "classpath",
+        shortName = "cp",
+        description = "Path to Java classpath."
+    ).required()
+
+    val mode by parser.option(
+        ArgType.Choice<Mode>(),
+        fullName = "mode",
+        description = "Mode to run inference in. Could be either CPU or GPU, the default is former."
+    ).default(Mode.CPU)
+
+    val methodFullName by parser.option(
+        ArgType.String,
+        fullName = "method",
+        description = "Full name of method to cover including package and class name."
+    ).required()
+
+    parser.parse(args)
+
+
+    val runner = JavaMethodRunner(
+        defaultOptions.copy(oracle = OnnxModelImpl<BasicBlock>(model, mode)),
+        classpath = path
+    )
+
+    // setLogLevel("ROOT", Level.DEBUG)
+
+
+    runner.cover(methodFullName)
+}
+
+// evil logger hack
+@Suppress("unused")
+private fun setLogLevel(loggerName: String, level: Level) {
+    val loggerContext = LoggerFactory.getILoggerFactory() as ch.qos.logback.classic.LoggerContext
+    val logger = loggerContext.getLogger(loggerName)
+    logger.level = level
+}
+
diff --git a/usvm-runner/src/main/kotlin/org/usvm/runner/util/DotGraph.kt b/usvm-runner/src/main/kotlin/org/usvm/runner/util/DotGraph.kt
@@ -1,3 +1,5 @@
+package org.usvm.runner.util
+
 import io.github.rchowell.dotlin.DotNodeShape
 import io.github.rchowell.dotlin.digraph
 import org.usvm.statistics.BasicBlock
diff --git a/usvm-runner/src/main/kotlin/org/usvm/runner/util/JacoDBContainer.kt b/usvm-runner/src/main/kotlin/org/usvm/runner/util/JacoDBContainer.kt
@@ -1,4 +1,4 @@
-package org.usvm.util
+package org.usvm.runner.util
 
 import kotlinx.coroutines.runBlocking
 import org.jacodb.api.jvm.JcClasspath
diff --git a/usvm-runner/src/main/kotlin/org/usvm/runner/util/Util.kt b/usvm-runner/src/main/kotlin/org/usvm/runner/util/Util.kt
@@ -1,4 +1,4 @@
-package org.usvm.util
+package org.usvm.runner.util
 
 import org.jacodb.api.jvm.JcClasspath
 import org.jacodb.api.jvm.JcTypedMethod
diff --git a/usvm-runner/src/main/resources/logback.xml b/usvm-runner/src/main/resources/logback.xml
@@ -0,0 +1,11 @@
+<configuration>
+    <appender name="STDOUT" class="ch.qos.logback.core.ConsoleAppender">
+        <encoder>
+            <pattern>%d{YYYY-MM-dd HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n</pattern>
+        </encoder>
+    </appender>
+    <root level="INFO">
+        <appender-ref ref="STDOUT" />
+    </root>
+    <logger name="io.netty" level="INFO"/>
+</configuration>

Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,7 @@ dependencies {`
`13`	`13`	`testImplementation(Libs.ksmt_yices)`
`14`	`14`
`15`	`15`	`implementation(Libs.onnxruntime)`
	`16`	`+ implementation(Libs.onnxruntime_gpu)`
`16`	`17`	`}`
`17`	`18`
`18`	`19`	`publishing {`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-package org.usvm.util`
	`1`	`+package org.usvm.gameserver.serialization`
`2`	`2`
`3`	`3`	`import org.usvm.gameserver.GameEdgeLabel`
`4`	`4`	`import org.usvm.gameserver.GameMapEdge`