Support targeting the Coral (#147)

* Add test to load from exported from tf 1.15

* Add PostTrainingQuantizationTask

* Add RunEdgeTpuCompilerTask

* Targeting MobileNetV2 for the Coral works

* Bump azure because of docker push

* Rename ModelDeploymentTarget.Normal

* Bump azure because of docker push

* Move failing test from Coral to Desktop

* Add disabled test that reproduces the conv2d issue targeting the coral

* Fix running training scripts

* Add a form to select the job target

* JobRunnerIntegTest working with new result uploading

* Cleanup

* Documentation and final cleanup

* Work on fixing windows build

* Work on fixing linux_docker build

* Working on fixing windows builds

* Don't work with paths manually

Author: Octogonapus

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

@@ -3,6 +3,7 @@ package edu.wpi.axon.aws
 import edu.wpi.axon.db.data.TrainingScriptProgress
 import edu.wpi.axon.tfdata.Dataset
 import edu.wpi.axon.util.FilePath
+import java.io.File
 import mu.KotlinLogging
 import org.apache.commons.lang3.RandomStringUtils
 import org.koin.core.KoinComponent
@@ -35,9 +36,6 @@ class EC2TrainingScriptRunner(
         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."
         }
@@ -80,14 +78,14 @@ class EC2TrainingScriptRunner(
             |apt-cache policy docker-ce
             |apt install -y docker-ce
             |systemctl status docker
-            |pip3 install https://github.com/wpilibsuite/axon-cli/releases/download/v0.1.12/axon-0.1.12-py2.py3-none-any.whl
+            |pip3 install https://github.com/wpilibsuite/axon-cli/releases/download/v0.1.14/axon-0.1.14-py2.py3-none-any.whl
             |axon create-heartbeat ${config.id}
             |axon update-training-progress ${config.id} "initializing"
             |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"
-            |axon upload-trained-model "${config.newModelName.filename}"
+            |docker run -v ${'$'}(eval "pwd"):/home wpilib/axon-ci:latest "/usr/bin/python3.6" "/home/$scriptFileName"
+            |axon upload-training-results ${config.id} "${config.workingDir}"
             |axon update-training-progress ${config.id} "completed"
             |axon remove-heartbeat ${config.id}
             |shutdown -h now
@@ -108,6 +106,11 @@ class EC2TrainingScriptRunner(
         canceller.addJob(config.id, instanceId)
     }
 
+    override fun listResults(id: Int): List<String> = s3Manager.listTrainingResults(id)
+
+    override fun getResult(id: Int, filename: String): File =
+        s3Manager.downloadTrainingResult(id, filename)
+
     fun getInstanceId(jobId: Int): String {
         requireJobIsInMaps(jobId)
         return instanceIds[jobId]!!

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

@@ -4,12 +4,12 @@ import edu.wpi.axon.db.data.TrainingScriptProgress
 import edu.wpi.axon.tfdata.Dataset
 import edu.wpi.axon.util.FilePath
 import edu.wpi.axon.util.createLocalProgressFilepath
+import edu.wpi.axon.util.getOutputModelName
 import edu.wpi.axon.util.runCommand
 import java.io.File
-import java.nio.file.Files
-import java.nio.file.Paths
 import kotlin.concurrent.thread
 import mu.KotlinLogging
+import org.apache.commons.lang3.RandomStringUtils
 
 /**
  * Runs the training script on the local machine. Assumes that Axon is running in the
@@ -31,9 +31,6 @@ class LocalTrainingScriptRunner(
         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."
         }
@@ -43,7 +40,8 @@ class LocalTrainingScriptRunner(
             }
         }
 
-        val scriptFile = Files.createTempFile("", ".py").toFile().apply {
+        val scriptFilename = "${config.workingDir}/${RandomStringUtils.randomAlphanumeric(20)}.py"
+        val scriptFile = File(scriptFilename).apply {
             createNewFile()
             writeText(config.scriptContents)
         }
@@ -83,8 +81,9 @@ class LocalTrainingScriptRunner(
                         """.trimMargin()
                     }
 
-                    val newModelFile =
-                        Paths.get(config.newModelName.path).toFile()
+                    val newModelFile = config.workingDir
+                        .resolve(getOutputModelName(config.oldModelName.filename))
+                        .toFile()
                     if (newModelFile.exists()) {
                         scriptProgressMap[config.id] = TrainingScriptProgress.Completed
                     } else {
@@ -100,6 +99,16 @@ class LocalTrainingScriptRunner(
         }
     }
 
+    override fun listResults(id: Int): List<String> {
+        requireJobIsInMaps(id)
+        return scriptDataMap[id]!!.workingDir.toFile().listFiles()!!.map { it.name }
+    }
+
+    override fun getResult(id: Int, filename: String): File {
+        requireJobIsInMaps(id)
+        return scriptDataMap[id]!!.workingDir.resolve(filename).toFile()
+    }
+
     override fun getTrainingProgress(jobId: Int) = progressReporter.getTrainingProgress(jobId)
 
     override fun overrideTrainingProgress(jobId: Int, progress: TrainingScriptProgress) =

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

@@ -2,22 +2,22 @@ package edu.wpi.axon.aws
 
 import edu.wpi.axon.tfdata.Dataset
 import edu.wpi.axon.util.FilePath
+import java.nio.file.Path
 
 /**
  * The configuration data needed to run a training script.
  *
  * @param oldModelName The path to the current model (that will be loaded).
- * @param newModelName The path to the new model (that will be saved).
  * @param dataset The path to the dataset.
  * @param scriptContents The contents of the training script.
  * @param epochs The number of epochs the model will be trained for. Must be greater than zero.
  * @param id The id of the Job this script is associated with.
  */
 data class RunTrainingScriptConfiguration(
     val oldModelName: FilePath,
-    val newModelName: FilePath,
     val dataset: Dataset,
     val scriptContents: String,
     val epochs: Int,
+    val workingDir: Path,
     val id: Int
 )

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

@@ -38,22 +38,23 @@ class S3Manager(
         downloadToLocalFile("axon-untrained-models/$filename")
 
     /**
-     * Uploads a trained model (one that the user wants to test with).
+     * Lists the training results for the Job.
      *
-     * @param file The local file containing the model to upload. The filename of the uploaded model
-     * will be the same as the filename of this file.
+     * @param jobId The ID of the Job.
+     * @return The filenames of the results.
      */
-    fun uploadTrainedModel(file: File) = uploadLocalFile(file, "axon-trained-models/${file.name}")
+    fun listTrainingResults(jobId: Int): List<String> =
+        listObjectsWithPrefixAndRemovePrefix("axon-training-results/$jobId/")
 
     /**
-     * Downloads a trained model (which was put in S3 by the training script when it
-     * ran on EC2). Meant to be used to download to the user's local machine.
+     * Downloads a training result to a local file.
      *
-     * @param filename The filename of the trained model file.
-     * @return A local file containing the trained model.
+     * @param jobId The ID of the Job.
+     * @param resultFilename The filename of the result to download.
+     * @return A local file containing the result.
      */
-    fun downloadTrainedModel(filename: String): File =
-        downloadToLocalFile("axon-trained-models/$filename")
+    fun downloadTrainingResult(jobId: Int, resultFilename: String): File =
+        downloadToLocalFile("axon-training-results/$jobId/$resultFilename")
 
     /**
      * Uploads a test data file.

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

@@ -1,5 +1,7 @@
 package edu.wpi.axon.aws
 
+import java.io.File
+
 interface TrainingScriptRunner : TrainingScriptProgressReporter, TrainingScriptCanceller {
 
     /**
@@ -8,4 +10,21 @@ interface TrainingScriptRunner : TrainingScriptProgressReporter, TrainingScriptC
      * @param config The data needed to start the script.
      */
     fun startScript(config: RunTrainingScriptConfiguration)
+
+    /**
+     * Lists the results from running the training script.
+     *
+     * @param id The Job ID.
+     * @return The names of the results.
+     */
+    fun listResults(id: Int): List<String>
+
+    /**
+     * Gets a result as a local file.
+     *
+     * @param id The Job ID.
+     * @param filename The name of the result to download.
+     * @return A local file containing the result.
+     */
+    fun getResult(id: Int, filename: String): File
 }

aws/src/main/kotlin/edu/wpi/axon/aws/preferences/Preferences.kt

@@ -13,7 +13,7 @@ import software.amazon.awssdk.services.ec2.model.InstanceType
  */
 @Serializable
 data class Preferences(
-    var defaultEC2NodeType: InstanceType = InstanceType.T2_MICRO,
+    var defaultEC2NodeType: InstanceType = InstanceType.T2_SMALL,
     var statusPollingDelay: Long = 5000
 ) {
 

aws/src/test/kotlin/edu/wpi/axon/aws/EC2TrainingScriptRunnerTest.kt

@@ -9,8 +9,10 @@ import io.kotlintest.shouldThrow
 import io.mockk.every
 import io.mockk.mockk
 import io.mockk.verify
+import java.io.File
 import org.apache.commons.lang3.RandomStringUtils
 import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.io.TempDir
 import software.amazon.awssdk.services.ec2.model.InstanceStateName
 import software.amazon.awssdk.services.ec2.model.InstanceType
 
@@ -23,8 +25,8 @@ internal class EC2TrainingScriptRunnerTest {
     )
 
     @Test
-    fun `test starting script and getting the state until completion`() {
-        val config = randomRunTrainingScriptConfigurationUsingAWS()
+    fun `test starting script and getting the state until completion`(@TempDir tempDir: File) {
+        val config = randomRunTrainingScriptConfigurationUsingAWS(tempDir)
         val instanceId = RandomStringUtils.randomAlphanumeric(10)
         val instanceType = InstanceType.T2_MICRO
 
@@ -89,8 +91,8 @@ internal class EC2TrainingScriptRunnerTest {
     }
 
     @Test
-    fun `test starting the script and cancelling it`() {
-        val config = randomRunTrainingScriptConfigurationUsingAWS()
+    fun `test starting the script and cancelling it`(@TempDir tempDir: File) {
+        val config = randomRunTrainingScriptConfigurationUsingAWS(tempDir)
         val instanceId = RandomStringUtils.randomAlphanumeric(10)
         val instanceType = InstanceType.T2_MICRO
 
@@ -150,63 +152,47 @@ internal class EC2TrainingScriptRunnerTest {
     }
 
     @Test
-    fun `test running with local old model`() {
+    fun `test running with local old model`(@TempDir tempDir: File) {
         shouldThrow<IllegalArgumentException> {
             runner.startScript(
                 RunTrainingScriptConfiguration(
                     FilePath.Local("a"),
-                    FilePath.S3("b"),
                     Dataset.ExampleDataset.FashionMnist,
                     "",
                     1,
+                    tempDir.toPath(),
                     1
                 )
             )
         }
     }
 
     @Test
-    fun `test running with local new model`() {
+    fun `test running with zero epochs`(@TempDir tempDir: File) {
         shouldThrow<IllegalArgumentException> {
             runner.startScript(
                 RunTrainingScriptConfiguration(
                     FilePath.S3("a"),
-                    FilePath.Local("b"),
-                    Dataset.ExampleDataset.FashionMnist,
-                    "",
-                    1,
-                    1
-                )
-            )
-        }
-    }
-
-    @Test
-    fun `test running with zero epochs`() {
-        shouldThrow<IllegalArgumentException> {
-            runner.startScript(
-                RunTrainingScriptConfiguration(
-                    FilePath.S3("a"),
-                    FilePath.S3("b"),
                     Dataset.ExampleDataset.FashionMnist,
                     "",
                     0,
+                    tempDir.toPath(),
                     1
                 )
             )
         }
     }
 
     @Test
-    fun `test running with local dataset`() {
+    fun `test running with local dataset`(@TempDir tempDir: File) {
         shouldThrow<IllegalArgumentException> {
             runner.startScript(
                 RunTrainingScriptConfiguration(
                     FilePath.S3("a"),
-                    FilePath.S3("b"),
                     Dataset.Custom(FilePath.Local("d"), "d"),
                     "",
                     1,
+                    tempDir.toPath(),
                     1
                 )
             )

aws/src/test/kotlin/edu/wpi/axon/aws/LocalTrainingScriptProgressReporterTest.kt

@@ -27,10 +27,10 @@ internal class LocalTrainingScriptProgressReporterTest {
         reporter.addJobAfterRestart(
             RunTrainingScriptConfiguration(
                 FilePath.Local("old.h5"),
-                FilePath.Local("new.h5"),
                 Dataset.ExampleDataset.Mnist,
                 "",
                 10,
+                tempDir.toPath(),
                 id
             )
         )
@@ -53,10 +53,10 @@ internal class LocalTrainingScriptProgressReporterTest {
         reporter.addJobAfterRestart(
             RunTrainingScriptConfiguration(
                 FilePath.Local("old.h5"),
-                FilePath.Local("new.h5"),
                 Dataset.ExampleDataset.Mnist,
                 "",
                 10,
+                tempDir.toPath(),
                 id
             )
         )
@@ -79,10 +79,10 @@ internal class LocalTrainingScriptProgressReporterTest {
         reporter.addJobAfterRestart(
             RunTrainingScriptConfiguration(
                 FilePath.Local("old.h5"),
-                FilePath.Local("new.h5"),
                 Dataset.ExampleDataset.Mnist,
                 "",
                 10,
+                tempDir.toPath(),
                 id
             )
         )