[ADAP-569] [Enhancement] only wait for Dataproc job to finish, not for the Dataproc serverless cluster to spin down

[wazi55]

Is this a new bug in dbt-bigquery?

• I believe this is a new bug in dbt-bigquery
• I have searched the existing issues, and I could not find an existing issue for this bug

Current Behavior

When working with DBT with Dataproc, the current behavior is DBT waits for the batch job complete the torn down of the cluster here, however, that adds an additional 1 or 2 minutes on the computation time. Instead of waiting for the job finish, the code should just wait until the job is finished by polling using the get_batch() method to check whether the job has succeeded or not.

Expected Behavior

The expected behavior is, DBT should continue on with the SQL models after the spark code succeed instead of waiting for the cluster to be torn down.

Steps To Reproduce

    def _submit_dataproc_job(self) -> dataproc_v1.types.jobs.Job:
        batch = self._configure_batch()
        parent = f"projects/{self.credential.execution_project}/locations/{self.credential.dataproc_region}"

        request = dataproc_v1.CreateBatchRequest(
            parent=parent,
            batch=batch,
        )
        # make the request
        operation = self.job_client.create_batch(request=request)  # type: ignore
        # this takes quite a while, waiting on GCP response to resolve
        # (not a google-api-core issue, more likely a dataproc serverless issue)
        response = operation.result(retry=self.retry)
        return response

Relevant log output

No response

Environment

- OS:
- Python:
- dbt-core:
- dbt-bigquery:

Additional Context

No response

[dbeatty10]

Thanks for reporting this @wazi55!

Are you interested in contributing a pull request for this, by any chance?

[dbeatty10]

I re-labeled this as an enhancement since I don't perceive there to be an error, flaw, failure or fault here -- this is more of an efficiency / optimization thing.

[dataders]

@wazi55 thanks so much for opening! This was on my to-do list to write up after an email thread last month with Big Query engineers.

Their recommendation was to use BatchControllerClient's .get_batch() instead of create_batch() to create DataProc jobs (relevant Google Python SDK docs).

afaict, this is not a drop-in replacement as some polling would have to be implemented to ensure that the response's state attribute (docs) is one of: SUCCEEDED, CANCELLED, or FAILED.

Perusing the docs, I also notice that there is a BatchControllerAsyncClient (docs), which might obviate the need for polling? I'm no expert here.

I'm going to ask the GCP team to validate my thinking.

[wazi55]

Hello 👋 get_batch is exactly the method I was going to suggest, basically check the returned BATCH.STATE should be the way to go.

[dataders]

resolved by #929