feat: implement Dataproc spark submitter utility #122
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,6 @@ | ||
| # configurations for testing | ||
| projectId: "canary-443022" | ||
| region: "us-central1" | ||
| clusterName: "canary-2" | ||
| jarUri: "gs://dataproc-temp-us-central1-703996152583-pqtvfptb/jars/cloud_gcp-assembly-0.1.0-SNAPSHOT.jar" | ||
|
Comment on lines
+2
to
+5
Contributor
There was a problem hiding this comment. 💡 Codebase verification Sensitive information needs to be externalized The verification reveals hardcoded sensitive information in two locations:
Recommendations:
🔗 Analysis chainVerify sensitive information usage Let's ensure these sensitive values aren't hardcoded elsewhere in the codebase. 🏁 Scripts executedThe following scripts were executed for the analysis: Script: #!/bin/bash
# Search for hardcoded project ID and bucket path
echo "Checking for hardcoded project ID..."
rg -l "canary-443022"
echo "Checking for hardcoded bucket path..."
rg -l "703996152583"
echo "Checking for other potential Dataproc configurations..."
fd -e yaml -e yml -e json -x grep -l "projectId\|clusterName\|jarUri"
Length of output: 671 Script: #!/bin/bash
# Let's check the content of the test file to understand the context
echo "Content of DataprocSubmitterTest.scala:"
cat "cloud_gcp/src/test/scala/ai/chronon/integrations/cloud_gcp/test/DataprocSubmitterTest.scala"
# Let's also check if there are any configuration templates or example files
echo -e "\nChecking for template or example configuration files:"
fd -e yaml.template -e yaml.example -e yml.template -e yml.example
Length of output: 2580 |
||
| mainClass: "ai.chronon.spark.Driver" | ||
tchow-zlai marked this conversation as resolved.
Show resolved
Hide resolved
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,105 @@ | ||
| package ai.chronon.integrations.cloud_gcp | ||
| import ai.chronon.spark.SparkAuth | ||
| import ai.chronon.spark.SparkSubmitter | ||
| import com.google.api.gax.rpc.ApiException | ||
| import com.google.cloud.dataproc.v1._ | ||
| import io.circe.generic.auto._ | ||
| import io.circe.yaml.parser | ||
|
|
||
| import scala.io.Source | ||
|
|
||
| import collection.JavaConverters._ | ||
|
|
||
| case class SubmitterConf( | ||
| projectId: String, | ||
| region: String, | ||
| clusterName: String, | ||
| jarUri: String, | ||
| mainClass: String | ||
| ) { | ||
|
|
||
| def endPoint: String = s"${region}-dataproc.googleapis.com:443" | ||
| } | ||
|
|
||
| case class GeneralJob( | ||
| jobName: String, | ||
| jars: String, | ||
| mainClass: String | ||
| ) | ||
nikhil-zlai marked this conversation as resolved.
Show resolved
Hide resolved
Comment on lines
+24
to
+28
Contributor
There was a problem hiding this comment. 🛠️ Refactor suggestion Remove unused The |
||
|
|
||
| class DataprocSubmitter(jobControllerClient: JobControllerClient, conf: SubmitterConf) extends SparkSubmitter { | ||
|
|
||
| override def status(jobId: String): Unit = { | ||
| try { | ||
| val currentJob: Job = jobControllerClient.getJob(conf.projectId, conf.region, jobId) | ||
| currentJob.getStatus.getState | ||
| } catch { | ||
| case e: ApiException => | ||
| println(s"Error monitoring job: ${e.getMessage}") | ||
|
Contributor
There was a problem hiding this comment. 🛠️ Refactor suggestion Use a logging framework instead of Using Import a logging framework and update error handling:
Also applies to: 71-71 |
||
| } | ||
| } | ||
|
Comment on lines
+32
to
+40
Contributor
There was a problem hiding this comment. Enhance job status monitoring implementation The status method:
Apply this diff: - override def status(jobId: String): Unit = {
+ override def status(jobId: String): JobState = {
try {
val currentJob: Job = jobControllerClient.getJob(conf.projectId, conf.region, jobId)
- currentJob.getStatus.getState
+ currentJob.getStatus.getState match {
+ case null => throw new IllegalStateException(s"No state found for job $jobId")
+ case state => JobState(state.toString)
+ }
} catch {
case e: ApiException =>
- println(s"Error monitoring job: ${e.getMessage}")
+ logger.error(s"Error monitoring job $jobId", e)
+ throw new JobMonitoringException(s"Failed to monitor job $jobId", e)
}
}
|
||
|
|
||
| override def kill(jobId: String): Unit = { | ||
| val job = jobControllerClient.cancelJob(conf.projectId, conf.region, jobId) | ||
| job.getDone | ||
| } | ||
|
Comment on lines
+42
to
+45
Contributor
There was a problem hiding this comment. Add error handling to kill operation The kill method lacks error handling and doesn't verify the completion status. Apply this diff: - override def kill(jobId: String): Unit = {
- val job = jobControllerClient.cancelJob(conf.projectId, conf.region, jobId)
- job.getDone
+ override def kill(jobId: String): Boolean = {
+ try {
+ val job = jobControllerClient.cancelJob(conf.projectId, conf.region, jobId)
+ if (!job.getDone) {
+ logger.warn(s"Job $jobId cancellation requested but not yet complete")
+ }
+ job.getDone
+ } catch {
+ case e: ApiException =>
+ logger.error(s"Error cancelling job $jobId", e)
+ throw new JobCancellationException(s"Failed to cancel job $jobId", e)
+ }
}
|
||
|
|
||
| override def submit(files: List[String], args: String*): String = { | ||
| val sparkJob = SparkJob | ||
| .newBuilder() | ||
| .setMainClass(conf.mainClass) | ||
| .addJarFileUris(conf.jarUri) | ||
| // .addJarFileUris("gs://spark-lib/bigquery/spark-bigquery-with-dependencies_2.12-0.41.0.jar") | ||
| .addAllFileUris(files.asJava) | ||
| .addAllArgs(args.toIterable.asJava) | ||
| .build() | ||
|
|
||
| val jobPlacement = JobPlacement | ||
| .newBuilder() | ||
| .setClusterName(conf.clusterName) | ||
| .build() | ||
|
|
||
| try { | ||
| val job = Job | ||
| .newBuilder() | ||
| .setReference(jobReference) | ||
| .setPlacement(jobPlacement) | ||
| .setSparkJob(sparkJob) | ||
| .build() | ||
|
|
||
| val submittedJob = jobControllerClient.submitJob(conf.projectId, conf.region, job) | ||
| submittedJob.getReference.getJobId | ||
|
|
||
| } catch { | ||
| case e: ApiException => | ||
| throw new RuntimeException(s"Failed to submit job: ${e.getMessage}") | ||
| } | ||
| } | ||
|
|
||
| def jobReference: JobReference = JobReference.newBuilder().build() | ||
|
Contributor
There was a problem hiding this comment. Remove unused method or implement proper JobReference generation The Either remove this method since JobReference is now generated in the submit method, or implement proper reference generation: - def jobReference: JobReference = JobReference.newBuilder().build() |
||
| } | ||
|
|
||
| object DataprocSubmitter { | ||
| def apply(): DataprocSubmitter = { | ||
| val conf = loadConfig | ||
| val jobControllerClient = JobControllerClient.create( | ||
| JobControllerSettings.newBuilder().setEndpoint(conf.endPoint).build() | ||
| ) | ||
| new DataprocSubmitter(jobControllerClient, conf) | ||
| } | ||
|
Comment on lines
+83
to
+89
Contributor
There was a problem hiding this comment. Add proper resource management for JobControllerClient The def apply(): DataprocSubmitter = {
val conf = loadConfig
- val jobControllerClient = JobControllerClient.create(
- JobControllerSettings.newBuilder().setEndpoint(conf.endPoint).build()
- )
- new DataprocSubmitter(jobControllerClient, conf)
+ try {
+ val jobControllerClient = JobControllerClient.create(
+ JobControllerSettings.newBuilder().setEndpoint(conf.endPoint).build()
+ )
+ new DataprocSubmitter(jobControllerClient, conf) {
+ override def finalize(): Unit = {
+ try {
+ jobControllerClient.close()
+ } catch {
+ case _: Exception => // Already closed or failed to close
+ }
+ super.finalize()
+ }
+ }
+ } catch {
+ case e: Exception =>
+ throw new RuntimeException("Failed to create JobControllerClient", e)
+ }
}
|
||
|
|
||
| def loadConfig: SubmitterConf = { | ||
| val is = getClass.getClassLoader.getResourceAsStream("dataproc-submitter-conf.yaml") | ||
|
Contributor
There was a problem hiding this comment. Handle possible
Add a null check for the input stream: val is = getClass.getClassLoader.getResourceAsStream("dataproc-submitter-conf.yaml")
+ if (is == null) {
+ throw new FileNotFoundException("Configuration file 'dataproc-submitter-conf.yaml' not found in classpath")
+ }
val confStr = Source.fromInputStream(is).mkString
|
||
| val confStr = Source.fromInputStream(is).mkString | ||
| val res: Either[io.circe.Error, SubmitterConf] = parser | ||
| .parse(confStr) | ||
| .flatMap(_.as[SubmitterConf]) | ||
| res match { | ||
|
|
||
| case Right(v) => v | ||
| case Left(e) => throw e | ||
| } | ||
| } | ||
| } | ||
|
|
||
| object DataprocAuth extends SparkAuth {} | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,6 @@ | ||
| package ai.chronon.integrations.cloud_gcp.test | ||
|
|
||
| import org.scalatest.funsuite.AnyFunSuite | ||
| import org.scalatestplus.mockito.MockitoSugar | ||
|
|
||
| class BigQueryCatalogTest extends AnyFunSuite with MockitoSugar {} |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,63 @@ | ||
| package ai.chronon.integrations.cloud_gcp.test | ||
|
|
||
| import ai.chronon.integrations.cloud_gcp.DataprocSubmitter | ||
| import ai.chronon.integrations.cloud_gcp.SubmitterConf | ||
| import com.google.api.gax.rpc.UnaryCallable | ||
| import com.google.cloud.dataproc.v1._ | ||
| import com.google.cloud.dataproc.v1.stub.JobControllerStub | ||
| import com.google.cloud.spark.bigquery.BigQueryUtilScala | ||
| import org.junit.Assert.assertEquals | ||
| import org.mockito.ArgumentMatchers._ | ||
| import org.mockito.Mockito._ | ||
| import org.scalatest.funsuite.AnyFunSuite | ||
| import org.scalatestplus.mockito.MockitoSugar | ||
|
|
||
| class DataprocSubmitterTest extends AnyFunSuite with MockitoSugar { | ||
|
|
||
| test("DataprocClient should return job id when a job is submitted") { | ||
|
|
||
| // Mock dataproc job client. | ||
| val jobId = "mock-job-id" | ||
| val mockJob = Job | ||
| .newBuilder() | ||
| .setReference(JobReference.newBuilder().setJobId(jobId)) | ||
| .setStatus(JobStatus.newBuilder().setState(JobStatus.State.DONE)) | ||
| .build() | ||
|
|
||
| val mockJobControllerStub = mock[JobControllerStub] | ||
| val mockSubmitJobCallable = mock[UnaryCallable[SubmitJobRequest, Job]] | ||
|
|
||
| when(mockSubmitJobCallable.call(any())) | ||
| .thenReturn(mockJob) | ||
|
|
||
| when(mockJobControllerStub.submitJobCallable) | ||
| .thenReturn(mockSubmitJobCallable) | ||
|
|
||
| val mockJobControllerClient = JobControllerClient.create(mockJobControllerStub) | ||
|
|
||
| // Test starts here. | ||
|
|
||
| val submitter = new DataprocSubmitter( | ||
| mockJobControllerClient, | ||
| SubmitterConf("test-project", "test-region", "test-cluster", "test-jar-uri", "test-main-class")) | ||
|
|
||
| val submittedJobId = submitter.submit(List.empty) | ||
| assertEquals(submittedJobId, jobId) | ||
| } | ||
|
Comment on lines
+17
to
+46
Contributor
There was a problem hiding this comment. 🛠️ Refactor suggestion Enhance test coverage with additional scenarios The current test only covers the happy path. Consider adding tests for:
Example test case for error handling: test("DataprocClient should handle job submission failure") {
val mockJobControllerStub = mock[JobControllerStub]
val mockSubmitJobCallable = mock[UnaryCallable[SubmitJobRequest, Job]]
when(mockSubmitJobCallable.call(any()))
.thenThrow(new RuntimeException("Submission failed"))
when(mockJobControllerStub.submitJobCallable)
.thenReturn(mockSubmitJobCallable)
val mockJobControllerClient = JobControllerClient.create(mockJobControllerStub)
val submitter = new DataprocSubmitter(
mockJobControllerClient,
SubmitterConf("test-project", "test-region", "test-cluster", "test-jar-uri", "test-main-class"))
assertThrows[RuntimeException] {
submitter.submit(List.empty)
}
} |
||
|
|
||
| test("Verify classpath with spark-bigquery-connector") { | ||
| BigQueryUtilScala.validateScalaVersionCompatibility() | ||
| } | ||
|
|
||
| ignore("Used to iterate locally. Do not enable this in CI/CD!") { | ||
|
|
||
| val submitter = DataprocSubmitter() | ||
| val submittedJobId = | ||
| submitter.submit(List("gs://dataproc-temp-us-central1-703996152583-pqtvfptb/jars/training_set.v1"), | ||
| "join", | ||
| "--end-date=2024-12-10", | ||
| "--conf-path=training_set.v1") | ||
| println(submittedJobId) | ||
| } | ||
|
Comment on lines
+52
to
+61
Contributor
There was a problem hiding this comment. Restructure integration test and remove hardcoded paths
Consider:
Example structure: class DataprocSubmitterIntegrationTest extends AnyFunSuite {
val config = ConfigFactory.load("integration-test.conf")
test("End-to-end job submission") {
val submitter = DataprocSubmitter()
val jarPath = config.getString("test.jar.path")
val confPath = config.getString("test.conf.path")
val submittedJobId = submitter.submit(
List(jarPath),
"join",
s"--end-date=${LocalDate.now}",
s"--conf-path=${confPath}")
// Add assertions for job completion
}
} |
||
|
|
||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,14 @@ | ||
| package ai.chronon.spark | ||
|
|
||
| trait SparkSubmitter { | ||
|
|
||
| def submit(files: List[String], args: String*): String | ||
|
|
||
| def status(jobId: String): Unit | ||
|
|
||
| def kill(jobId: String): Unit | ||
| } | ||
|
|
||
| abstract class SparkAuth { | ||
| def token(): Unit = {} | ||
| } |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
💡 Codebase verification
Tests need improvement for better coverage
The DataprocSubmitter implementation has unit tests, but they are insufficient:
submitmethod is tested with a basic happy pathsubmitmethodstatusandkillmethodsRequired test coverage:
🔗 Analysis chain
Add unit tests for the Dataproc submitter implementation
The PR objectives indicate missing unit tests, yet mockito-core is added as a test dependency. Please implement unit tests for the Dataproc submitter functionality.
Also applies to: 207-207
Let me check the content of the test file to verify if it adequately tests the DataprocSubmitter implementation.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
Length of output: 6273
Script:
Length of output: 2081