Add support for SGC run resumption from Databricks Jobs by artjen · Pull Request #19015 · mlflow/mlflow

artjen · 2025-11-24T22:06:10Z

🛠 DevTools 🛠

Install mlflow from this PR

# mlflow
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/19015/merge
# mlflow-skinny
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/19015/merge#subdirectory=libs/skinny

For Databricks, use the following command:

%sh curl -LsSf https://raw.githubusercontent.com/mlflow/mlflow/HEAD/dev/install-skinny.sh | sh -s pull/19015/merge

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

Automatically associate a Databricks jobId with a MLflow run_id when specifying Databricks job parameter SERVERLESS_GPU_COMPUTE_ASSOCIATED_JOB_RUN_ID. This allows MLflow runs to be resumed across Databricks Job invocations automatically.

How is this PR tested?

Existing unit/integration tests
New unit/integration tests
Manual tests: https://e2-spog.staging.cloud.databricks.com/editor/notebooks/2450900058003568?o=6051921418418893

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

No. You can skip the rest of this section.
Yes. Give a description of this change to be included in the release notes for MLflow users.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

How should the PR be classified in the release notes? Choose one:

rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
rn/feature - A new user-facing feature worth mentioning in the release notes
rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
rn/documentation - A user-facing documentation change worth mentioning in the release notes

Should this PR be included in the next patch release?

Yes should be selected for bug fixes, documentation updates, and other small changes. No should be selected for new features and larger changes. If you're unsure about the release classification of this PR, leave this unchecked to let the maintainers decide.

What is a minor/patch release?

Minor release: a release that increments the second part of the version number (e.g., 1.2.0 -> 1.3.0).
Bug fixes, doc updates and new features usually go into minor releases.
Patch release: a release that increments the third part of the version number (e.g., 1.2.0 -> 1.2.1).
Bug fixes and doc updates usually go into patch releases.

Yes (this PR will be cherry-picked and included in the next patch release)
No (this PR will be included in the next minor release)

Signed-off-by: Arpit Jasapara <arpit.jasapara@databricks.com>

github-actions · 2025-11-24T22:06:26Z

@artjen Thank you for the contribution! Could you fix the following issue(s)?

⚠ DCO check

The DCO check failed. Please sign off your commit(s) by following the instructions here. See https://github.com/mlflow/mlflow/blob/master/CONTRIBUTING.md#sign-your-work for more details.

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

github-actions · 2025-11-24T22:34:15Z

Documentation preview for 336d0c0 is available at:

https://pr-19015--mlflow-docs-preview.netlify.app/docs/latest/

More info

Ignore this comment if this PR does not change the documentation.
The preview is updated when a new commit is pushed to this PR.
This comment was created by this workflow run.
The documentation was built by this workflow run.

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

mlflow/tracking/fluent.py

mlflow/utils/databricks_utils.py

tests/tracking/fluent/test_fluent.py

mlflow/environment_variables.py

mlflow/tracking/fluent.py

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

mlflow/tracking/fluent.py

mlflow/environment_variables.py

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

mlflow/tracking/fluent.py

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

mlflow/tracking/fluent.py

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

harupy · 2025-11-25T03:15:31Z

mlflow/tracking/fluent.py

+                client.set_experiment_tag(
+                    exp_id_for_run, sgc_job_run_id_tag_key, active_run_obj.info.run_id
+                )


On Databricks, maximum allowed tags per experiment is limited 1000. Could this be an issue?

Users have to specify the SERVERLESS_GPU_COMPUTE_ASSOCIATED_JOB_RUN_ID parameter in the running job to take effect:

try: return dbutils.jobs.taskValues.get( "SERVERLESS_GPU_COMPUTE_ASSOCIATED_JOB_RUN_ID", debugValue=None ) except Exception as e: _logger.debug(f"Failed to retrieve SGC job run ID from task values: {e}", exc_info=True) return None

This means that non-SGC users will not add new experiment tags by default.

For SGC users, we can increase the tag limit on a case-by-case basis if this becomes an issue.

Sounds good!

mlflow/tracking/fluent.py

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

XiaohanZhangCMU · 2025-11-25T19:25:54Z

mlflow/utils/databricks_utils.py

+        return dbutils.jobs.taskValues.get(
+            "SERVERLESS_GPU_COMPUTE_ASSOCIATED_JOB_RUN_ID", debugValue=None
+        )
+    except Exception as e:


[non blocker] is it possible to have more specific exceptions here?

I want to ensure this will not disrupt the normal start_run flow. The error is reprinted in a debug message below for debugging if necessary.

mlflow/utils/databricks_utils.py

XiaohanZhangCMU

thanks looks nice! left two nits.

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

harupy

LGTM!

arpitjasa-db added 2 commits November 21, 2025 14:06

Add support for SGC run resumption from Databricks Jobs

ca85872

Signed-off-by: Arpit Jasapara <arpit.jasapara@databricks.com>

fix lint

bd44dda

