Add GCS token refresh for log_artifacts #13397

kriscon-db · 2024-10-11T16:01:25Z

🛠 DevTools 🛠

Install mlflow from this PR

pip install git+https://github.com/mlflow/mlflow.git@refs/pull/13397/merge

Checkout with GitHub CLI

gh pr checkout 13397

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

Added GCP token refresh attempt to log_artifacts in the gcs_artifact_repo

How is this PR tested?

Existing unit/integration tests
New unit/integration tests
Manual tests

Unit test added, updated existing tests, and verified on a GCP workspace that the refresh works to get a large model registered.

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

Interface

area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
area/docker: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
area/windows: Windows support

Language

language/r: R APIs and clients
language/java: Java APIs and clients
language/new: Proposals for new client languages

Integrations

integrations/azure: Azure and Azure ML integrations
integrations/sagemaker: SageMaker integrations
integrations/databricks: Databricks integrations

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)

github-actions · 2024-10-11T16:01:48Z

Documentation preview for 81681b5 will be available when this CircleCI job
completes successfully.

More info

Ignore this comment if this PR does not change the documentation.
It takes a few minutes for the preview to be available.
The preview is updated when a new commit is pushed to this PR.
This comment was created by https://github.com/mlflow/mlflow/actions/runs/11394412382.

github-actions · 2024-10-11T16:01:51Z

@kriscon-db 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.

kriscon-db · 2024-10-13T01:42:37Z

/autoformat

mlflow/store/artifact/gcs_artifact_repo.py

tests/store/artifact/test_gcs_artifact_repo.py

mlflow/store/artifact/gcs_artifact_repo.py

tests/store/artifact/test_gcs_artifact_repo.py

harupy · 2024-10-16T08:19:18Z

mlflow/store/artifact/gcs_artifact_repo.py

+            _logger.info(
+                "Failed to complete upload request, possibly due to credential expiration."
+                f" Attempting to refresh credentials and trying again... (Error: {e})"
+            )
+            if self.credential_refresh_def is not None:
+                new_token = self.credential_refresh_def()


Suggested change

_logger.info(

"Failed to complete upload request, possibly due to credential expiration."

f" Attempting to refresh credentials and trying again... (Error: {e})"

)

if self.credential_refresh_def is not None:

new_token = self.credential_refresh_def()

if self.credential_refresh_def is not None:

_logger.info(

"Failed to complete upload request, possibly due to credential expiration."

f" Attempting to refresh credentials and trying again... (Error: {e})"

)

new_token = self.credential_refresh_def()

This logger message is to indicate that we failed due to a potential credential expiration and that we are attempting to refresh again. I don't think it should be in the if block.

mlflow/store/artifact/artifact_repo.py

kriscon-db · 2024-10-17T01:55:20Z

/autoformat

tests/store/artifact/test_gcs_artifact_repo.py

harupy · 2024-10-17T16:02:17Z

mlflow/store/artifact/artifact_repo.py

