Models in r2 rw support #9490

shichengzhou-db · 2023-08-30T22:11:43Z

Related Issues/PRs

N/A

What changes are proposed in this pull request?

Added a R2ArtifactStore class to support registering uc models to be stored in R2 Storage. R2ArtifactStore subclasses S3ArtifactStore to do most of the heavy lifting but overrides some methods to 1) parse (bucket, path) from r2 uri 2) instantiate s3 client with virtual addressing style

How is this patch tested?

Added unit test for R2ArtifactStore to specifically test for parsing of R2 URI and if S3 client is initialized correctly for R2 artifact store.
Manually created a R2 Storage backed UC Catalog and successfully registered a model into it (Version 3 in screenshot).

Existing unit/integration tests
New unit/integration tests
Manual tests (describe details, including test results, below)

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.

Support registering model in a R2 Storage backed Databricks UC Catalog.

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/breaking-change - The PR will be mentioned in the "Breaking Changes" section
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/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

github-actions · 2023-08-30T22:14:00Z

@shichengzhou-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.

mlflow-automation · 2023-08-30T22:14:18Z

Documentation preview for dc8c6ac will be available here 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/6032495621.

smurching

Looks great @shichengzhou-db ! Were you also able to test loading the model that you registered to R2 via e.g. mlflow.pyfunc.load_model?

smurching · 2023-08-30T22:22:24Z

@jerrylian-db FYI there were some gotchas we ran into (e.g. needing to use virtual host style addressing) to get reading/writing model artifacts to R2 to work using UC temporary credentials. You might want to work with Shicheng to test that R2 model read/write still work once you have a PR ready to refactor S3ArtifactRepo to share the fast upload/download logic in DatabricksArtifactRepo

smurching · 2023-08-30T22:23:36Z

mlflow/store/artifact/r2_artifact_repo.py

+        return bucket, path
+
+    def _get_s3_client(self, addressing_style="virtual"):
+        return super()._get_s3_client(addressing_style=addressing_style)


You may need to run black . to fix lint errors (e.g. having newlines at EOF)

Signed-off-by: shichengzhou-db <[email protected]>

shichengzhou-db · 2023-08-30T23:37:32Z

Looks great @shichengzhou-db ! Were you also able to test loading the model that you registered to R2 via e.g. mlflow.pyfunc.load_model?

yep loaded the model and did a prediction, works as expected (added this manual test code in the notebook shared with you earlier)

Signed-off-by: shichengzhou-db <[email protected]>

kriscon-db

LGTM

shichengzhou-db · 2023-08-31T02:15:51Z

@mlflow-automation autoformat

…port

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

harupy · 2023-08-31T03:16:58Z

mlflow/store/artifact/r2_artifact_repo.py

+        if path.startswith("/"):
+            path = path[1:]


Suggested change

if path.startswith("/"):

path = path[1:]

path = path.lstrip("/")

can we use lstrip here?

harupy · 2023-08-31T03:18:59Z

mlflow/store/artifact/r2_artifact_repo.py

+        # r2 uri format(virtual): r2://<bucket-name>@<account-id>.r2.cloudflarestorage.com/<path>
+        parsed = urllib.parse.urlparse(uri)
+        if parsed.scheme != "r2":
+            raise Exception(f"Not an R2 URI: {uri}")


Suggested change

raise Exception(f"Not an R2 URI: {uri}")

raise MlflowException.invalid_parameter_value(f"Not an R2 URI: {uri}")

harupy · 2023-08-31T03:27:47Z

mlflow/store/artifact/r2_artifact_repo.py

+    def __init__(
+        self, artifact_uri, access_key_id=None, secret_access_key=None, session_token=None
+    ):
+        super().__init__(
+            artifact_uri,
+            access_key_id=access_key_id,
+            secret_access_key=secret_access_key,
+            session_token=session_token,
+        )


Do we need to override __init__?

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

Signed-off-by: shichengzhou-db <[email protected]> Signed-off-by: mlflow-automation <[email protected]> Co-authored-by: mlflow-automation <[email protected]> Signed-off-by: Lu Peng <[email protected]>

shichengzhou-db requested review from kriscon-db and smurching August 30, 2023 22:13

github-actions bot added area/artifacts Artifact stores and artifact logging area/model-registry Model registry, model registry APIs, and the fluent client calls for model registry integrations/databricks Databricks integrations rn/feature Mention under Features in Changelogs. labels Aug 30, 2023

smurching approved these changes Aug 30, 2023

View reviewed changes

smurching reviewed Aug 30, 2023

View reviewed changes

Support r2 r/w model version

cf5b780

Signed-off-by: shichengzhou-db <[email protected]>

shichengzhou-db force-pushed the models_in_r2_rw_support branch from 6ff7637 to 28a81f9 Compare August 30, 2023 23:04

Support r2 r/w model version

b9c1b4e

Signed-off-by: shichengzhou-db <[email protected]>

shichengzhou-db force-pushed the models_in_r2_rw_support branch from 28a81f9 to b9c1b4e Compare August 30, 2023 23:43

kriscon-db approved these changes Aug 31, 2023

View reviewed changes

mlflow-automation added 2 commits August 31, 2023 02:53

Merge remote-tracking branch 'mlflow/master' into models_in_r2_rw_sup…

245d9cb

…port

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

dc8c6ac

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

harupy reviewed Aug 31, 2023

View reviewed changes

shichengzhou-db merged commit 419d0df into mlflow:master Aug 31, 2023
26 checks passed

jerrylian-db pushed a commit that referenced this pull request Sep 4, 2023

Models in r2 rw support (#9490)

e09c90e

Signed-off-by: shichengzhou-db <[email protected]> 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

Models in r2 rw support #9490

Models in r2 rw support #9490

shichengzhou-db commented Aug 30, 2023 •

edited

Loading

github-actions bot commented Aug 30, 2023

mlflow-automation commented Aug 30, 2023 •

edited

Loading

smurching left a comment

smurching commented Aug 30, 2023

smurching Aug 30, 2023

shichengzhou-db commented Aug 30, 2023

kriscon-db left a comment

shichengzhou-db commented Aug 31, 2023

harupy Aug 31, 2023 •

edited

Loading

harupy Aug 31, 2023

harupy Aug 31, 2023

	if path.startswith("/"):
	path = path[1:]
	path = path.lstrip("/")

	raise Exception(f"Not an R2 URI: {uri}")
	raise MlflowException.invalid_parameter_value(f"Not an R2 URI: {uri}")

Models in r2 rw support #9490

Models in r2 rw support #9490

Conversation

shichengzhou-db commented Aug 30, 2023 • edited Loading

Related Issues/PRs

What changes are proposed in this pull request?

How is this patch 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:

github-actions bot commented Aug 30, 2023

⚠ DCO check

mlflow-automation commented Aug 30, 2023 • edited Loading

smurching left a comment

Choose a reason for hiding this comment

smurching commented Aug 30, 2023

smurching Aug 30, 2023

Choose a reason for hiding this comment

shichengzhou-db commented Aug 30, 2023

kriscon-db left a comment

Choose a reason for hiding this comment

shichengzhou-db commented Aug 31, 2023

harupy Aug 31, 2023 • edited Loading

Choose a reason for hiding this comment

harupy Aug 31, 2023

Choose a reason for hiding this comment

harupy Aug 31, 2023

Choose a reason for hiding this comment

shichengzhou-db commented Aug 30, 2023 •

edited

Loading

mlflow-automation commented Aug 30, 2023 •

edited

Loading

harupy Aug 31, 2023 •

edited

Loading