Signed-off-by: Arpit Jasapara <arpit.jasapara@databricks.com>

github-actions bot added the v3.6.1 label Nov 24, 2025

artjen requested a review from harupy November 24, 2025 22:06

github-actions bot added area/tracking Tracking service, tracking client APIs, autologging rn/feature Mention under Features in Changelogs. labels Nov 24, 2025

artjen mentioned this pull request Nov 24, 2025

Add support for SGC run resumption from Databricks Jobs #18965

Closed

29 tasks

address review comments

880fb6b

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

artjen force-pushed the sgcRunResume branch from 83ef930 to 880fb6b Compare November 24, 2025 22:25

lint

6b3000d

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

harupy reviewed Nov 25, 2025

View reviewed changes

mlflow/tracking/fluent.py Outdated Show resolved Hide resolved

harupy reviewed Nov 25, 2025

View reviewed changes

mlflow/tracking/fluent.py Outdated Show resolved Hide resolved

harupy reviewed Nov 25, 2025

View reviewed changes

mlflow/tracking/fluent.py Outdated Show resolved Hide resolved

harupy reviewed Nov 25, 2025

View reviewed changes

mlflow/tracking/fluent.py Outdated Show resolved Hide resolved

harupy reviewed Nov 25, 2025

View reviewed changes

mlflow/utils/databricks_utils.py Outdated Show resolved Hide resolved

harupy reviewed Nov 25, 2025

View reviewed changes

tests/tracking/fluent/test_fluent.py Outdated Show resolved Hide resolved

harupy reviewed Nov 25, 2025

View reviewed changes

tests/tracking/fluent/test_fluent.py Outdated Show resolved Hide resolved

harupy reviewed Nov 25, 2025

View reviewed changes

mlflow/environment_variables.py Outdated Show resolved Hide resolved

harupy reviewed Nov 25, 2025

View reviewed changes

mlflow/tracking/fluent.py Outdated Show resolved Hide resolved

harupy reviewed Nov 25, 2025

View reviewed changes

mlflow/tracking/fluent.py Outdated Show resolved Hide resolved

default true and comments

9493c57

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

harupy reviewed Nov 25, 2025

View reviewed changes

mlflow/tracking/fluent.py Outdated Show resolved Hide resolved

harupy reviewed Nov 25, 2025

View reviewed changes

mlflow/tracking/fluent.py Outdated Show resolved Hide resolved

harupy reviewed Nov 25, 2025

View reviewed changes

mlflow/environment_variables.py Outdated Show resolved Hide resolved

lint and imports and tests

af41d74

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

harupy reviewed Nov 25, 2025

View reviewed changes

mlflow/tracking/fluent.py Outdated Show resolved Hide resolved

typing

e0cc03b

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

harupy reviewed Nov 25, 2025

View reviewed changes

mlflow/tracking/fluent.py Outdated Show resolved Hide resolved

whitespace

3dc61eb

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

artjen requested a review from harupy November 25, 2025 02:34

artjen added 2 commits November 25, 2025 03:01

system tag convention

440f1e8

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

lint

59bf48e

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

harupy reviewed Nov 25, 2025

View reviewed changes

mlflow/tracking/fluent.py Show resolved Hide resolved

init str | None

f86012c

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

artjen requested a review from harupy November 25, 2025 18:06

XiaohanZhangCMU reviewed Nov 25, 2025

View reviewed changes

mlflow/utils/databricks_utils.py Show resolved Hide resolved

XiaohanZhangCMU reviewed Nov 25, 2025

View reviewed changes

widgets and comments

bf350af

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

artjen requested a review from XiaohanZhangCMU November 25, 2025 23:03

remove Py4JJavaError

336d0c0

Signed-off-by: Arthur Jenoudet <arthur.jenoudet@databricks.com>

harupy approved these changes Nov 26, 2025

View reviewed changes

artjen added this pull request to the merge queue Nov 26, 2025

Merged via the queue into mlflow:master with commit f08459e Nov 26, 2025
48 of 50 checks passed

artjen deleted the sgcRunResume branch November 26, 2025 02:29

artjen mentioned this pull request Dec 1, 2025

Include environment variable fallback for SGC run resumption #19143

Merged

29 tasks

Conversation

artjen commented Nov 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Install mlflow from this PR

Related Issues/PRs

What changes are proposed in this pull request?

How is this PR tested?

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

What component(s), interfaces, languages, and integrations does this PR affect?

How should the PR be classified in the release notes? Choose one:

Should this PR be included in the next patch release?

Uh oh!

github-actions bot commented Nov 24, 2025

⚠ DCO check

Uh oh!

github-actions bot commented Nov 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

harupy Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

artjen Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

harupy Nov 26, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

XiaohanZhangCMU Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

artjen Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

XiaohanZhangCMU left a comment

Choose a reason for hiding this comment

Uh oh!

harupy left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

artjen commented Nov 24, 2025 •

edited

Loading

github-actions bot commented Nov 24, 2025 •

edited

Loading