+    provided by creds_func.
+    """
+    try:
+        first_creds = creds_func() if orig_creds is None else orig_creds


is there any chance that creds_func() throws?

Yes but that would be a UC error and we should just fail the upload at that point on subsequent retries.

harupy · 2024-10-17T16:03:37Z

mlflow/utils/_unity_catalog_utils.py

+            return {
+                "oauth_token": new_gcp_creds.oauth_token,
+            }


do we need to return a dict? can we return new_gcp_creds.oauth_token?

This is to match the pattern for credential refresh for other repositories.

kriscon-db · 2024-10-17T23:06:01Z

/autoformat

Signed-off-by: mlflow-automation <[email protected]>

Signed-off-by: mlflow-automation <[email protected]> Co-authored-by: mlflow-automation <[email protected]>

Signed-off-by: mlflow-automation <[email protected]> Co-authored-by: mlflow-automation <[email protected]> Signed-off-by: k99kurella <[email protected]>

Add GCS token refresh for log_artifacts

24e617f

github-actions bot added area/artifacts Artifact stores and artifact logging rn/none List under Small Changes in Changelogs. labels Oct 11, 2024

kriscon-db added 4 commits October 12, 2024 11:51

Bug

02bb924

Fix credential refresh logic

7a0357a

fix rest store test

3ae43d1

fix artifact repo test

44e919b

kriscon-db requested review from harupy and artjen October 13, 2024 01:41

github-actions bot added the patch-2.17.0 label Oct 13, 2024

artjen reviewed Oct 14, 2024

View reviewed changes

mlflow/store/artifact/gcs_artifact_repo.py Outdated Show resolved Hide resolved

tests/store/artifact/test_gcs_artifact_repo.py Outdated Show resolved Hide resolved

CR comments

2e907cb

kriscon-db force-pushed the gcp_credential_refresh branch from 247a183 to 2e907cb Compare October 15, 2024 22:17

github-actions bot added the patch-2.17.1 label Oct 15, 2024

fix test

d23f50c

kriscon-db requested a review from artjen October 15, 2024 23:41

artjen reviewed Oct 16, 2024

View reviewed changes

mlflow/store/artifact/gcs_artifact_repo.py Outdated Show resolved Hide resolved

mlflow/store/artifact/gcs_artifact_repo.py Outdated Show resolved Hide resolved

harupy reviewed Oct 16, 2024

View reviewed changes

tests/store/artifact/test_gcs_artifact_repo.py Outdated Show resolved Hide resolved

harupy reviewed Oct 16, 2024

View reviewed changes

mlflow/store/artifact/artifact_repo.py Outdated Show resolved Hide resolved

harupy reviewed Oct 16, 2024

View reviewed changes

mlflow/store/artifact/artifact_repo.py Outdated Show resolved Hide resolved

kriscon-db added 3 commits October 16, 2024 19:22

CR fixes and logging for testing

bc6e405

Reset the gcs bucket after each file in case we updated credentials

35c2649

CR comments

f341540

kriscon-db requested review from artjen and harupy October 17, 2024 01:54

harupy approved these changes Oct 17, 2024

View reviewed changes

harupy reviewed Oct 17, 2024

View reviewed changes

tests/store/artifact/test_gcs_artifact_repo.py Outdated Show resolved Hide resolved

harupy reviewed Oct 17, 2024

View reviewed changes

artjen approved these changes Oct 17, 2024

View reviewed changes

CR comment

72d8c88

kriscon-db force-pushed the gcp_credential_refresh branch from 3c483f9 to 72d8c88 Compare October 17, 2024 23:05

mlflow-automation added 2 commits October 17, 2024 23:07

Merge remote-tracking branch 'base/master' into gcp_credential_refresh

1c1a7b9

Autoformat: https://github.com/mlflow/mlflow/actions/runs/11394394996

81681b5

Signed-off-by: mlflow-automation <[email protected]>

kriscon-db merged commit 0862ba2 into mlflow:master Oct 18, 2024
40 checks passed

serena-ruan pushed a commit to serena-ruan/mlflow that referenced this pull request Oct 25, 2024

Add GCS token refresh for log_artifacts (mlflow#13397)

da13335

Signed-off-by: mlflow-automation <[email protected]> Co-authored-by: mlflow-automation <[email protected]>

serena-ruan pushed a commit to serena-ruan/mlflow that referenced this pull request Oct 25, 2024

Add GCS token refresh for log_artifacts (mlflow#13397)

fbe4557

Signed-off-by: mlflow-automation <[email protected]> Co-authored-by: mlflow-automation <[email protected]>

serena-ruan pushed a commit that referenced this pull request Oct 25, 2024

Add GCS token refresh for log_artifacts (#13397)

90b227e

Signed-off-by: mlflow-automation <[email protected]> Co-authored-by: mlflow-automation <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add GCS token refresh for log_artifacts #13397

Add GCS token refresh for log_artifacts #13397

kriscon-db commented Oct 11, 2024 •

edited

Loading

github-actions bot commented Oct 11, 2024 •

edited

Loading

github-actions bot commented Oct 11, 2024

kriscon-db commented Oct 13, 2024

harupy Oct 16, 2024

kriscon-db Oct 16, 2024

kriscon-db commented Oct 17, 2024

harupy Oct 17, 2024

kriscon-db Oct 17, 2024 •

edited

Loading

harupy Oct 17, 2024

artjen Oct 17, 2024

kriscon-db commented Oct 17, 2024

Add GCS token refresh for log_artifacts #13397

Add GCS token refresh for log_artifacts #13397

Conversation

kriscon-db commented Oct 11, 2024 • edited Loading

Install mlflow from this PR

Checkout with GitHub CLI

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?

github-actions bot commented Oct 11, 2024 • edited Loading

github-actions bot commented Oct 11, 2024

⚠ DCO check

kriscon-db commented Oct 13, 2024

harupy Oct 16, 2024

Choose a reason for hiding this comment

kriscon-db Oct 16, 2024

Choose a reason for hiding this comment

kriscon-db commented Oct 17, 2024

harupy Oct 17, 2024

Choose a reason for hiding this comment

kriscon-db Oct 17, 2024 • edited Loading

Choose a reason for hiding this comment

harupy Oct 17, 2024

Choose a reason for hiding this comment

artjen Oct 17, 2024

Choose a reason for hiding this comment

kriscon-db commented Oct 17, 2024

kriscon-db commented Oct 11, 2024 •

edited

Loading

github-actions bot commented Oct 11, 2024 •

edited

Loading

kriscon-db Oct 17, 2024 •

edited

Loading