[Endpoints] [2/x] Add cryptography implementation for secrets storage by BenWilson2 · Pull Request #19003 · mlflow/mlflow

BenWilson2 · 2025-11-24T17:57:47Z

🥞 Stacked PR

Use this link to review incremental changes.

stack/endpoints/crypto [Files changed]
- stack/endpoints/entities [Files changed]
  - stack/endpoints/abstract [Files changed]
    - stack/endpoints/sql-store [Files changed]
      - stack/endpoints/rest [Files changed]
        
        stack/endpoints/rest-2 [Files changed]
        
        stack/endpoints/cache [Files changed]
        
        stack/endpoints/litellm [Files changed]
        
        stack/endpoints/ui-apis [Files changed]
        
        stack/endpoints/ui-create-endpoint [Files changed]
        stack/endpoints/ui-tabs [Files changed]
        stack/endpoints/key-management [Files changed]
        stack/endpoints/passphrase [Files changed]

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

Adds the encryption / decryption layer for secrets management for endpoints

How is this PR tested?

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

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)

github-actions · 2025-11-24T18:29:22Z

Documentation preview for 1da3509 is available at:

https://pr-19003--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: Ben Wilson <benjamin.wilson@databricks.com>

Copilot

Pull request overview

This PR adds a comprehensive cryptography implementation for secure secrets storage in MLflow, including encryption/decryption using envelope encryption (KEK/DEK pattern) and a CLI tool for KEK rotation operations.

Key Changes:

Implements AES-256-GCM encryption with PBKDF2-derived KEK for secrets management
Adds envelope encryption pattern where each secret has a unique DEK wrapped by a master KEK
Provides CLI command mlflow crypto rotate-kek for secure KEK passphrase rotation

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
`mlflow/utils/cryptography.py`	Core cryptography module implementing KEK management, encryption/decryption functions, key wrapping/unwrapping, and secret masking utilities
`mlflow/cli/cryptography.py`	CLI commands for cryptographic operations including KEK rotation with database transaction support
`mlflow/cli/__init__.py`	Integrates cryptography CLI commands with optional import handling
`tests/utils/test_cryptography.py`	Comprehensive test suite covering encryption, decryption, key rotation, and edge cases
`tests/cli/test_cryptography.py`	CLI-specific tests including rotation workflows, error handling, and user interaction scenarios

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2025-12-04T02:00:16Z

mlflow/cli/cryptography.py

+@click.group("crypto", help="Commands for managing MLflow's cryptographic passphrase.")
+def commands():
+    """
+    MLflow cryptopgraphic management CLI. Allows for the management of the envelope


Spelling error: "cryptopgraphic" should be "cryptographic"

Suggested change

MLflow cryptopgraphic management CLI. Allows for the management of the envelope

MLflow cryptographic management CLI. Allows for the management of the envelope

mlflow/utils/cryptography.py

Copilot · 2025-12-04T02:00:17Z

mlflow/cli/cryptography.py

+            click.echo(
+                f"\n✓ Successfully rotated {rotated_count} encryption keys "


Grammar issue with pluralization: when rotated_count is 1, the message will say "Successfully rotated 1 encryption keys" which is grammatically incorrect. Consider using singular/plural form based on count:

key_word = "key" if rotated_count == 1 else "keys" click.echo( f"\n✓ Successfully rotated {rotated_count} encryption {key_word} " f"from KEK v{old_version} to v{new_version}\n" )

Suggested change

click.echo(

f"\n✓ Successfully rotated {rotated_count} encryption keys "

key_word = "key" if rotated_count == 1 else "keys"

click.echo(

f"\n✓ Successfully rotated {rotated_count} encryption {key_word} "

TomeHirata · 2025-12-04T12:01:44Z

mlflow/cli/cryptography.py

+        # Step 6: Restart server
+        $ systemctl start mlflow-server
+    """
+    old_passphrase = os.getenv("MLFLOW_CRYPTO_KEK_PASSPHRASE")


