|
| 1 | +package org.thp.cortex.services |
| 2 | + |
| 3 | +import java.util.concurrent.TimeUnit |
| 4 | +import java.nio.file._ |
| 5 | + |
| 6 | +import scala.concurrent.duration.FiniteDuration |
| 7 | +import scala.util.{Try} |
| 8 | +import scala.collection.JavaConverters._ |
| 9 | + |
| 10 | +import play.api.{Configuration, Logger} |
| 11 | + |
| 12 | +import akka.actor.ActorSystem |
| 13 | +import io.fabric8.kubernetes.client.{DefaultKubernetesClient} |
| 14 | +import io.fabric8.kubernetes.api.model.batch.{JobBuilder => KJobBuilder} |
| 15 | +import io.fabric8.kubernetes.api.model.{PersistentVolumeClaimVolumeSourceBuilder} |
| 16 | +import javax.inject.{Inject, Singleton} |
| 17 | +import org.thp.cortex.models._ |
| 18 | + |
| 19 | +@Singleton |
| 20 | +class K8sJobRunnerSrv( |
| 21 | + client: DefaultKubernetesClient, |
| 22 | + config: Configuration, |
| 23 | + autoUpdate: Boolean, |
| 24 | + jobBaseDirectory: Path, |
| 25 | + persistentVolumeClaimName: String, |
| 26 | + implicit val system: ActorSystem |
| 27 | +) { |
| 28 | + |
| 29 | + @Inject() |
| 30 | + def this(config: Configuration, system: ActorSystem) = |
| 31 | + this( |
| 32 | + new DefaultKubernetesClient(), |
| 33 | + config, |
| 34 | + config.getOptional[Boolean]("job.kubernetes.autoUpdate").getOrElse(true), |
| 35 | + Paths.get(config.get[String]("job.directory")), |
| 36 | + config.get[String]("job.kubernetes.persistentVolumeClaimName"), |
| 37 | + system: ActorSystem |
| 38 | + ) |
| 39 | + |
| 40 | + lazy val logger = Logger(getClass) |
| 41 | + |
| 42 | + lazy val isAvailable: Boolean = |
| 43 | + Try { |
| 44 | + val ver = client.getVersion() |
| 45 | + logger.info(s"Kubernetes is available: major ${ver.getMajor()} minor ${ver.getMinor()} git ${ver.getGitCommit()}") |
| 46 | + true |
| 47 | + }.recover { |
| 48 | + case error => |
| 49 | + logger.info(s"Kubernetes is not available", error) |
| 50 | + false |
| 51 | + }.get |
| 52 | + |
| 53 | + def run(jobDirectory: Path, dockerImage: String, job: Job, timeout: Option[FiniteDuration]): Try[Unit] = { |
| 54 | + val cacertsFile = jobDirectory.resolve("input").resolve("cacerts") |
| 55 | + val relativeJobDirectory = jobBaseDirectory.relativize(jobDirectory).toString() |
| 56 | + // make the default longer than likely values, but still not infinite |
| 57 | + val timeout_or_default = timeout getOrElse new FiniteDuration(8, TimeUnit.HOURS) |
| 58 | + // https://kubernetes.io/docs/concepts/overview/working-with-objects/names/ |
| 59 | + // FIXME: this collapses case, jeopardizing the uniqueness of the |
| 60 | + // identifier. LDH: lowercase, digits, hyphens. |
| 61 | + val ldh_jobid = "_".r.replaceAllIn(job.id.map(_.toLower), "-") |
| 62 | + val kjobName = "neuron-job-" + ldh_jobid |
| 63 | + val pvcvs = new PersistentVolumeClaimVolumeSourceBuilder() |
| 64 | + .withClaimName(persistentVolumeClaimName) |
| 65 | + .withReadOnly(false) |
| 66 | + .build(); |
| 67 | + val kjob1 = new KJobBuilder() |
| 68 | + .withApiVersion("batch/v1") |
| 69 | + .withNewMetadata() |
| 70 | + .withName(kjobName) |
| 71 | + .withLabels(Map( |
| 72 | + "cortex-job-id" -> job.id, |
| 73 | + "cortex-worker-id" -> job.workerId(), |
| 74 | + "cortex-neuron-job" -> "true").asJava) |
| 75 | + .endMetadata() |
| 76 | + .withNewSpec() |
| 77 | + .withNewTemplate() |
| 78 | + .withNewSpec() |
| 79 | + .addNewVolume() |
| 80 | + .withName("job-directory") |
| 81 | + .withPersistentVolumeClaim(pvcvs) |
| 82 | + .endVolume() |
| 83 | + .addNewContainer() |
| 84 | + .withName("neuron") |
| 85 | + .withImage(dockerImage) |
| 86 | + .withArgs("/job") |
| 87 | + .addNewEnv() |
| 88 | + .withName("CORTEX_JOB_FOLDER") |
| 89 | + .withValue(relativeJobDirectory) |
| 90 | + .endEnv(); |
| 91 | + val kjob2 = if (Files.exists(cacertsFile)) { |
| 92 | + kjob1.addNewEnv() |
| 93 | + .withName("REQUESTS_CA_BUNDLE") |
| 94 | + .withValue("/job/input/cacerts") |
| 95 | + .endEnv() |
| 96 | + } else { |
| 97 | + kjob1 |
| 98 | + } |
| 99 | + val kjob3 = kjob2 |
| 100 | + .addNewVolumeMount() |
| 101 | + .withName("job-directory") |
| 102 | + .withSubPathExpr("$(CORTEX_JOB_FOLDER)/input") |
| 103 | + .withMountPath("/job/input") |
| 104 | + .withReadOnly(true) |
| 105 | + .endVolumeMount() |
| 106 | + .addNewVolumeMount() |
| 107 | + .withName("job-directory") |
| 108 | + .withSubPathExpr("$(CORTEX_JOB_FOLDER)/output") |
| 109 | + .withMountPath("/job/output") |
| 110 | + .withReadOnly(false) |
| 111 | + .endVolumeMount() |
| 112 | + .endContainer() |
| 113 | + .withRestartPolicy("Never") |
| 114 | + .endSpec() |
| 115 | + .endTemplate() |
| 116 | + .endSpec() |
| 117 | + .build(); |
| 118 | + |
| 119 | + val execution = Try { |
| 120 | + val created_kjob = client.batch().jobs().create(kjob3) |
| 121 | + val created_env = created_kjob |
| 122 | + .getSpec().getTemplate().getSpec().getContainers().get(0) |
| 123 | + .getEnv().asScala; |
| 124 | + logger.info( |
| 125 | + s"Created Kubernetes Job ${created_kjob.getMetadata().getName()}\n" + |
| 126 | + s" timeout: ${timeout_or_default.toString}\n" + |
| 127 | + s" image : $dockerImage\n" + |
| 128 | + s" mount : pvc ${persistentVolumeClaimName} subdir ${relativeJobDirectory} as /job" + |
| 129 | + created_env.map(ev => s"\n env : ${ev.getName()} = ${ev.getValue()}").mkString) |
| 130 | + val ended_kjob = client.batch().jobs().withLabel("cortex-job-id", job.id) |
| 131 | + .waitUntilCondition((x => Option(x).flatMap(j => |
| 132 | + Option(j.getStatus).flatMap(s => |
| 133 | + Some(s.getConditions.asScala.map(_.getType).filter(t => |
| 134 | + t.equals("Complete") || t.equals("Failed")).nonEmpty))) |
| 135 | + getOrElse false), |
| 136 | + timeout_or_default.length, timeout_or_default.unit); |
| 137 | + if(ended_kjob != null) { |
| 138 | + logger.info(s"Kubernetes Job ${ended_kjob.getMetadata().getName()} " + |
| 139 | + s"(for job ${job.id}) status is now ${ended_kjob.getStatus().toString()}") |
| 140 | + } else { |
| 141 | + logger.info(s"Kubernetes Job for ${job.id} no longer exists") |
| 142 | + } |
| 143 | + } |
| 144 | + // let's find the job by the attribute we know is fundamentally |
| 145 | + // unique, rather than one constructed from it |
| 146 | + val deleted = client.batch().jobs().withLabel("cortex-job-id", job.id).delete() |
| 147 | + if(deleted) { |
| 148 | + logger.info(s"Deleted Kubernetes Job for job ${job.id}") |
| 149 | + } else { |
| 150 | + logger.info(s"While trying to delete Kubernetes Job for ${job.id}, the job was not found; this is OK") |
| 151 | + } |
| 152 | + execution |
| 153 | + } |
| 154 | +} |
0 commit comments