Initial skeleton code for Agent service

api/thrift/agent.thrift

@@ -75,7 +75,7 @@ union JobBase {
 }
 
 struct Job {
-    1: optional string jobId
+    1: optional JobInfo jobInfo
     2: optional JobBase jobUnion
     3: optional i32 statusReportInterval
     4: optional i32 maxRetries
@@ -93,11 +93,12 @@ struct JobListResponse {
 }
 
 enum JobStatusType {
-    PENDING = 0,
-    RUNNING = 1,
-    SUCCEEDED = 2,
-    FAILED = 3,
-    STOPPED = 4
+    UNKNOWN = 0,
+    PENDING = 1,
+    RUNNING = 2,
+    SUCCEEDED = 3,
+    FAILED = 4,
+    STOPPED = 5
 }
 
 struct ResourceUsage {

orchestration/BUILD.bazel

@@ -112,5 +112,6 @@ scala_test_suite(
 jvm_binary(
     name = "orchestration_assembly",
     main_class = "ai.chronon.service.ChrononServiceLauncher",
+    resources = glob(["src/main/resources/**/*"]),
     runtime_deps = [":lib"],
 )

orchestration/src/main/scala/ai/chronon/orchestration/agent/AgentConfig.scala

@@ -0,0 +1,18 @@
+package ai.chronon.orchestration.agent
+
+/** Configuration for Agent components.
+  *
+  * TODO: To move these to a separate file which can be configured per client
+  */
+object AgentConfig {
+
+  val orchestrationServiceHostname: String = "localhost"
+
+  val orchestrationServicePort: Int = 3903
+
+  val pollingIntervalMs: Long = 5 * 60 * 1000
+
+  val statusReportingIntervalMs: Long = 5 * 60 * 1000
+
+  val topicId: String = "test-topic"
+}

orchestration/src/main/scala/ai/chronon/orchestration/agent/JobExecutionService.scala

@@ -0,0 +1,31 @@
+package ai.chronon.orchestration.agent
+
+import ai.chronon.api.{Job, JobStatusType}
+
+/** TODO: To extend this further if needed and have implementations for both GCP and AWS.
+  * Just a place holder will change a lot in coming PRs
+  */
+trait JobExecutionService {
+
+  def submitJob(job: Job): Unit
+
+  def getJobStatus(jobId: String): JobStatusType
+}
+
+/** In-memory implementation of JobExecutionService for testing and development.
+  */
+class InMemoryJobExecutionService extends JobExecutionService {
+  private val jobStatuses = scala.collection.mutable.Map[String, JobStatusType]()
+
+  override def submitJob(job: Job): Unit = {
+    jobStatuses.put(job.jobInfo.getJobId, JobStatusType.RUNNING)
+  }
+
+  override def getJobStatus(jobId: String): JobStatusType = {
+    jobStatuses.getOrElse(jobId, JobStatusType.UNKNOWN)
+  }
+}
+
+object JobExecutionService {
+  def createInMemory(): JobExecutionService = new InMemoryJobExecutionService()
+}

orchestration/src/main/scala/ai/chronon/orchestration/agent/KVStore.scala

@@ -0,0 +1,73 @@
+package ai.chronon.orchestration.agent
+
+import ai.chronon.api.{Job, JobStatusType}
+
+/** Trait representing a key-value store service.
+  *
+  * This service provides abstraction for key-value storage operations,
+  * allowing the application to work with different storage implementations.
+  *
+  * TODO: To use our chronon KVStore, extend if needed and have implementations for both GCP and AWS
+  * Just a place holder will change a lot in coming PRs
+  */
+trait KVStore {
+
+  /** Stores a job with the given ID.
+    *
+    * @param jobId The unique identifier for the job
+    * @param job The job object to store
+    */
+  def storeJob(jobId: String, job: Job): Unit
+
+  /** Retrieves a job with the given ID.
+    *
+    * @param jobId The unique identifier for the job
+    * @return Option containing the job if found, None otherwise
+    */
+  def getJob(jobId: String): Option[Job]
+
+  /** Retrieves all jobs stored in the system.
+    *
+    * @return List of all jobs
+    */
+  def getAllJobs: List[Job]
+
+  /** Updates the status of a job with the given ID.
+    *
+    * @param jobId The unique identifier for the job
+    * @param status The new status of the job
+    */
+  def updateJobStatus(jobId: String, status: JobStatusType): Unit
+}
+
+/** In-memory implementation of KVStore for testing and development.
+  */
+class InMemoryKVStore extends KVStore {
+  private val store = scala.collection.mutable.Map[String, Job]()
+
+  override def storeJob(jobId: String, job: Job): Unit = {
+    store.put(jobId, job)
+  }
+
+  override def getJob(jobId: String): Option[Job] = {
+    store.get(jobId)
+  }
+
+  override def getAllJobs: List[Job] = {
+    store.values.toList
+  }
+
+  override def updateJobStatus(jobId: String, status: JobStatusType): Unit = {
+    store.get(jobId).foreach { job =>
+      job.jobInfo.setCurrentStatus(status)
+    }
+  }
+}
+
+object KVStore {
+
+  /** Factory method to create an in-memory KVStore instance.
+    * @return An in-memory KVStore implementation
+    */
+  def createInMemory(): KVStore = new InMemoryKVStore()
+}

orchestration/src/main/scala/ai/chronon/orchestration/agent/handlers/JobPollingHandler.scala

@@ -0,0 +1,109 @@
+package ai.chronon.orchestration.agent.handlers
+
+import ai.chronon.api.{Job, JobInfo}
+import ai.chronon.orchestration.agent.{AgentConfig, JobExecutionService, KVStore}
+import io.vertx.core.{AsyncResult, Handler}
+import io.vertx.core.buffer.Buffer
+import io.vertx.core.json.{JsonArray, JsonObject}
+import io.vertx.ext.web.client.{HttpResponse, WebClient}
+import org.slf4j.{Logger, LoggerFactory}
+
+import scala.util.{Failure, Success, Try}
+
+/** Handler for job polling operations. */
+class JobPollingHandler(
+    webClient: WebClient,
+    kvStore: KVStore,
+    jobExecutionService: JobExecutionService
+) extends Handler[java.lang.Long] {
+  private val logger: Logger = LoggerFactory.getLogger(classOf[JobPollingHandler])
+
+  /** Vert.x Handler implementation that gets called periodically.
+    *
+    * @param timerId The ID of the timer that triggered this handler
+    */
+  override def handle(timerId: java.lang.Long): Unit = {
+    pollForJobs()
+  }
+
+  /** Polls the orchestration service for new jobs. */
+  private def pollForJobs(): Unit = {
+    logger.info(s"Polling for jobs with topicId: ${AgentConfig.topicId}")
+
+    try {
+      val host = AgentConfig.orchestrationServiceHostname
+      val port = AgentConfig.orchestrationServicePort
+
+      logger.info(s"Connecting to host:$host at port:$port")
+
+      // Make HTTP request to orchestration service
+      webClient
+        .get(port, host, "/jobs/list")
+        .addQueryParam("topic_id", AgentConfig.topicId)
+        .send(createJobPollingResponseHandler())
+    } catch {
+      case e: Exception =>
+        logger.error("Error while polling for jobs", e)
+    }
+  }
+
+  /** Creates a response handler for job polling HTTP requests */
+  private def createJobPollingResponseHandler(): Handler[AsyncResult[HttpResponse[Buffer]]] = {
+    (ar: AsyncResult[HttpResponse[Buffer]]) =>
+      {
+        if (ar.succeeded()) {
+          val response = ar.result()
+          if (response.statusCode() == 200) {
+            val responseBody = response.bodyAsJsonObject()
+            processJobsResponse(responseBody)
+          } else {
+            logger.error(s"Failed to poll for jobs: ${response.statusCode()} ${response.statusMessage()}")
+          }
+        } else {
+          logger.error("Failed to poll for jobs", ar.cause())
+        }
+      }
+  }
+
+  /** Processes the response from the orchestration service containing jobs. */
+  private def processJobsResponse(responseJson: JsonObject): Unit = {
+    Try {
+      // Extract jobs to start
+      val jobsToStartJson = responseJson.getJsonArray("jobsToStart", new JsonArray())
+      val jobsToStart = (0 until jobsToStartJson.size()).map { i =>
+        val jobJson = jobsToStartJson.getJsonObject(i)
+
+        // Create JobInfo
+        val jobInfo = new JobInfo()
+        jobInfo.setJobId(jobJson.getString("jobId"))
+        // Set other job info properties as needed
+
+        // Create Job with JobInfo
+        val job = new Job()
+        job.setJobInfo(jobInfo)
+        // Set other job properties as needed
+
+        job
+      }
+
+      logger.info(s"Received ${jobsToStart.size} jobs to start")
+
+      // Process each job
+      jobsToStart.foreach { job =>
+        // TODO: To handle all other cases
+        if (kvStore.getJob(job.jobInfo.getJobId).isEmpty) {
+          // Store job in KV store
+          kvStore.storeJob(job.jobInfo.getJobId, job)
+
+          // Submit job to cluster
+          jobExecutionService.submitJob(job)
+        }
+      }
+    } match {
+      case Success(_) =>
+        logger.info("Successfully processed jobs response")
+      case Failure(e) =>
+        logger.error("Error processing jobs response", e)
+    }
+  }
+}

orchestration/src/main/scala/ai/chronon/orchestration/agent/handlers/StatusReportingHandler.scala

@@ -0,0 +1,95 @@
+package ai.chronon.orchestration.agent.handlers
+
+import ai.chronon.api.JobStatusType
+import ai.chronon.orchestration.agent.{AgentConfig, JobExecutionService, KVStore}
+import io.vertx.core.{AsyncResult, Handler}
+import io.vertx.core.buffer.Buffer
+import io.vertx.core.json.JsonObject
+import io.vertx.ext.web.client.{HttpResponse, WebClient}
+import org.slf4j.{Logger, LoggerFactory}
+
+import scala.util.{Failure, Success, Try}
+
+/** Handler for status reporting operations. */
+class StatusReportingHandler(
+    webClient: WebClient,
+    kvStore: KVStore,
+    jobExecutionService: JobExecutionService
+) extends Handler[java.lang.Long] {
+  private val logger: Logger = LoggerFactory.getLogger(classOf[StatusReportingHandler])
+
+  /** Vert.x Handler implementation that gets called periodically.
+    *
+    * @param timerId The ID of the timer that triggered this handler
+    */
+  override def handle(timerId: java.lang.Long): Unit = {
+    reportJobStatuses()
+  }
+
+  /** Checks the status of all jobs and reports changes. */
+  private def reportJobStatuses(): Unit = {
+    logger.debug("Checking for job status updates")
+
+    try {
+      // Get all jobs from KV store
+      val jobs = kvStore.getAllJobs
+
+      // For each job, check current status and report if changed
+      jobs.foreach { job =>
+        val jobId = job.jobInfo.getJobId
+        val currentStatus = job.jobInfo.getCurrentStatus
+        val latestStatus = jobExecutionService.getJobStatus(jobId)
+
+        if (currentStatus != latestStatus) {
+          logger.info(s"Job $jobId status changed from $currentStatus to $latestStatus")
+
+          // Update status in KV store
+          kvStore.updateJobStatus(jobId, latestStatus)
+
+          // Report status change to orchestration service
+          reportStatusToOrchestration(jobId, latestStatus)
+        }
+      }
+    } catch {
+      case e: Exception =>
+        logger.error("Error while reporting job statuses", e)
+    }
+  }
+
+  /** Reports a job status update to the orchestration service. */
+  private def reportStatusToOrchestration(jobId: String, status: JobStatusType): Unit = {
+    // Create status update payload
+    val statusUpdate = new JsonObject()
+      .put("jobId", jobId)
+      .put("status", status.toString)
+      .put("timestamp", System.currentTimeMillis())
+
+    val host = AgentConfig.orchestrationServiceHostname
+    val port = AgentConfig.orchestrationServicePort
+
+    logger.info(s"Connecting to host:$host at port:$port")
+
+    // Send status update to orchestration service
+    webClient
+      .put(port, host, s"/jobs/$jobId/status")
+      .sendJsonObject(statusUpdate, createStatusUpdateHandler(jobId, status))
+  }
+
+  /** Creates a response handler for status update HTTP requests */
+  private def createStatusUpdateHandler(
+      jobId: String,
+      status: JobStatusType): Handler[AsyncResult[HttpResponse[Buffer]]] = { (ar: AsyncResult[HttpResponse[Buffer]]) =>
+    {
+      if (ar.succeeded()) {
+        val response = ar.result()
+        if (response.statusCode() == 200) {
+          logger.info(s"Successfully reported status update for job $jobId: $status")
+        } else {
+          logger.error(s"Failed to report status update: ${response.statusCode()} ${response.statusMessage()}")
+        }
+      } else {
+        logger.error("Failed to report status update", ar.cause())
+      }
+    }
+  }
+}