Add LocalTrainingScriptRunner (#138)

* Remove anything related to interfacing with S3 in the training script

* Initial implementation of LocalTrainingScriptRunner

* Fix SaveModelTaskTest

* Add LocalProgressReportingCallbackTask

* Fix integration tests

* Add missing dependency

* Fix issues with injecting the bucket. Fix bug in LocalPreferencesManager. Add LocalPreferencesManagerTest.

* Add a job for running locally. Improve logging. Working on a test failure.

* Patch the scripts

* Fix Conv32321IntegrationTest

* Improve progress reporting for LocalTrainingScriptRunner

* Convert Dataset.Custom to use a FilePath instead of assuming it is in S3

* Update util/src/main/kotlin/edu/wpi/axon/util/Util.kt

* Address PR comments

* Simplify local progress file creation logic

* Fix running jobs with EC2. Improve status reporting. Eagerly load AWS bucket. Download untrained models from S3 if necessary.

* Add comment

Author: Octogonapus

aws/aws.gradle.kts

@@ -21,6 +21,7 @@ dependencies {
 
     api(project(":tf-data"))
     api(project(":db-data"))
+    api(project(":training"))
 
     api(
         group = "software.amazon.awssdk",

aws/src/main/kotlin/edu/wpi/axon/aws/EC2TrainingScriptRunner.kt

@@ -2,6 +2,7 @@ package edu.wpi.axon.aws
 
 import edu.wpi.axon.dbdata.TrainingScriptProgress
 import edu.wpi.axon.tfdata.Dataset
+import edu.wpi.axon.util.FilePath
 import java.lang.NumberFormatException
 import java.util.Base64
 import java.util.concurrent.atomic.AtomicLong
@@ -10,6 +11,7 @@ import org.apache.commons.lang3.RandomStringUtils
 import org.koin.core.KoinComponent
 import software.amazon.awssdk.services.ec2.Ec2Client
 import software.amazon.awssdk.services.ec2.model.Ec2Exception
+import software.amazon.awssdk.services.ec2.model.Filter
 import software.amazon.awssdk.services.ec2.model.InstanceStateName
 import software.amazon.awssdk.services.ec2.model.InstanceType
 import software.amazon.awssdk.services.ec2.model.ShutdownBehavior
@@ -20,11 +22,12 @@ import software.amazon.awssdk.services.ec2.model.ShutdownBehavior
  * this class will handle all of that. The script should just load and save the model from/to its
  * current directory.
  *
+ * @param bucketName The name of the S3 bucket to use.
  * @param instanceType The type of the EC2 instance to run the training script on.
  */
 class EC2TrainingScriptRunner(
     bucketName: String,
-    private val instanceType: InstanceType // TODO: Move this to [startScript]
+    private val instanceType: InstanceType
 ) : TrainingScriptRunner, KoinComponent {
 
     private val ec2 by lazy { Ec2Client.builder().build() }
@@ -35,30 +38,32 @@ class EC2TrainingScriptRunner(
     private val scriptDataMap = mutableMapOf<Long, RunTrainingScriptConfiguration>()
 
     override fun startScript(
-        runTrainingScriptConfiguration: RunTrainingScriptConfiguration
+        config: RunTrainingScriptConfiguration
     ): Long {
-        // Check for if the script uses the CLI to manage the model in S3. This class is supposed to
-        // own working with S3.
-        require(
-            !runTrainingScriptConfiguration.scriptContents.contains("download_model") &&
-                !runTrainingScriptConfiguration.scriptContents.contains("upload_model")
-        ) {
-            """
-            |Cannot start the script because it interfaces with AWS:
-            |${runTrainingScriptConfiguration.scriptContents}
-            |
-            """.trimMargin()
+        require(config.oldModelName is FilePath.S3) {
+            "Must start from a model in S3. Got: ${config.oldModelName}"
+        }
+        require(config.newModelName is FilePath.S3) {
+            "Must export to a model in S3. Got: ${config.newModelName}"
+        }
+        require(config.epochs > 0) {
+            "Must train for at least one epoch. Got ${config.epochs} epochs."
+        }
+        when (config.dataset) {
+            is Dataset.Custom -> require(config.dataset.path is FilePath.S3) {
+                "Custom datasets must be in S3. Got non-local dataset: ${config.dataset}"
+            }
         }
 
         // The file name for the generated script
         val scriptFileName = "${RandomStringUtils.randomAlphanumeric(20)}.py"
 
-        val newModelName = runTrainingScriptConfiguration.newModelName
-        val datasetName = runTrainingScriptConfiguration.dataset.nameForS3ProgressReporting
+        val newModelName = config.newModelName.filename
+        val datasetName = config.dataset.progressReportingName
 
         s3Manager.uploadTrainingScript(
             scriptFileName,
-            runTrainingScriptConfiguration.scriptContents
+            config.scriptContents
         )
 
         // Reset the training progress so the script doesn't start in the completed state
@@ -69,10 +74,10 @@ class EC2TrainingScriptRunner(
 
         // We need to download custom datasets from S3. Example datasets will be downloaded
         // by the script using Keras.
-        val downloadDatasetString = when (runTrainingScriptConfiguration.dataset) {
+        val downloadDatasetString = when (config.dataset) {
             is Dataset.ExampleDataset -> ""
             is Dataset.Custom ->
-                """axon download-dataset "${runTrainingScriptConfiguration.dataset.pathInS3}""""
+                """axon download-dataset "${config.dataset.path.path}""""
         }
 
         val scriptForEC2 = """
@@ -91,7 +96,7 @@ class EC2TrainingScriptRunner(
             |pip3 install https://github.com/wpilibsuite/axon-cli/releases/download/v0.1.11/axon-0.1.11-py2.py3-none-any.whl
             |axon create-heartbeat "$newModelName" "$datasetName"
             |axon update-training-progress "$newModelName" "$datasetName" "initializing"
-            |axon download-untrained-model "${runTrainingScriptConfiguration.oldModelName}"
+            |axon download-untrained-model "${config.oldModelName.path}"
             |$downloadDatasetString
             |axon download-training-script "$scriptFileName"
             |docker run -v ${'$'}(eval "pwd"):/home wpilib/axon-ci:latest "/usr/bin/python3.6 /home/$scriptFileName"
@@ -122,7 +127,7 @@ class EC2TrainingScriptRunner(
 
         val scriptId = nextScriptId.getAndIncrement()
         instanceIds[scriptId] = runInstancesResponse.instances().first().instanceId()
-        scriptDataMap[scriptId] = runTrainingScriptConfiguration
+        scriptDataMap[scriptId] = config
         return scriptId
     }
 
@@ -132,14 +137,24 @@ class EC2TrainingScriptRunner(
         require(scriptId in scriptDataMap.keys)
 
         val runTrainingScriptConfiguration = scriptDataMap[scriptId]!!
-        val newModelName = runTrainingScriptConfiguration.newModelName
-        val datasetName = runTrainingScriptConfiguration.dataset.nameForS3ProgressReporting
+        val newModelName = runTrainingScriptConfiguration.newModelName.filename
+        val datasetName = runTrainingScriptConfiguration.dataset.progressReportingName
 
         val status = try {
             ec2.describeInstanceStatus {
                 it.instanceIds(instanceIds[scriptId]!!)
+                    .includeAllInstances(true)
+                    .filters(
+                        Filter.builder().name("instance-state-name").values(
+                            "pending",
+                            "running",
+                            "shutting-down",
+                            "stopping"
+                        ).build()
+                    )
             }.instanceStatuses().firstOrNull()?.instanceState()?.name()
         } catch (ex: Ec2Exception) {
+            LOGGER.warn(ex) { "Failed to get instance status." }
             null
         }
 
@@ -166,13 +181,29 @@ class EC2TrainingScriptRunner(
             status: InstanceStateName?,
             epochs: Int
         ): TrainingScriptProgress {
+            LOGGER.debug {
+                """
+                |Heartbeat: $heartbeat
+                |Progress: $progress
+                |Instance status: $status
+                """.trimMargin()
+            }
+
             val progressAssumingEverythingIsFine = computeProgressAssumingEverythingIsFine(
                 heartbeat,
                 progress,
                 status,
                 epochs
             )
 
+            if ((status == InstanceStateName.SHUTTING_DOWN ||
+                    status == InstanceStateName.TERMINATED ||
+                    status == InstanceStateName.STOPPING) &&
+                (heartbeat != "0" || progress != "completed")
+            ) {
+                return TrainingScriptProgress.Error
+            }
+
             return when (heartbeat) {
                 "0" -> when (progress) {
                     "not started", "completed" -> progressAssumingEverythingIsFine
@@ -185,8 +216,7 @@ class EC2TrainingScriptRunner(
 
                     else -> when (status) {
                         InstanceStateName.SHUTTING_DOWN, InstanceStateName.TERMINATED,
-                        InstanceStateName.STOPPING, InstanceStateName.STOPPED, null ->
-                            TrainingScriptProgress.Error
+                        InstanceStateName.STOPPING -> TrainingScriptProgress.Error
 
                         else -> progressAssumingEverythingIsFine
                     }
@@ -205,7 +235,7 @@ class EC2TrainingScriptRunner(
             "not started" -> when (status) {
                 InstanceStateName.PENDING -> TrainingScriptProgress.Creating
                 InstanceStateName.RUNNING -> TrainingScriptProgress.Initializing
-                else -> TrainingScriptProgress.NotStarted
+                else -> TrainingScriptProgress.Creating
             }
 
             "initializing" -> TrainingScriptProgress.Initializing

aws/src/main/kotlin/edu/wpi/axon/aws/LoadPreferences.kt

@@ -1,29 +1,33 @@
 package edu.wpi.axon.aws
 
+import arrow.core.None
+import arrow.core.Option
+import arrow.core.Some
 import mu.KotlinLogging
 import software.amazon.awssdk.core.exception.SdkClientException
 import software.amazon.awssdk.services.s3.S3Client
 
 /**
- * Finds the S3 bucket Axon will work out of. Returns `null` if there is no matching bucket. The AWS
- * region MUST be auto-detectable from the environment (like when running on ECS). To use AWS when
- * running locally, set `AWS_REGION` to your preferred region. To not use AWS when running locally,
- * do not set `AWS_REGION`.
+ * Finds the S3 bucket Axon will work out of. Returns [None] if there is no matching bucket, which
+ * causes Axon to run locally and not interface with AWS. The AWS region MUST be auto-detectable
+ * from the environment (like when running on ECS). To use AWS when running locally, set
+ * `AWS_REGION` to your preferred region. To not use AWS when running locally, do not set
+ * `AWS_REGION`.
  *
- * @return The name of the bucket or `null` if the bucket could not be found.
+ * @return The name of the bucket or [None] if the bucket could not be found.
  */
-fun findAxonS3Bucket() = try {
+fun findAxonS3Bucket(): Option<String> = try {
     val s3Client = S3Client.builder().build()
 
     val bucket = s3Client.listBuckets().buckets().first {
         it.name().startsWith("axon-autogenerated-")
     }
 
     LOGGER.info { "Starting with S3 bucket: $bucket" }
-    bucket.name()
+    Some(bucket.name())
 } catch (e: SdkClientException) {
     LOGGER.info(e) { "Not loading credentials because of this exception." }
-    null
+    None
 }
 
 private val LOGGER = KotlinLogging.logger { }

aws/src/main/kotlin/edu/wpi/axon/aws/LocalTrainingScriptRunner.kt

@@ -0,0 +1,146 @@
+package edu.wpi.axon.aws
+
+import edu.wpi.axon.dbdata.TrainingScriptProgress
+import edu.wpi.axon.tfdata.Dataset
+import edu.wpi.axon.util.FilePath
+import edu.wpi.axon.util.createProgressFilePath
+import edu.wpi.axon.util.runCommand
+import java.io.File
+import java.lang.NumberFormatException
+import java.nio.file.Files
+import java.nio.file.Paths
+import java.util.concurrent.atomic.AtomicLong
+import kotlin.concurrent.thread
+import mu.KotlinLogging
+
+/**
+ * Runs the training script on the local machine. Assumes that Axon is running in the
+ * wpilib/axon-hosted Docker container.
+ */
+class LocalTrainingScriptRunner : TrainingScriptRunner {
+
+    private val nextScriptId = AtomicLong()
+    private val scriptDataMap = mutableMapOf<Long, RunTrainingScriptConfiguration>()
+    private val scriptProgressMap = mutableMapOf<Long, TrainingScriptProgress>()
+    private val scriptThreadMap = mutableMapOf<Long, Thread>()
+
+    override fun startScript(config: RunTrainingScriptConfiguration): Long {
+        require(config.oldModelName is FilePath.Local) {
+            "Must start from a local model. Got: ${config.oldModelName}"
+        }
+        require(config.newModelName is FilePath.Local) {
+            "Must export to a local model. Got: ${config.newModelName}"
+        }
+        require(config.epochs > 0) {
+            "Must train for at least one epoch. Got ${config.epochs} epochs."
+        }
+        when (config.dataset) {
+            is Dataset.Custom -> require(config.dataset.path is FilePath.Local) {
+                "Custom datasets must be local. Got non-local dataset: ${config.dataset}"
+            }
+        }
+
+        val scriptFile = Files.createTempFile("", ".py").toFile()
+        scriptFile.createNewFile()
+        scriptFile.writeText(config.scriptContents)
+
+        val modelName = config.newModelName.filename
+        val datasetName = config.dataset.progressReportingName
+
+        // Clear the progress file if there was a previous run
+        val progressFile = File(createProgressFilePath(modelName, datasetName))
+        progressFile.parentFile.mkdirs()
+        progressFile.createNewFile()
+        progressFile.writeText("0.0")
+
+        val scriptId = nextScriptId.getAndIncrement()
+        scriptDataMap[scriptId] = config
+        scriptProgressMap[scriptId] = TrainingScriptProgress.Creating
+
+        scriptThreadMap[scriptId] = thread {
+            scriptProgressMap[scriptId] = TrainingScriptProgress.Initializing
+
+            runCommand(
+                listOf("python3.6", scriptFile.absolutePath),
+                emptyMap(),
+                null
+            ).attempt().unsafeRunSync().fold(
+                {
+                    LOGGER.warn(it) { "Training script failed." }
+                    scriptProgressMap[scriptId] = TrainingScriptProgress.Error
+                },
+                { (exitCode, stdOut, stdErr) ->
+                    LOGGER.info {
+                        """
+                        |Training script completed.
+                        |Process exit code: $exitCode
+                        |Process std out:
+                        |$stdOut
+                        |
+                        |Process std err:
+                        |$stdErr
+                        |
+                        """.trimMargin()
+                    }
+
+                    val newModelFile =
+                        Paths.get(config.newModelName.path).toFile()
+                    if (newModelFile.exists()) {
+                        scriptProgressMap[scriptId] = TrainingScriptProgress.Completed
+                    } else {
+                        scriptProgressMap[scriptId] = TrainingScriptProgress.Error
+                    }
+                }
+            )
+        }
+
+        return scriptId
+    }
+
+    override fun getTrainingProgress(scriptId: Long): TrainingScriptProgress {
+        require(scriptId in scriptDataMap.keys)
+        require(scriptId in scriptThreadMap.keys)
+        require(scriptId in scriptProgressMap.keys)
+
+        return if (scriptThreadMap[scriptId]!!.isAlive) {
+            // Training thread is still running. Try to read the progress file.
+            when (scriptProgressMap[scriptId]!!) {
+                // These statuses are reasonable
+                is TrainingScriptProgress.Initializing -> TrainingScriptProgress.Initializing
+                is TrainingScriptProgress.Completed -> TrainingScriptProgress.Completed
+                is TrainingScriptProgress.Error -> TrainingScriptProgress.Error
+                else -> {
+                    // Otherwise it must be InProgress
+                    val config = scriptDataMap[scriptId]!!
+                    val modelName = config.newModelName.filename
+                    val datasetName = config.dataset.progressReportingName
+                    val progressFile = File(createProgressFilePath(modelName, datasetName))
+                    if (progressFile.exists()) {
+                        try {
+                            TrainingScriptProgress.InProgress(
+                                progressFile.readText().toDouble() / config.epochs
+                            )
+                        } catch (ex: NumberFormatException) {
+                            TrainingScriptProgress.Error
+                        }
+                    } else {
+                        TrainingScriptProgress.InProgress(0.0)
+                    }
+                }
+            }
+        } else {
+            // Training thread died or is not started yet. If it dies, either it finished and wrote
+            // Completed to scriptProgressMap or exploded and didn't write Completed. If it is not
+            // started yet, then the status will still be Creating.
+            when (scriptProgressMap[scriptId]!!) {
+                is TrainingScriptProgress.Creating -> TrainingScriptProgress.Creating
+                is TrainingScriptProgress.Completed -> TrainingScriptProgress.Completed
+                else -> TrainingScriptProgress.Error
+            }
+        }
+    }
+
+    companion object {
+        private val LOGGER = KotlinLogging.logger { }
+    }
+}