[BUG] Flyte task keeps running forever when running a Databricks job

[rambrus]

Describe the bug

BACKGROUND:
I'm trying to run a simplified Flyte task using Databricks plugin.

PREREQUISITES:

• Flyte 1.5.0 installed on EKS cluster
• Flyte Databricks plugins is configured as described here: https://docs.flyte.org/en/latest/deployment/plugins/webapi/databricks.html
• entrypoint.py is uploaded to DBFS
• Databricks cluster is created with the following settings: DBR 11.3 LTS, installed libs: flytekit==1.5.0, flytekitplugins-spark==1.5.0
• A simplified Flyte workflow is created:

@task(
    task_config=Databricks(
        databricks_conf={
           "run_name": "dbx simplified example",
           "existing_cluster_id": "<my-existing-cluster-id>",
           "timeout_seconds": 3600,
           "max_retries": 1,
       }
    ),
    limits=Resources(mem="2000M"),
    cache_version="1",
)
def print_spark_config():
    spark = flytekit.current_context().spark_session
    print(spark.sparkContext.getConf().getAll())

@workflow
def my_databricks_job():
    print_spark_config()

STEPS:

• Run workflow: pyflyte --verbose run --remote --destination-dir . dbx_simplified_example.py my_databricks_job

ISSUE:
Databricks job run triggered and successfully completed, but the Flyte job status is not updated, it is stuck in RUNNING state.

Expected behavior

Flyte job status keeps updated.

Additional context to reproduce

Please note that we tried to upgrade Flyte to v1.7.0:

• when I also upgraded the Flyte packages on the Databricks cluster to 1.7.0 -> [BUG] loader must define exec_module() when running Databricks task #3853
• when I used the 1.5.0 Flyte packages on the Databricks cluster -> same issue - Databricks job successfully completed, but Flyte job keeps running forever

Screenshots

No response

Are you sure this issue hasn't been raised already?

• Yes

Have you read the Code of Conduct?

• Yes

[rambrus]

Tested Databricks API response on this:

curl --netrc --request GET --header "Authorization: Bearer $DATABRICKS_TOEN" \
'https://dbc-32fcad04-13c2.cloud.databricks.com/api/2.0/jobs/runs/get?run_id=306'

Response:

{
    "attempt_number": 0,
    "cleanup_duration": 0,
    "cluster_instance": {
        "cluster_id": "<my-cluster-id>",
        "spark_context_id": "<my-spark-context-id>"
    },
    "cluster_spec": {
        "existing_cluster_id": "<my-cluster-id>"
    },
    "creator_user_name": "<my-username>",
    "end_time": 1688987784820,
    "execution_duration": 223000,
    "format": "SINGLE_TASK",
    "job_id": 1060720031042619,
    "number_in_job": 574539,
    "run_id": 574539,
    "run_name": "dbx simplified example",
    "run_page_url": "<my-run-page-url>",
    "run_type": "SUBMIT_RUN",
    "setup_duration": 41000,
    "start_time": 1688987520036,
    "state": {
        "life_cycle_state": "TERMINATED",
        "result_state": "SUCCESS",
        "state_message": "",
        "user_cancelled_or_timedout": false
    },
    "task": {
        "spark_python_task": {
            "parameters": [
                "pyflyte-fast-execute",
                "--additional-distribution",
                "s3://<my-s3-bucket>/flytesnacks/development/UMZ6XPNM4L6KL4YALV56QDMSX4======/script_mode.tar.gz",
                "--dest-dir",
                ".",
                "--",
                "pyflyte-execute",
                "--inputs",
                "s3://<my-s3-bucket>/metadata/propeller/flytesnacks-development-ff83ea058624d44ddbe9/n0/data/inputs.pb",
                "--output-prefix",
                "s3://<my-s3-bucket>/metadata/propeller/flytesnacks-development-ff83ea058624d44ddbe9/n0/data/0",
                "--raw-output-data-prefix",
                "s3://<my-s3-bucket>/raw_data/sh/ff83ea058624d44ddbe9-n0-0",
                "--checkpoint-path",
                "s3://<my-s3-bucket>/raw_data/sh/ff83ea058624d44ddbe9-n0-0/_flytecheckpoints",
                "--prev-checkpoint",
                "\"\"",
                "--resolver",
                "flytekit.core.python_auto_container.default_task_resolver",
                "--",
                "task-module",
                "dbx_simplified_example",
                "task-name",
                "print_spark_config"
            ],
            "python_file": "dbfs:/tmp/flyte/entrypoint.py"
        }
    }
}

[eapolinario]

Fixed in #4206.