Transient Clusters Triggered by teams.py

api/python/test/canary/teams.py

@@ -23,9 +23,9 @@
             "CUSTOMER_ID": "dev",
             "GCP_PROJECT_ID": "canary-443022",
             "GCP_REGION": "us-central1",
-            "GCP_DATAPROC_CLUSTER_NAME": "zipline-canary-cluster",
             "GCP_BIGTABLE_INSTANCE_ID": "zipline-canary-instance",
             "FLINK_STATE_URI": "gs://zipline-warehouse-canary/flink-state",
+            "GCP_DATAPROC_CLUSTER_NAME": "zipline-canary-cluster"
         },
     ),
 )
@@ -61,9 +61,21 @@
             "CUSTOMER_ID": "dev",
             "GCP_PROJECT_ID": "canary-443022",
             "GCP_REGION": "us-central1",
-            "GCP_DATAPROC_CLUSTER_NAME": "zipline-canary-cluster",
             "GCP_BIGTABLE_INSTANCE_ID": "zipline-canary-instance",
+            "GCP_DATAPROC_CLUSTER_NAME": "zipline-canary-cluster",
         },
+        modeEnvironments={
+            RunMode.BACKFILL: {
+                "GCP_CREATE_DATAPROC": "true",
+                "GCP_DATAPROC_NUM_WORKERS": "2",
+                "ARTIFACT_PREFIX": "gs://zipline-artifacts-canary",
+            },
+            RunMode.UPLOAD: {
+                "GCP_CREATE_DATAPROC": "true",
+                "GCP_DATAPROC_NUM_WORKERS": "2",
+                "ARTIFACT_PREFIX": "gs://zipline-artifacts-canary",
+            }
+        }
     ),
     conf=ConfigProperties(
         common={

cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/DataprocSubmitter.scala

@@ -3,6 +3,7 @@ import ai.chronon.spark.submission.JobSubmitterConstants._
 import ai.chronon.spark.submission.{JobSubmitter, JobType, FlinkJob => TypeFlinkJob, SparkJob => TypeSparkJob}
 import com.google.api.gax.rpc.ApiException
 import com.google.cloud.dataproc.v1._
+import com.google.protobuf.Duration
 import org.apache.hadoop.fs.Path
 import org.json4s._
 import org.json4s.jackson.JsonMethods._
@@ -351,8 +352,16 @@ object DataprocSubmitter {
   private def initializeDataprocSubmitter(): DataprocSubmitter = {
     val projectId = sys.env.getOrElse(GcpProjectIdEnvVar, throw new Exception(s"$GcpProjectIdEnvVar not set"))
     val region = sys.env.getOrElse(GcpRegionEnvVar, throw new Exception(s"$GcpRegionEnvVar not set"))
-    val clusterName = sys.env
-      .getOrElse(GcpDataprocClusterNameEnvVar, throw new Exception(s"$GcpDataprocClusterNameEnvVar not set"))
+    val clusterName = if (sys.env.getOrElse(GcpCreateDataprocEnvVar, "false").toBoolean) {
+      val dataprocClient = ClusterControllerClient.create(
+        ClusterControllerSettings.newBuilder().setEndpoint(s"$region-dataproc.googleapis.com:443").build())
+      createDataprocCluster(projectId, region, dataprocClient)
+    } else if (sys.env.contains(GcpDataprocClusterNameEnvVar)) {
+      sys.env
+        .getOrElse(GcpDataprocClusterNameEnvVar, throw new Exception(s"$GcpDataprocClusterNameEnvVar not set"))
+    } else {
+      throw new Exception(s"Either $GcpDataprocClusterNameEnvVar or $GcpCreateDataprocEnvVar must be set, but neither are")
+    }
 
     val submitterConf = SubmitterConf(
       projectId,
@@ -517,6 +526,179 @@ object DataprocSubmitter {
     }
   }
 
+
+  private def buildClusterConfig(projectId: String, artifact_prefix: String): ClusterConfig = {
+    val numWorkers = sys.env
+      .getOrElse(GcpDataprocNumWorkersEnvVar, throw new Exception(s"$GcpCreateDataprocEnvVar is true but $GcpDataprocNumWorkersEnvVar not set"))
+      .toInt
+    val hostType = sys.env
+      .getOrElse(GcpDataprocHostTypeEnvVar, "n2-highmem-4")
+    val networkUri = sys.env
+      .getOrElse(GcpDataprocNetworkEnvVar, "default")
+    val initializationActions = sys.env
+      .getOrElse(GcpDataprocInitializationActionsEnvVar, "").split(",").toList
+    val tags = sys.env
+      .getOrElse(GcpDataprocTagsEnvVar, "").split(",").toList
+
+    // Build the cluster configuration with autoscaling
+    val config = ClusterConfig
+      .newBuilder()
+      .setMasterConfig(
+        InstanceGroupConfig
+          .newBuilder()
+          .setNumInstances(1)
+          .setMachineTypeUri("n2-highmem-64") // Adjust machine type as needed
+          .setDiskConfig(
+            DiskConfig
+              .newBuilder()
+              .setBootDiskType("pd-standard") // Use SSD for better performance
+              .setBootDiskSizeGb(1024) // Adjust disk size as needed
+              .build()
+          )
+          .build()
+      )
+      .setWorkerConfig(
+        InstanceGroupConfig
+          .newBuilder()
+          .setNumInstances(numWorkers) // Initial number of worker nodes. Autoscaling will adjust this
+          .setMachineTypeUri(hostType)
+          .setDiskConfig(
+            DiskConfig
+              .newBuilder()
+              .setBootDiskType("pd-standard")
+              .setBootDiskSizeGb(64)
+              .setNumLocalSsds(2)
+              .build()
+          )
+          .build()
+      )
+
+      val gceClusterConfig = GceClusterConfig
+        .newBuilder()
+        .setNetworkUri(networkUri)
+        .setServiceAccount(f"dataproc@$projectId.iam.gserviceaccount.com")
+        .addAllServiceAccountScopes(
+          List(
+            "https://www.googleapis.com/auth/cloud-platform",
+            "https://www.googleapis.com/auth/cloud.useraccounts.readonly",
+            "https://www.googleapis.com/auth/devstorage.read_write",
+            "https://www.googleapis.com/auth/logging.write"
+          ).asJava
+        )
+        .putMetadata("hive-version", "3.1.2")
+        .putMetadata("SPARK_BQ_CONNECTOR_URL", "gs://spark-lib/bigquery/spark-3.5-bigquery-0.42.1.jar")
+        .putMetadata("artifact_prefix", artifact_prefix)
+        .setInternalIpOnly(true)
+
+      for(tag <- tags if tag != "") {
+        gceClusterConfig
+        .addTags(tag)
+      }
+
+      config
+      .setGceClusterConfig(
+        gceClusterConfig.build()
+      )
+      .setSoftwareConfig(
+        SoftwareConfig
+          .newBuilder()
+          .setImageVersion("2.2.50-debian12")
+          .addOptionalComponents(Component.FLINK)
+          .addOptionalComponents(Component.JUPYTER)
+          .putProperties("flink:env.java.opts.client",
+            "-Djava.net.preferIPv4Stack=true -Djava.security.properties=/etc/flink/conf/java.security")
+          .build()
+      )
+      .setEndpointConfig(
+        EndpointConfig
+          .newBuilder()
+          .setEnableHttpPortAccess(true)
+          .build()
+      )
+      .setLifecycleConfig(
+        LifecycleConfig
+          .newBuilder()
+          .setIdleDeleteTtl(
+            Duration
+              .newBuilder()
+              .setSeconds(7200) // 2 hours
+              .build()
+          )
+          .build()
+      )
+      .addInitializationActions(
+        NodeInitializationAction
+          .newBuilder()
+          .setExecutableFile(f"$artifact_prefix/scripts/copy_java_security.sh")
+          .build()
+      )
+
+    for(action <- initializationActions if action != "") {
+        config.addInitializationActions(
+            NodeInitializationAction
+            .newBuilder()
+            .setExecutableFile(action)
+            .build()
+        )
+    }
+
+    config.build()
+  }
+
+  private[cloud_gcp] def createDataprocCluster(projectId: String, region: String,
+                                              dataprocClient: ClusterControllerClient): String = {
+    val artifact_prefix = sys.env
+      .getOrElse(ArtifactPrefixEnvVar, throw new Exception(s"$ArtifactPrefixEnvVar not set"))
+
+    val clusterConfig = buildClusterConfig(projectId, artifact_prefix)
+
+    val clusterName = s"zipline-transient-cluster-${System.currentTimeMillis()}"
+
+    val cluster: Cluster = Cluster
+      .newBuilder()
+      .setClusterName(clusterName)
+      .setProjectId(projectId)
+      .setConfig(clusterConfig)
+      .build()
+
+    val createRequest = CreateClusterRequest
+      .newBuilder()
+      .setProjectId(projectId)
+      .setRegion(region)
+      .setCluster(cluster)
+      .build()
+
+    // Asynchronously create the cluster and wait for it to be ready
+    dataprocClient
+      .createClusterAsync(createRequest)
+      .get(5, java.util.concurrent.TimeUnit.MINUTES) match {
+      case null =>
+        throw new RuntimeException("Failed to create Dataproc cluster.")
+      case _ =>
+        println(s"Created Dataproc cluster: $clusterName")
+    }
+    // Check status of the cluster creation
+    var currentState = dataprocClient.getCluster(projectId, region, clusterName).getStatus.getState
+    while (
+      currentState != ClusterStatus.State.RUNNING &&
+        currentState != ClusterStatus.State.ERROR &&
+        currentState != ClusterStatus.State.STOPPING
+    ) {
+      println(s"Waiting for Dataproc cluster $clusterName to be in RUNNING state. Current state: $currentState")
+      Thread.sleep(30000) // Wait for 30 seconds before checking again
+      currentState = dataprocClient.getCluster(projectId, region, clusterName).getStatus.getState
+    }
+    currentState match {
+      case ClusterStatus.State.RUNNING =>
+        println(s"Dataproc cluster $clusterName is running.")
+        clusterName
+      case ClusterStatus.State.ERROR =>
+        throw new RuntimeException(s"Failed to create Dataproc cluster $clusterName: ERROR state.")
+      case _ =>
+        throw new RuntimeException(s"Dataproc cluster $clusterName is in unexpected state: $currentState.")
+    }
+  }
+
   private[cloud_gcp] def run(args: Array[String],
                              submitter: DataprocSubmitter,
                              envMap: Map[String, Option[String]] = Map.empty): Unit = {

cloud_gcp/src/test/scala/ai/chronon/integrations/cloud_gcp/DataprocSubmitterTest.scala

@@ -3,6 +3,9 @@ package ai.chronon.integrations.cloud_gcp
 import ai.chronon.spark
 import ai.chronon.spark.submission
 import ai.chronon.spark.submission.JobSubmitterConstants._
+import com.google.api.core.ApiFuture
+import com.google.api.gax.longrunning.{OperationFuture, OperationSnapshot}
+import com.google.api.gax.retrying.RetryingFuture
 import com.google.api.gax.rpc.UnaryCallable
 import com.google.cloud.dataproc.v1.JobControllerClient.ListJobsPagedResponse
 import com.google.cloud.dataproc.v1._
@@ -15,9 +18,18 @@ import org.scalatest.flatspec.AnyFlatSpec
 import org.scalatestplus.mockito.MockitoSugar
 
 import java.nio.file.Paths
+import java.util.concurrent.TimeUnit
 import scala.jdk.CollectionConverters._
 
 class DataprocSubmitterTest extends AnyFlatSpec with MockitoSugar {
+  def setEnv(key: String, value: String): Unit = {
+    val env = System.getenv()
+    val field = env.getClass.getDeclaredField("m")
+    field.setAccessible(true)
+    val writableEnv = field.get(env).asInstanceOf[java.util.Map[String, String]]
+    writableEnv.put(key, value)
+  }
+
   it should "test buildFlinkJob with the expected flinkStateUri and savepointUri" in {
     val submitter = new DataprocSubmitter(jobControllerClient = mock[JobControllerClient],
                                           conf = SubmitterConf("test-project", "test-region", "test-cluster"))
@@ -715,6 +727,39 @@ class DataprocSubmitterTest extends AnyFlatSpec with MockitoSugar {
 
   }
 
+  it should "create a Dataproc cluster successfully" in {
+    val mockDataprocClient = mock[ClusterControllerClient]
+
+    val mockOperationFuture = mock[OperationFuture[Cluster, ClusterOperationMetadata]]
+    val mockRetryingFuture = mock[RetryingFuture[OperationSnapshot]]
+    val mockMetadataFuture = mock[ApiFuture[ClusterOperationMetadata]]
+    val mockCluster = Cluster
+      .newBuilder()
+      .setStatus(ClusterStatus.newBuilder().setState(ClusterStatus.State.RUNNING))
+      .build()
+
+    when(mockDataprocClient.createClusterAsync(any[CreateClusterRequest]))
+      .thenReturn(mockOperationFuture)
+    when(mockOperationFuture.getPollingFuture).thenReturn(mockRetryingFuture)
+    when(mockOperationFuture.peekMetadata()).thenReturn(mockMetadataFuture)
+    when(mockOperationFuture.get(anyLong(), any[TimeUnit])).thenReturn(mockCluster)
+
+    when(mockDataprocClient.createClusterAsync(any[CreateClusterRequest]))
+      .thenReturn(mockOperationFuture)
+
+    when(mockDataprocClient.getCluster(any[String], any[String], any[String])).thenReturn(mockCluster)
+
+
+    val region = "test-region"
+    val projectId = "test-project"
+    setEnv(ArtifactPrefixEnvVar, "gs://test-bucket")
+    setEnv(GcpDataprocNumWorkersEnvVar, "2")
+
+    val clusterName = DataprocSubmitter.createDataprocCluster(region, projectId, mockDataprocClient)
+
+    verify(mockDataprocClient).createClusterAsync(any())
+  }
+
   it should "test getZiplineVersionOfDataprocJob successfully" in {
     val jobId = "mock-job-id"
     val mockJob = mock[Job]

spark/src/main/scala/ai/chronon/spark/submission/JobSubmitter.scala

@@ -178,6 +178,16 @@ object JobSubmitterConstants {
   val GcpRegionEnvVar = "GCP_REGION"
   val GcpDataprocClusterNameEnvVar = "GCP_DATAPROC_CLUSTER_NAME"
 
+  // Transient Dataproc creation
+  val GcpCreateDataprocEnvVar = "GCP_CREATE_DATAPROC"
+  val GcpDataprocNumWorkersEnvVar = "GCP_DATAPROC_NUM_WORKERS"
+  val GcpDataprocHostTypeEnvVar = "GCP_DATAPROC_HOST_TYPE"
+  val GcpDataprocNetworkEnvVar = "GCP_DATAPROC_NETWORK"
+  val GcpDataprocTagsEnvVar = "GCP_DATAPROC_TAGS"
+  val GcpDataprocInitializationActionsEnvVar = "GCP_DATAPROC_INITIALIZATION_ACTIONS"
+
+  val ArtifactPrefixEnvVar = "ARTIFACT_PREFIX"
+
   val CheckIfJobIsRunning = "check-if-job-is-running"
   val StreamingDeploy = "deploy"