q: we don't define this env var in mlflow.environment_variables because it's not used in the tracking server logic?

It's not used client-side at all (that would be dangerous). This passphrase must only reside server-side :)

TomeHirata · 2025-12-04T12:07:52Z

mlflow/cli/cryptography.py

+                        click.echo(
+                            f"\n✗ Failed to rotate encryption key {secret.secret_id}: {e}", err=True
+                        )
+                        session.rollback()


Isn't rollback handled by ManagedSessionMaker?

TomeHirata · 2025-12-04T12:08:03Z

mlflow/cli/cryptography.py

+                            "No changes were made. Fix the issue and re-run the command."
+                        ) from e
+
+            session.commit()


TomeHirata · 2025-12-04T12:10:37Z

mlflow/utils/cryptography.py

+
+        return result
+
+    if not isinstance(secret_value, str):


What type falls into this branch? The type hint only expects str or dict.

TomeHirata · 2025-12-04T12:12:06Z

mlflow/utils/cryptography.py

+            return json.loads(plaintext)
+        except json.JSONDecodeError:
+            return plaintext
+


Suggested change

TomeHirata · 2025-12-04T12:12:26Z

mlflow/utils/cryptography.py

+        new_wrapped_dek = wrap_dek(dek, new_kek)
+
+        return RotatedSecret(encrypted_value=encrypted_value, wrapped_dek=new_wrapped_dek)
+


Suggested change

TomeHirata · 2025-12-04T12:13:13Z

mlflow/utils/cryptography.py

