[dagster-databricks] fix issue with prior_attempts_count (#8639)

dagster-io · Jun 27, 2022 · 54b173f · 54b173f
1 parent adeda18
commit 54b173f
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/python_modules/libraries/dagster-databricks/dagster_databricks/databricks.py b/python_modules/libraries/dagster-databricks/dagster_databricks/databricks.py
@@ -271,7 +271,7 @@ def poll_run_state(
             log.error(error_message)
             raise DatabricksError(error_message)
     else:
-        log.info("Run %s in state %s" % (databricks_run_id, run_state))
+        log.debug("Run %s in state %s" % (databricks_run_id, run_state))
     if time.time() - start_poll_time > max_wait_time_sec:
         raise DatabricksError(
             "Job run {} took more than {}s to complete; failing".format(

diff --git a/python_modules/libraries/dagster-databricks/dagster_databricks/databricks_step_main.py b/python_modules/libraries/dagster-databricks/dagster_databricks/databricks_step_main.py
@@ -92,9 +92,15 @@ def main(
         print("Running dagster job")  # noqa pylint: disable=print-call
 
         step_run_dir = os.path.dirname(step_run_ref_filepath)
+        if step_run_ref.known_state is not None:
+            attempt_count = step_run_ref.known_state.get_retry_state().get_attempt_count(
+                step_run_ref.step_key
+            )
+        else:
+            attempt_count = 0
         events_filepath = os.path.join(
             step_run_dir,
-            f"{step_run_ref.prior_attempts_count}_{PICKLED_EVENTS_FILE_NAME}",
+            f"{attempt_count}_{PICKLED_EVENTS_FILE_NAME}",
         )
         stdout_filepath = os.path.join(step_run_dir, "stdout")
         stderr_filepath = os.path.join(step_run_dir, "stderr")