+    )
+
+
+def decrypt_secret(


Suggested change

def decrypt_secret(

def _decrypt_secret(

TomeHirata · 2025-12-04T12:13:25Z

mlflow/utils/cryptography.py

+    return f"{prefix}...{suffix}"
+
+
+def encrypt_secret(


Suggested change

def encrypt_secret(

def _encrypt_secret(

TomeHirata · 2025-12-04T12:13:39Z

mlflow/utils/cryptography.py

+    return aad_str.encode("utf-8")
+
+
+def mask_secret_value(secret_value: str | dict[str, Any]) -> str:


Suggested change

def mask_secret_value(secret_value: str | dict[str, Any]) -> str:

def _mask_secret_value(

TomeHirata · 2025-12-04T12:13:49Z

mlflow/utils/cryptography.py

+        ) from e
+
+
+def create_aad(secret_id: str, secret_name: str) -> bytes:


Suggested change

def create_aad(secret_id: str, secret_name: str) -> bytes:

def _create_aad(secret_id: str, secret_name: str) -> bytes:

TomeHirata · 2025-12-04T12:14:43Z

mlflow/utils/cryptography.py

+        return self._kek_version
+
+
+def generate_dek() -> bytes:


Shall we make this private? Same for other methods not used by other files

Suggested change

def generate_dek() -> bytes:

def _generate_dek() -> bytes:

TomeHirata · 2025-12-04T12:18:00Z

mlflow/utils/cryptography.py

+def _check_cryptography_available():
+    """Check if cryptography is installed and raise helpful error if not."""
+    try:
+        import cryptography  # noqa: F401


Can we use importlib.util.find_spec?

TomeHirata · 2025-12-04T12:21:57Z

Overall looks great, left some style comments

BenWilson2 changed the title ~~Add crypto capabilities~~ [Endpoints] [2/x] Add cryptography implementation for secrets storage Nov 24, 2025

BenWilson2 marked this pull request as ready for review November 24, 2025 18:19

BenWilson2 mentioned this pull request Nov 24, 2025

[Endpoints] [8/x] Add credential cache #19014

Merged

29 tasks

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

BenWilson2 force-pushed the stack/endpoints/crypto branch 2 times, most recently from 1775fa6 to bbb7755 Compare November 26, 2025 03:39

This was referenced Nov 26, 2025

[Endpoints] [11/x] Main Gateway page and create endpoint #19042

Closed

[Endpoints] [12/x] Add Left hand nav and management stub #19070

Closed

[Endpoints] [13/x] Add the API Key management pages #19071

Closed

BenWilson2 force-pushed the stack/endpoints/crypto branch from bbb7755 to edb2277 Compare November 27, 2025 06:31

BenWilson2 mentioned this pull request Nov 27, 2025

[Endpoints] [14/x] Models management UI #19074

Closed

29 tasks

BenWilson2 force-pushed the stack/endpoints/crypto branch 2 times, most recently from e663668 to 272cf2d Compare December 2, 2025 02:14

BenWilson2 added the team-review Trigger a team review request label Dec 2, 2025

github-actions bot requested review from B-Step62, TomeHirata, daniellok-db, harupy, kevin-lyn and serena-ruan December 2, 2025 16:03

github-actions bot requested review from WeichenXu123 and xq-yin December 2, 2025 16:03

BenWilson2 mentioned this pull request Dec 2, 2025

[Endpoints] [14/x] Add encryption availability endpoints #19176

Closed

29 tasks

BenWilson2 force-pushed the stack/endpoints/crypto branch 3 times, most recently from 489066d to 79cefae Compare December 3, 2025 22:35

Add crypto capabilities

1da3509

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

BenWilson2 force-pushed the stack/endpoints/crypto branch from 79cefae to 1da3509 Compare December 4, 2025 01:54

Copilot AI review requested due to automatic review settings December 4, 2025 01:54

Copilot started reviewing on behalf of BenWilson2 December 4, 2025 01:55 View session

Copilot finished reviewing on behalf of BenWilson2 December 4, 2025 01:57

Copilot AI reviewed Dec 4, 2025

View reviewed changes

TomeHirata reviewed Dec 4, 2025

View reviewed changes

mlflow/cli/cryptography.py

"No changes were made. Fix the issue and re-run the command."

) from e

session.commit()

Copy link

Collaborator

TomeHirata Dec 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ditto

TomeHirata reviewed Dec 4, 2025

View reviewed changes

mlflow/utils/cryptography.py

return json.loads(plaintext)

except json.JSONDecodeError:

return plaintext

Copy link

Collaborator

TomeHirata Dec 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change

TomeHirata reviewed Dec 4, 2025

View reviewed changes

mlflow/utils/cryptography.py

)

def decrypt_secret(

Copy link

Collaborator

TomeHirata Dec 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change

def decrypt_secret(

def _decrypt_secret(

TomeHirata reviewed Dec 4, 2025

View reviewed changes

BenWilson2 closed this Dec 4, 2025

	MLflow cryptopgraphic management CLI. Allows for the management of the envelope
	MLflow cryptographic management CLI. Allows for the management of the envelope

		click.echo(
		f"\n✓ Successfully rotated {rotated_count} encryption keys "

		new_wrapped_dek = wrap_dek(dek, new_kek)

		return RotatedSecret(encrypted_value=encrypted_value, wrapped_dek=new_wrapped_dek)

		return aad_str.encode("utf-8")


		def mask_secret_value(secret_value: str \| dict[str, Any]) -> str:

	def mask_secret_value(secret_value: str \| dict[str, Any]) -> str:
	def _mask_secret_value(

		) from e


		def create_aad(secret_id: str, secret_name: str) -> bytes:

	def create_aad(secret_id: str, secret_name: str) -> bytes:
	def _create_aad(secret_id: str, secret_name: str) -> bytes:

Conversation

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

Uh oh!

🥞 Stacked 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 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Dec 4, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Copilot AI Dec 4, 2025

Choose a reason for hiding this comment

Uh oh!

TomeHirata Dec 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

TomeHirata Dec 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

TomeHirata commented Dec 4, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

BenWilson2 commented Nov 24, 2025 •

edited

Loading

github-actions bot commented Nov 24, 2025 •

edited

Loading

TomeHirata Dec 4, 2025 •

edited

Loading

TomeHirata Dec 4, 2025 •

edited

Loading