WIP: Database migration CI check#7
Conversation
WalkthroughAdds workspace multi‑tenancy (DB schema, ORM, migration, tests), removes the Diviner flavor and docs/tests, introduces scorer permissions and telemetry for received traces, adds Vercel AI OTEL translator, repo-aware Python inference, a subprocess lint rule, many test/CI updates, and tracing GetTrace endpoints and store changes. (50 words) Changes
Sequence Diagram(s)sequenceDiagram
participant Client as OTLP client
participant API as mlflow OTEL API
participant Store as Tracking Store
participant Telemetry as Telemetry Recorder
Client->>API: POST /v1/traces (with User-Agent)
API->>API: process spans -> detect completed traces
API->>Store: persist spans (batch)
API->>Telemetry: _record_event(TracesReceivedByServerEvent(source, count))
Telemetry-->>API: acknowledged
API-->>Client: 200 OK
sequenceDiagram
actor User
participant CLI as mlflow server CLI
participant Utils as tracking utils
participant FS as local mlruns
participant DB as DEFAULT_TRACKING_URI
User->>CLI: start server (no backend-store-uri)
CLI->>Utils: _get_default_tracking_uri()
Utils->>FS: check mlruns/0/meta.yaml
alt mlruns exists with data
FS-->>Utils: found
Utils-->>CLI: use file-based mlruns
else
Utils-->>CLI: use DEFAULT_TRACKING_URI (sqlite:///mlflow.db)
end
CLI-->>User: server started with chosen backend_store_uri
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Areas needing extra attention:
Poem
Pre-merge checks and finishing touches❌ Failed checks (2 warnings, 1 inconclusive)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
mprahl
force-pushed
the
workspaces-db-model
branch
from
85861c2 to
10278ca
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (27)
.github/workflows/preview-comment.js (1)
132-151: LGTM! The conditional rendering improves UX.The implementation correctly branches on the number of changed pages, providing a collapsible section for longer lists and a simpler display for shorter ones. The logic is sound and the markdown/HTML formatting is appropriate for GitHub.
Consider extracting the threshold value as a named constant for better maintainability:
+const MAX_PAGES_BEFORE_COLLAPSE = 5; + function getCommentTemplate({ commitSha, workflowRunLink, docsWorkflowRunUrl, mainMessage, changedPages, }) { let changedPagesSection = ""; if (changedPages && changedPages.length > 0) { const pageLinks = changedPages.map(({ link, status }) => `- ${link} (${status})`).join("\n"); // Only collapse if there are more than 5 changed pages - if (changedPages.length > 5) { + if (changedPages.length > MAX_PAGES_BEFORE_COLLAPSE) { changedPagesSection = ` <details> <summary>Changed Pages (${changedPages.length})</summary> ${pageLinks} </details> `; } else { changedPagesSection = ` **Changed Pages (${changedPages.length})** ${pageLinks} `; } }tests/gateway/providers/test_bedrock.py (1)
108-168: Anthropic Bedrock fixtures correctly assertanthropic_version, consider de‑duplicating the literalThe added
anthropic_version: "bedrock-2023-05-31"fields in the three Anthropic fixtures line up with the adapter behavior and make the_requestexpectations precise. To avoid future drift between tests andAmazonBedrockAnthropicAdapter, consider extracting this version into a shared constant inmlflow/gateway/providers/bedrock.pyand referencing it here instead of repeating the string literal in multiple places.mlflow/gateway/providers/bedrock.py (2)
19-26: Chat adapter now stampsanthropic_version; consider using a shared constant / defaultingAdding
payload["anthropic_version"] = "bedrock-2023-05-31"inchat_to_modelmakes the Bedrock Anthropic chat payloads explicit and consistent with completions. To reduce duplication (and keep tests in sync), you might want to promote this string to a module-level constant and reference it here. If you ever plan to allow callers to choose a different version, usingpayload.setdefault("anthropic_version", <CONST>)would avoid overwriting a caller-specified value.
29-43: Completions adapteranthropic_versionaddition aligns with tests; same constant/default considerationsThe new
payload["anthropic_version"] = "bedrock-2023-05-31"incompletions_to_model, alongside the existingstop_sequencesandmax_tokens_to_samplenormalization, matches whattests/gateway/providers/test_bedrock.pynow expects for Anthropic models. As above, consider extracting the version into a single constant used by both adapter methods and the tests, and optionally switching tosetdefaultif preserving a user-provided version becomes important later.mlflow/tracking/fluent.py (1)
407-410: Error message improved but could be clearer.The change from "active run ID" to "active experiment ID" is a good improvement. However, the phrase "environment run ID" at the end might still be confusing since the actual comparison is between two experiment IDs (the active experiment ID set via
set_experiment()vs. the experiment ID of the run specified by theMLFLOW_RUN_IDenvironment variable).Consider rewording for even better clarity:
- f"Cannot start run with ID {existing_run_id} because active experiment ID " - "does not match environment run ID. Make sure --experiment-name " + f"Cannot start run with ID {existing_run_id} because active experiment ID " + "does not match the experiment ID of the run from environment. Make sure --experiment-name "tests/crewai/test_crewai_autolog.py (2)
118-118: Consider the trade-off of usingANYfortasks_output.Replacing the explicit task output structure with
ANYmakes the test more version-tolerant but significantly reduces precision. This could mask regressions in task output format.If the task output structure varies legitimately across CrewAI versions, consider version-conditional assertions instead of a blanket
ANY.
268-268: Significant reduction in test assertion precision.Multiple span output assertions have been relaxed from checking exact dictionary structures to simply verifying
is not None. While this improves test stability across CrewAI versions, it substantially reduces test coverage.Consider:
- Documenting why exact assertions aren't feasible across versions
- Implementing version-conditional assertions for critical output fields where possible
- Adding complementary integration tests that validate end-to-end output quality
This pattern appears intentional for version tolerance, but be aware it may not catch output structure regressions.
Also applies to: 396-396, 463-463, 473-473, 515-515, 594-594, 718-718, 807-807, 886-886
tests/semantic_kernel/test_semantic_kernel_autolog.py (5)
103-114: LLM span name assertion now supports both naming schemesAllowing
spans[3].nameto be either"chat.completions gpt-4o-mini"or"chat gpt-4o-mini"keeps this test stable across instrumentation variants while still tightly checking the actual LLM span.To avoid repeating this tuple across the file, consider extracting a small constant and using it here and in the other tests that assert the same names, e.g.:
CHAT_SPAN_NAME_CANDIDATES = ("chat.completions gpt-4o-mini", "chat gpt-4o-mini")- assert spans[3].name in ("chat.completions gpt-4o-mini", "chat gpt-4o-mini") + assert spans[3].name in CHAT_SPAN_NAME_CANDIDATES
319-331: Active-span test shares the same chat span name tupleThis assertion mirrors the earlier one in
test_sk_invoke_simple, again allowing both"chat.completions gpt-4o-mini"and"chat gpt-4o-mini".If you introduce a shared
CHAT_SPAN_NAME_CANDIDATESconstant as suggested earlier, it would be nice to use it here too to keep these tests in sync:- assert spans[4].name in ("chat.completions gpt-4o-mini", "chat gpt-4o-mini") + assert spans[4].name in CHAT_SPAN_NAME_CANDIDATES
376-420: Substring-basedspan_name_patternworks but could become ambiguous laterSwitching the parametrization to
"chat"/"text"and then matching with:for span in trace.data.spans: if span_name_pattern in span.name: target_span = span breakmakes the test tolerant to changes in the full span name, which is good. The only caveat is future ambiguity if additional spans’ names also contain
"chat"or"text".To future‑proof this a bit, you could optionally narrow candidates using
span.span_type(orSpanAttributeKey.SPAN_TYPE) so you always pick the LLM/chat span even if multiple names match, e.g.:from mlflow.entities import SpanType for span in trace.data.spans: if span_name_pattern == "chat" and span.span_type == SpanType.CHAT_MODEL: target_span = span break if span_name_pattern == "text" and span.span_type == SpanType.LLM: target_span = span breakNot required right now, but would reduce the chance of flakiness as instrumentation evolves.
432-435: Selecting the chat span via"chat" in s.nameis fine; consider using span_typeUsing:
child_span = next(s for s in traces[0].data.spans if "chat" in s.name)is a pragmatic way to handle multiple possible chat span names here, and it’s safe given the current span set. To make it more robust and self‑documenting, you could instead key off
SpanType.CHAT_MODEL(imported at the top) which directly represents the semantic role:- child_span = next(s for s in traces[0].data.spans if "chat" in s.name) + child_span = next( + s for s in traces[0].data.spans if s.span_type == SpanType.CHAT_MODEL + )This aligns the test with how spans are typed in
mlflow.entities.span.SpanTypeand avoids relying on name substrings.
468-478: Kernel function-object chat span assertion duplicates allowed name tupleAs in the other tests, accepting both
"chat.completions gpt-4o-mini"and"chat gpt-4o-mini"is a good way to accommodate instrumentation differences while still asserting on the correct child span.If you decide to introduce a shared
CHAT_SPAN_NAME_CANDIDATESconstant, this assertion can be simplified too:- assert chat_span.name in ("chat.completions gpt-4o-mini", "chat gpt-4o-mini") + assert chat_span.name in CHAT_SPAN_NAME_CANDIDATESdocs/docs/classic-ml/traditional-ml/prophet/index.mdx (2)
47-98: Add installation and setup instructions.The documentation jumps directly to code examples without explaining how to install Prophet or enable the MLflow Prophet integration. Consider adding a setup section before "Basic Model Logging" that covers installation steps and any required configuration.
Suggested addition (before line 47):
## Installation Before using MLflow with Prophet, install both libraries: \`\`\`bash pip install prophet mlflow \`\`\` Ensure you have the MLflow Prophet flavor available: \`\`\`bash pip install mlflow[prophet] \`\`\`
152-153: Add missing closing section.The code block starting at line 152 (Optuna import section) appears to have an extra blank line (line 153). While minor, ensure this doesn't break the code example rendering.
This is a nitpick: consider removing the extra blank line at line 153 for cleaner code formatting:
import optuna - def objective(trial, df):dev/set_matrix.py (2)
368-370: Consider using a proper TOML parser.The regex-based TOML parsing is fragile and may fail on valid TOML files with:
- Multiline strings
- Escaped quotes in values
- Comments within strings
- Alternative quote styles
Consider using the
tomlilibrary (standard in Python 3.11+) ortoml:+import tomli + def _requires_python_from_repo(repo_url: str) -> str | None: ... try: resp = requests.get(raw_url, timeout=10) resp.raise_for_status() except requests.HTTPError as e: if e.response.status_code == 404: print(f" pyproject.toml not found at {raw_url}", file=sys.stderr) return None raise - if match := re.search(r'requires-python\s*=\s*["\']([^"\']+)["\']', resp.text): - print(f" Found requires-python: {match.group(1)}", file=sys.stderr) - return match.group(1) + try: + data = tomli.loads(resp.text) + if requires_python := data.get("project", {}).get("requires-python"): + print(f" Found requires-python: {requires_python}", file=sys.stderr) + return requires_python + except tomli.TOMLDecodeError: + print(" Failed to parse pyproject.toml", file=sys.stderr) + return None print(" requires-python field not found in pyproject.toml", file=sys.stderr) return None
359-366: Note: Limited exception handling for network errors.The function only catches
HTTPError. Other exceptions likeTimeout,ConnectionError, orRequestExceptionwill propagate uncaught. This may be intentional for CI failure-fast behavior, but consider whether graceful degradation (returningNone) would be more appropriate.If graceful degradation is preferred:
try: resp = requests.get(raw_url, timeout=10) resp.raise_for_status() - except requests.HTTPError as e: + except requests.RequestException as e: - if e.response.status_code == 404: + if isinstance(e, requests.HTTPError) and e.response.status_code == 404: print(f" pyproject.toml not found at {raw_url}", file=sys.stderr) - return None + else: + print(f" Network error fetching pyproject.toml: {e}", file=sys.stderr) - raise + return Nonemlflow/ml-package-versions.yml (1)
1-1067: NewrepoURLs for dev installs are consistent with existing install_dev commandsAdding
package_info.repofor these packages (both non‑GenAI and GenAI) aligns with the existinginstall_devcommands and uses the same GitHub org/repo paths and subdirectories. That should make any tooling that surfaces source locations or links much more useful with minimal risk.One thing to keep in mind: using
/tree/HEADmeans these links always track the default branch, which is great for discoverability but less reproducible if upstream restructures directories. If you ever need stronger determinism for CI tooling that consumesrepo, consider pinning to a specific branch, tag, or commit instead.docs/docs/classic-ml/traditional-ml/sparkml/index.mdx (1)
162-166: Consider adding a version requirement note.The
array_to_vectorfunction was introduced in Spark 3.1. Consider adding a brief note about the Spark version requirement for this example.Add a note similar to the datasource autologging section:
```python from mlflow.models import infer_signature from pyspark.ml.functions import array_to_vector +# Note: array_to_vector requires Spark 3.1+ vector_data = spark.createDataFrame(docs/docusaurus.config.ts (1)
800-806: Consider consolidating TensorFlow redirects.For consistency with the PyTorch redirect pattern (lines 757-764), consider combining these two TensorFlow redirect entries into a single entry with both from-paths.
Apply this diff to consolidate the redirects:
- { - to: '/ml/deep-learning/tensorflow', - from: ['/deep-learning/tensorflow/guide'], - }, - { - to: '/ml/deep-learning/tensorflow', - from: ['/deep-learning/tensorflow/quickstart/quickstart_tensorflow'], - }, + { + to: '/ml/deep-learning/tensorflow', + from: [ + '/deep-learning/tensorflow/guide', + '/ml/deep-learning/tensorflow/guide', + '/deep-learning/tensorflow/quickstart/quickstart_tensorflow', + '/ml/deep-learning/tensorflow/quickstart/quickstart-tensorflow', + ], + },tests/tracing/otel/test_span_translation.py (1)
19-75: Vercel span-kind translation coverage looks correctThe added cases for
VercelAITranslator(ai.generateText→SpanType.LLM,ai.toolCall→SpanType.TOOL) align withSPAN_KIND_TO_MLFLOW_TYPEand follow the existing JSON‑serialized pattern in this test. Looks good as-is; you could later add an embedding case if you want full SPAN_KIND coverage, but it’s not required here.mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
20-223: Upgrade path and dialect-specific unique/FK handling looks solidThe upgrade logic correctly:
- Normalizes unique detection across dialects (including MSSQL via
_fetch_mssql_unique_metadata) and safely drops pre‑existing uniqueness onnamebefore introducing workspace-aware keys.- Handles SQLite vs non‑SQLite by pre‑dropping FKs only on non‑SQLite and letting
_with_batchtake care of FK drops on SQLite, avoiding double-drops.- Adds
workspacewith a consistentserver_default 'default'and updates PKs/FKs to includeworkspace, matching the intended multi-tenant schema.One optional improvement would be to hoist the
"default"workspace name into a module-level constant and reuse it for both the column defaults and the seed insert to avoid drift if the default workspace name ever changes.tests/db/test_schema.py (1)
2-107: Dialect-aware unique-constraint reattachment is appropriate for schema dumpsThe new
_reattach_missing_unique_constraints+_get_unique_constraintslogic cleanly normalizes unique-constraint reflection quirks across MySQL and MSSQL before generating schema strings. Limiting this to a small, explicit set of constraint names and guarding withduplicates_indexfor MySQL keeps the behavior safe and targeted, and the warning on missing columns is a good diagnostic. Overall this is a solid way to stabilizedump_schemaacross dialects.mlflow/tracing/otel/translation/vercel_ai.py (1)
12-32: Translator behavior looks correct; consider a small note on JSON depth handlingThe span kind/type mapping and chat vs non‑chat branching in
get_input_value/get_output_valuelook consistent with the Vercel AI schema and the baseOtelSchemaTranslatorcontract. The_unpack_attributes_with_prefixandtoolshandling should give you a clean, JSON‑encoded aggregate ofai.prompt.*/ai.response.*attributes for chat spans, while deferring to the base logic for non‑chat spans.One minor point:
_safe_load_jsononly recurses while the intermediate value is still a string; once it becomes a dict/list, it returns it as‑is. That’s fine for handling double‑encoded JSON strings (your main goal) but it won’t walk nested dicts/lists to decode any inner stringified JSON. If you ever need deeper normalization of nested structures, this is the place to extend, but it’s not required for correctness as written.Also applies to: 34-48, 50-62, 64-88
mlflow/store/model_registry/dbmodels/models.py (2)
1-1: Simplifyserver_defaultfor workspace columnsFor string defaults you can rely on SQLAlchemy’s own quoting/escaping and drop the
sa.text(f"'{DEFAULT_WORKSPACE_NAME}'")wrapper, which also lets you avoid the extrasqlalchemy as saimport.- workspace = Column( - String(63), - nullable=False, - server_default=sa.text(f"'{DEFAULT_WORKSPACE_NAME}'"), - ) + workspace = Column( + String(63), + nullable=False, + server_default=DEFAULT_WORKSPACE_NAME, + )(Same pattern applies to the other workspace columns in this file.)
Also applies to: 42-46, 87-91, 156-160, 193-197, 232-236, 296-300
173-180: Consider reordering tag PK columns to match likely access patternsBoth tag tables currently use PKs prefixed by
(workspace, key, ...); given the pre-workspace keys were effectively(name[, version], key), it’s probably more efficient and intuitive to keep(workspace, name[, version], key)so the index prefix matches typical lookups by model/version identity.- PrimaryKeyConstraint("workspace", "key", "name", name="registered_model_tag_pk"), + PrimaryKeyConstraint("workspace", "name", "key", name="registered_model_tag_pk"),- PrimaryKeyConstraint("workspace", "key", "name", "version", name="model_version_tag_pk"), + PrimaryKeyConstraint("workspace", "name", "version", "key", name="model_version_tag_pk"),If you make this change, remember to update the migrations and test schemas (e.g., sqlite.sql) to match.
Also applies to: 207-214
tests/db/schemas/sqlite.sql (2)
106-110: Optional: add FKs from workspace columns toworkspacesRight now
workspacesis a standalone catalog; if you want stronger integrity guarantees, consider adding foreign keys from columns likeexperiments.workspaceandregistered_models.workspacetoworkspaces(name)to prevent dangling workspace names. If intentional, no change needed.
218-221: Revisit PK column order for tag tables in sqlite schemaThe PKs for
registered_model_tagsandmodel_version_tagsare(workspace, key, name[, version]). To better preserve the original(name[, version], key)ordering and index prefix, consider switching to(workspace, name[, version], key), and keep this in sync with the ORM definitions.- CONSTRAINT registered_model_tag_pk PRIMARY KEY (workspace, key, name), + CONSTRAINT registered_model_tag_pk PRIMARY KEY (workspace, name, key),- CONSTRAINT model_version_tag_pk PRIMARY KEY (workspace, key, name, version), + CONSTRAINT model_version_tag_pk PRIMARY KEY (workspace, name, version, key),Also applies to: 372-375
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (3)
docs/static/images/traditional-ml-ui.gifis excluded by!**/*.gifmlflow/server/js/nohup.outis excluded by!**/*.outuv.lockis excluded by!**/*.lock
📒 Files selected for processing (86)
.github/actions/setup-python/action.yml(1 hunks).github/workflows/master.yml(0 hunks).github/workflows/preview-comment.js(2 hunks)dev/clint/src/clint/linter.py(1 hunks)dev/clint/src/clint/rules/__init__.py(2 hunks)dev/clint/src/clint/rules/subprocess_check_call.py(1 hunks)dev/clint/tests/rules/test_subprocess_check_call.py(1 hunks)dev/create_release_tag.py(1 hunks)dev/set_matrix.py(5 hunks)docs/api_reference/api_inventory.txt(0 hunks)docs/api_reference/source/python_api/mlflow.diviner.rst(0 hunks)docs/docs/classic-ml/deep-learning/pytorch/guide/index.mdx(0 hunks)docs/docs/classic-ml/deep-learning/pytorch/index.mdx(0 hunks)docs/docs/classic-ml/deep-learning/pytorch/quickstart/quickstart-pytorch.ipynb(0 hunks)docs/docs/classic-ml/model/index.mdx(0 hunks)docs/docs/classic-ml/traditional-ml/index.mdx(2 hunks)docs/docs/classic-ml/traditional-ml/prophet/guide/index.mdx(0 hunks)docs/docs/classic-ml/traditional-ml/prophet/index.mdx(2 hunks)docs/docs/classic-ml/traditional-ml/sparkml/guide/index.mdx(0 hunks)docs/docs/classic-ml/traditional-ml/sparkml/index.mdx(1 hunks)docs/docs/classic-ml/tutorials-and-examples/index.mdx(0 hunks)docs/docs/genai/tracing/integrations/listing/langchain.mdx(6 hunks)docs/docs/genai/tracing/integrations/listing/langgraph.mdx(6 hunks)docs/docusaurus.config.ts(3 hunks)docs/sidebarsClassicML.ts(2 hunks)examples/diviner/MLproject(0 hunks)examples/diviner/python_env.yaml(0 hunks)examples/diviner/train.py(0 hunks)mlflow/__init__.py(0 hunks)mlflow/diviner/__init__.py(0 hunks)mlflow/entities/__init__.py(2 hunks)mlflow/entities/workspace.py(1 hunks)mlflow/gateway/providers/bedrock.py(2 hunks)mlflow/mcp/server.py(1 hunks)mlflow/ml-package-versions.yml(30 hunks)mlflow/ml_package_versions.py(1 hunks)mlflow/models/__init__.py(0 hunks)mlflow/semantic_kernel/tracing_utils.py(1 hunks)mlflow/server/js/src/common/constants.tsx(1 hunks)mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py(1 hunks)mlflow/store/model_registry/dbmodels/models.py(11 hunks)mlflow/store/tracking/dbmodels/models.py(5 hunks)mlflow/store/workspace/dbmodels/__init__.py(1 hunks)mlflow/store/workspace/dbmodels/models.py(1 hunks)mlflow/telemetry/constant.py(0 hunks)mlflow/tracing/otel/translation/__init__.py(2 hunks)mlflow/tracing/otel/translation/vercel_ai.py(1 hunks)mlflow/tracking/fluent.py(1 hunks)mlflow/types/utils.py(0 hunks)mlflow/utils/workspace_utils.py(1 hunks)pyproject.toml(0 hunks)requirements/constraints.txt(0 hunks)requirements/extra-ml-requirements.txt(0 hunks)tests/check_mlflow_lazily_imports_ml_packages.py(0 hunks)tests/conftest.py(2 hunks)tests/crewai/test_crewai_autolog.py(16 hunks)tests/db/check_migration.py(3 hunks)tests/db/schemas/mssql.sql(6 hunks)tests/db/schemas/mysql.sql(6 hunks)tests/db/schemas/postgresql.sql(6 hunks)tests/db/schemas/sqlite.sql(6 hunks)tests/db/test_schema.py(3 hunks)tests/db/test_workspace_migration.py(1 hunks)tests/diviner/test_diviner_model_export.py(0 hunks)tests/examples/test_examples.py(0 hunks)tests/gateway/providers/test_bedrock.py(3 hunks)tests/mcp/test_mcp.py(2 hunks)tests/pyfunc/docker/conftest.py(1 hunks)tests/pyfunc/docker/test_docker_flavors.py(0 hunks)tests/resources/example_mlflow_1x_sklearn_model/MLmodel(1 hunks)tests/resources/example_mlflow_1x_sklearn_model/conda.yaml(1 hunks)tests/resources/example_mlflow_1x_sklearn_model/python_env.yaml(1 hunks)tests/resources/example_virtualenv_conda_project/conda.yaml(1 hunks)tests/resources/example_virtualenv_conda_project/entrypoint.py(1 hunks)tests/resources/example_virtualenv_no_python_env/entrypoint.py(1 hunks)tests/resources/example_virtualenv_no_python_env/python_env.yaml(1 hunks)tests/resources/example_virtualenv_project/entrypoint.py(1 hunks)tests/resources/example_virtualenv_project/python_env.yaml(1 hunks)tests/semantic_kernel/test_semantic_kernel_autolog.py(7 hunks)tests/sklearn/test_sklearn_model_export.py(2 hunks)tests/tracing/conftest.py(1 hunks)tests/tracing/otel/test_span_translation.py(2 hunks)tests/tracing/otel/test_vercel_ai_translator.py(1 hunks)tests/tracking/fluent/test_fluent.py(1 hunks)tests/utils/test_request_utils.py(1 hunks)tests/webhooks/test_e2e.py(1 hunks)
💤 Files with no reviewable changes (25)
- .github/workflows/master.yml
- mlflow/types/utils.py
- examples/diviner/MLproject
- docs/docs/classic-ml/deep-learning/pytorch/index.mdx
- mlflow/models/init.py
- mlflow/telemetry/constant.py
- requirements/extra-ml-requirements.txt
- docs/api_reference/api_inventory.txt
- docs/docs/classic-ml/tutorials-and-examples/index.mdx
- examples/diviner/train.py
- docs/api_reference/source/python_api/mlflow.diviner.rst
- docs/docs/classic-ml/traditional-ml/prophet/guide/index.mdx
- pyproject.toml
- docs/docs/classic-ml/traditional-ml/sparkml/guide/index.mdx
- tests/diviner/test_diviner_model_export.py
- docs/docs/classic-ml/model/index.mdx
- tests/examples/test_examples.py
- docs/docs/classic-ml/deep-learning/pytorch/guide/index.mdx
- mlflow/init.py
- docs/docs/classic-ml/deep-learning/pytorch/quickstart/quickstart-pytorch.ipynb
- tests/pyfunc/docker/test_docker_flavors.py
- tests/check_mlflow_lazily_imports_ml_packages.py
- mlflow/diviner/init.py
- requirements/constraints.txt
- examples/diviner/python_env.yaml
🧰 Additional context used
🧬 Code graph analysis (18)
dev/clint/src/clint/rules/__init__.py (1)
dev/clint/src/clint/rules/subprocess_check_call.py (1)
SubprocessCheckCall(7-43)
mlflow/store/workspace/dbmodels/__init__.py (1)
mlflow/store/workspace/dbmodels/models.py (1)
SqlWorkspace(10-22)
mlflow/semantic_kernel/tracing_utils.py (1)
mlflow/entities/span.py (1)
SpanType(44-63)
mlflow/tracing/otel/translation/__init__.py (1)
mlflow/tracing/otel/translation/vercel_ai.py (1)
VercelAITranslator(8-88)
mlflow/entities/__init__.py (1)
mlflow/entities/workspace.py (1)
Workspace(9-13)
mlflow/tracing/otel/translation/vercel_ai.py (2)
mlflow/entities/span.py (1)
SpanType(44-63)mlflow/tracing/otel/translation/base.py (1)
OtelSchemaTranslator(15-171)
dev/clint/src/clint/rules/subprocess_check_call.py (2)
dev/clint/src/clint/resolver.py (1)
Resolver(6-89)dev/clint/src/clint/rules/base.py (1)
Rule(11-31)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
tests/db/test_schema.py (1)
_get_unique_constraints(90-107)
tests/tracking/fluent/test_fluent.py (1)
mlflow/exceptions.py (1)
MlflowException(67-113)
tests/tracing/otel/test_span_translation.py (2)
mlflow/tracing/otel/translation/vercel_ai.py (1)
VercelAITranslator(8-88)mlflow/entities/span.py (1)
SpanType(44-63)
tests/db/test_schema.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
_get_unique_constraints(144-159)
tests/tracing/otel/test_vercel_ai_translator.py (1)
mlflow/tracing/otel/translation/__init__.py (1)
translate_span_when_storing(35-71)
tests/semantic_kernel/test_semantic_kernel_autolog.py (2)
mlflow/entities/span.py (9)
name(131-133)name(796-797)parent_id(146-150)parent_id(812-813)span_id(126-128)span_id(792-793)span_type(168-170)SpanType(44-63)get_attribute(215-225)tests/semantic_kernel/resources.py (2)
_create_and_invoke_text_completion(95-108)_create_and_invoke_chat_completion_direct(123-142)
tests/crewai/test_crewai_autolog.py (2)
dev/set_matrix.py (1)
Version(57-79)mlflow/entities/span.py (1)
outputs(163-165)
mlflow/store/workspace/dbmodels/models.py (1)
mlflow/entities/workspace.py (1)
Workspace(9-13)
tests/db/test_workspace_migration.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (2)
upgrade(138-421)downgrade(424-678)
dev/clint/src/clint/linter.py (1)
dev/clint/src/clint/rules/subprocess_check_call.py (2)
SubprocessCheckCall(7-43)check(15-43)
dev/clint/tests/rules/test_subprocess_check_call.py (3)
dev/clint/src/clint/linter.py (3)
Position(52-59)Range(62-95)lint_file(918-959)dev/clint/src/clint/rules/subprocess_check_call.py (1)
SubprocessCheckCall(7-43)dev/clint/tests/rules/conftest.py (1)
index_path(8-12)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (30)
- GitHub Check: windows (4)
- GitHub Check: windows (2)
- GitHub Check: windows (3)
- GitHub Check: windows (1)
- GitHub Check: pyfunc (3)
- GitHub Check: pyfunc (4)
- GitHub Check: models (2)
- GitHub Check: pyfunc (2)
- GitHub Check: models (1)
- GitHub Check: pyfunc (1)
- GitHub Check: evaluate (2)
- GitHub Check: database
- GitHub Check: evaluate (1)
- GitHub Check: java
- GitHub Check: genai
- GitHub Check: flavors
- GitHub Check: python (2)
- GitHub Check: python (1)
- GitHub Check: python (3)
- GitHub Check: python-skinny
- GitHub Check: build (tracing)
- GitHub Check: build (dev)
- GitHub Check: build (skinny)
- GitHub Check: r
- GitHub Check: build
- GitHub Check: core
- GitHub Check: js (ubuntu-latest, --testPathPattern)
- GitHub Check: js (windows-latest, --testPathIgnorePatterns)
- GitHub Check: js (ubuntu-latest, --testPathIgnorePatterns)
- GitHub Check: js (windows-latest, --testPathPattern)
| # Cross-validation | ||
| cv_results = cross_validation( | ||
| model, initial="730 days", period="180 days", horizon="365 days" | ||
| model, | ||
| initial="730 days", | ||
| period="180 days", | ||
| horizon="365 days", | ||
| ) | ||
|
|
||
| # Log performance metrics | ||
| metrics = performance_metrics(cv_results) | ||
| avg_metrics = metrics[["mse", "rmse", "mae", "mape"]].mean().to_dict() | ||
| mlflow.log_metrics(avg_metrics) | ||
| mlflow.log_metrics(metrics[["mse", "rmse", "mae", "mape"]].mean().to_dict()) | ||
|
|
||
| # Log the model | ||
| # Log model | ||
| mlflow.prophet.log_model( | ||
| pr_model=model, name="prophet_model", input_example=df[["ds"]].head() | ||
| ) | ||
| ``` |
There was a problem hiding this comment.
🧩 Analysis chain
Clarify required dependencies for code examples.
The Basic Model Logging example imports cross_validation and performance_metrics from prophet.diagnostics and uses mlflow.prophet.log_model(). Ensure the documentation mentions these required packages and that the metric column names ("mse", "rmse", "mae", "mape") are accurate for the current Prophet version.
🌐 Web query:
What are the output column names from Prophet's performance_metrics function in recent versions?
💡 Result:
Typical output columns from Prophet's performance_metrics (recent versions) are:
- horizon
- mse
- rmse
- mae
- mape
- mdape
- smape
- coverage
(Which metrics appear can depend on the package/language and the metrics argument; by default the set includes mse, rmse, mae, mape, mdape, smape, and coverage.) [1][2][3]
Need to verify the remaining APIs used in the code example.
Code example APIs verified as correct; add dependency clarification to documentation.
The code correctly uses cross_validation from prophet.diagnostics with parameters initial, period, and horizon. The performance_metrics utility computes mean squared error (MSE), root mean squared error (RMSE), mean absolute error (MAE), mean absolute percent error (MAPE) along with additional metrics, so the four metric names selected are valid. mlflow.prophet.log_model() correctly accepts pr_model and input_example parameters.
The example works as written. To improve clarity, explicitly mention that users need the prophet and mlflow packages for this example (the basic Prophet import example also implicitly requires pandas).
🤖 Prompt for AI Agents
In docs/docs/classic-ml/traditional-ml/prophet/index.mdx around lines 82–98, the
code example works but lacks a dependency note; add a short sentence before or
after the code block stating that the example requires the prophet and mlflow
Python packages (and pandas for DataFrame handling), e.g., mentioning pip
install prophet mlflow pandas or equivalent, so readers know which libraries to
install before running the snippet.
| @@ -1,205 +1,276 @@ | |||
| import Link from "@docusaurus/Link"; | |||
| import { CardGroup, PageCard } from "@site/src/components/Card"; | |||
| import { APILink } from "@site/src/components/APILink"; | |||
There was a problem hiding this comment.
🧩 Analysis chain
Verify if APILink import is used.
The APILink component is imported but doesn't appear to be used in the document body.
Run the following script to check if this import is used:
🏁 Script executed:
#!/bin/bash
# Check if APILink is used in the SparkML index file
rg -n "APILink" docs/docs/classic-ml/traditional-ml/sparkml/index.mdxLength of output: 57
Remove unused APILink import.
The APILink component is imported on line 1 but is not used anywhere in the document. Remove the unused import.
🤖 Prompt for AI Agents
In docs/docs/classic-ml/traditional-ml/sparkml/index.mdx around line 1, the
APILink component is imported but never used; remove the unused import line
`import { APILink } from "@site/src/components/APILink";` from the top of the
file to clean up the code.
| ```typescript | ||
| import { OpenAI } from "@langchain/openai"; | ||
| import { PromptTemplate } from "@langchain/core/prompts"; | ||
|
|
||
| const model = new OpenAI("gpt-4o-mini"); | ||
| const prompt = PromptTemplate.fromTemplate("What is a good name for a company that makes {product}?"); | ||
| const chain = prompt.pipe({ llm: model }); | ||
|
|
||
| const res = await chain.invoke({ product: "colorful socks" }); | ||
| console.log({ res }); | ||
| ``` |
There was a problem hiding this comment.
🧩 Analysis chain
Incorrect TypeScript v0 pipe syntax—likely to fail when users reproduce the example.
Line 236 uses prompt.pipe({ llm: model }), but the correct LangChain.js syntax (v0 and v1) is prompt.pipe(model). The object-wrapped syntax does not match documented patterns and will cause errors at runtime.
Apply this diff to fix the chain definition:
-const chain = prompt.pipe({ llm: model });
+const chain = prompt.pipe(model);Additionally, Line 234 uses a positional argument for the model instantiation. Ensure this matches the OpenAI constructor API for LangChain 0.x. If the current syntax is valid, consider adding a brief comment for clarity.
🌐 Web query:
LangChain.js pipe syntax prompt.pipe model chainable
💡 Result:
Short answer: LangChain.js uses the Runnable/LCEL pipe API so prompts, models, parsers and other Runables are chainable with .pipe(). The output of each runnable becomes the input of the next, and the resulting RunnableSequence can be invoked, streamed, or further chained.
Example (typical pattern):
- create a prompt runnable
- pipe it into a chat model runnable
- pipe into an output parser (optional)
- invoke/stream the final chain
Code:
import { ChatPromptTemplate } from "@langchain/core/prompts";
import { ChatOpenAI } from "@langchain/openai";
import { StringOutputParser } from "@langchain/core/output_parsers";
const prompt = ChatPromptTemplate.fromTemplate("Tell me a joke about {topic}");
const model = new ChatOpenAI({});
const chain = prompt.pipe(model).pipe(new StringOutputParser());
const result = await chain.invoke({ topic: "bears" });
Notes:
- You can bind model kwargs with .bind() before piping (e.g., stop sequences or function schemas). [1][3]
- Chains support streaming (.stream()) and further composition (RunnablePick, RunnableLambda, RunnablePassthrough, etc.). [1][2]
- This is part of LangChain’s Expression Language / Runnable API; docs show examples across integrations. [1][4]
Sources:
[1] LangChain docs — “How to chain runnables / The pipe method”.
[2] LangChain docs — “Get started / Prompt + LLM” examples.
[3] LangChain docs — “Prompt + LLM / Attaching kwargs with .bind()”.
[4] @langchain/core API examples (pipe + stream).
🌐 Web query:
LangChain OpenAI constructor positional arguments v0
💡 Result:
Short answer: In LangChain v0 the OpenAI/ChatOpenAI constructors expect keyword arguments (e.g., model/model_name, temperature, max_tokens, api_key/base_url/organization, timeout, max_retries, stream_options, etc.) rather than requiring a fixed set of positional parameters. Example instantiation:
from langchain_openai import ChatOpenAI
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.0, max_retries=2, api_key="...")
Docs / reference for the constructor parameters: [1], [2].
Sources:
[1] LangChain reference — ChatOpenAI / OpenAI constructor params.
[2] LangChain v0 API docs — OpenAI llm constructor details.
🌐 Web query:
LangChain.js pipe object parameter llm model syntax
💡 Result:
Short answer — pipe accepts a "runnable-like" (an LLM/ChatModel instance or other runnable). Create the LLM with the model name in its constructor, then pass that instance to .pipe. To set per-call kwargs (stop, functions, etc.) use .bind(...) on the model runnable.
Examples:
-
Basic (string LLM / chat model):
const model = new ChatOpenAI({ model: "gpt-4o-mini", temperature: 0.2 });
const prompt = PromptTemplate.fromTemplate("Tell me a joke about {topic}");
const chain = prompt.pipe(model);
const out = await chain.invoke({ topic: "cats" }); -
Attach call-time options via bind:
const chain = prompt.pipe(model.bind({ stop: ["\n"], functions: myFunctionSchema })); -
Use RunnableSequence.from for explicit sequences:
const chain = RunnableSequence.from([prompt, model, outputParser]);
References: LangChain.js prompt+LLM examples and API docs. [1][2][3]
Sources
[1] Prompt + LLM examples (LangChain docs).
[2] Quickstart / model constructor (LangChain docs).
[3] Runnable / pipe API reference (LangChain.js API).
🌐 Web query:
@langchain/openai OpenAI constructor TypeScript JavaScript parameters
💡 Result:
Short answer — constructor signatures and main params
-
Signatures:
- new OpenAI(fields?: Partial & BaseLLMParams & { configuration?: ClientOptions & LegacyOpenAIInput }) → OpenAI. (v02.api.js.langchain.com)
- (older/alternate docs) new OpenAI(fields?: Partial & BaseLLMParams & { configuration?: ClientOptions }) → OpenAI. (v03.api.js.langchain.com)
-
Common fields you can pass (in the fields object)
- model (string) — model id (e.g. "gpt-4o", "gpt-3.5-turbo"). (docs.langchain.com)
- temperature (number), maxTokens / max_tokens, timeout, maxRetries. (js.langchain.com)
- apiKey (string) — or rely on OPENAI_API_KEY env var. (docs.langchain.com)
- streamUsage / stream_options related flags for streaming behavior. (docs.langchain.com)
- modelName (string) for calling fine-tuned models (e.g. ft:...). (docs.langchain.com)
- azureADTokenProvider (for Azure) when using Azure OpenAI. (v02.api.js.langchain.com)
-
configuration (ClientOptions) — passed either as fields.configuration or as a second optional arg; accepts official OpenAI SDK ClientOptions such as:
- baseURL, defaultHeaders (or custom headers), other ClientOptions accepted by openai SDK. Useful for proxies/custom endpoints or Azure. (v02.api.js.langchain.com)
Example (JS/TS):
- new OpenAI({ model: "gpt-3.5-turbo", temperature: 0.7, apiKey: process.env.OPENAI_API_KEY, configuration: { baseURL: "https://proxy.example.com", defaultHeaders: { Authorization: "Bearer X" } } }); (docs.langchain.com)
If you want, I can paste the exact TypeScript types (OpenAIInput, BaseLLMParams, ClientOptions) from the repo/docs.
Citations:
- 1: https://v02.api.js.langchain.com/classes/_langchain_openai.OpenAI.html?utm_source=openai
- 2: https://v03.api.js.langchain.com/classes/langchain_openai.OpenAI.html?utm_source=openai
- 3: https://docs.langchain.com/oss/javascript/integrations/llms/openai/?utm_source=openai
- 4: https://js.langchain.com/docs/integrations/llms/openai/?utm_source=openai
- 5: https://docs.langchain.com/oss/javascript/integrations/llms/openai/?utm_source=openai
- 6: https://docs.langchain.com/oss/javascript/integrations/chat/openai?utm_source=openai
- 7: https://docs.langchain.com/oss/javascript/integrations/chat/openai?utm_source=openai
- 8: https://v02.api.js.langchain.com/classes/_langchain_openai.OpenAI.html?utm_source=openai
- 9: https://v02.api.js.langchain.com/classes/_langchain_openai.OpenAI.html?utm_source=openai
- 10: https://docs.langchain.com/oss/javascript/integrations/llms/openai/?utm_source=openai
Fix two syntax errors in the TypeScript v0 LangChain example that will cause runtime failures.
Line 236 uses incorrect pipe syntax: prompt.pipe({ llm: model }). The correct LangChain.js pattern is prompt.pipe(model) where the LLM instance is passed directly to pipe.
Line 234 also has a syntax error: new OpenAI("gpt-4o-mini"). The OpenAI constructor requires an object parameter with fields like model, temperature, and other configuration options, not a positional string argument.
Apply these diffs:
-const model = new OpenAI("gpt-4o-mini");
+const model = new OpenAI({ model: "gpt-4o-mini" });
const prompt = PromptTemplate.fromTemplate("What is a good name for a company that makes {product}?");
-const chain = prompt.pipe({ llm: model });
+const chain = prompt.pipe(model);🤖 Prompt for AI Agents
In docs/docs/genai/tracing/integrations/listing/langchain.mdx around lines 230
to 240, the TypeScript v0 LangChain example has two syntax errors: the OpenAI
constructor is called with a positional string and should be called with an
options object (e.g., new OpenAI({ model: "gpt-4o-mini", ... })) and the pipe
call incorrectly passes an object; change prompt.pipe({ llm: model }) to
prompt.pipe(model) so the LLM instance is passed directly. Ensure any additional
OpenAI options (temperature, apiKey, etc.) are provided in the options object as
needed.
| if p.default is not None and ( | ||
| # In click >= 8.3.0, the default value is set to `Sentinel.UNSET` when no default is | ||
| # provided. Skip setting the default in this case. | ||
| # See https://github.com/pallets/click/pull/3030 for more details. | ||
| not isinstance(p.default, str) and repr(p.default) != "Sentinel.UNSET" | ||
| ): |
There was a problem hiding this comment.
🧩 Analysis chain
Critical logic error: legitimate string defaults are excluded.
The condition not isinstance(p.default, str) and repr(p.default) != "Sentinel.UNSET" will exclude all string defaults from the JSON schema because if p.default is a string, not isinstance(p.default, str) evaluates to False, making the entire expression False.
Since Sentinel.UNSET is a sentinel object (not a string), string defaults are always legitimate and should be included in the schema. The correct logic should include defaults that are either strings or non-sentinel values.
Apply this diff to fix the logic:
- if p.default is not None and (
- # In click >= 8.3.0, the default value is set to `Sentinel.UNSET` when no default is
- # provided. Skip setting the default in this case.
- # See https://github.com/pallets/click/pull/3030 for more details.
- not isinstance(p.default, str) and repr(p.default) != "Sentinel.UNSET"
- ):
+ if p.default is not None and (
+ # In click >= 8.3.0, the default value is set to `Sentinel.UNSET` when no default is
+ # provided. Skip setting the default in this case.
+ # See https://github.com/pallets/click/pull/3030 for more details.
+ isinstance(p.default, str) or repr(p.default) != "Sentinel.UNSET"
+ ):
schema["default"] = p.defaultVerify whether any Click commands in traces_cli or scorers_cli have string defaults that would be affected:
🏁 Script executed:
#!/bin/bash
# Description: Search for Click options/arguments with string defaults in CLI commands
# Search for default string values in Click decorators
rg -n -A 2 -B 2 '@click\.(option|argument)' mlflow/cli/traces/commands.py mlflow/cli/scorers/commands.py | rg -C 2 'default\s*=\s*["\']'Length of output: 68
Let me regenerate the verification script with corrected syntax:
🏁 Script executed:
#!/bin/bash
# Search for Click options/arguments with string defaults in mlflow/cli
rg -n -B 2 -A 2 "default\s*=\s*['\"]" mlflow/cliLength of output: 1604
🏁 Script executed:
#!/bin/bash
# Also check the structure of the CLI directory
fd -e py . mlflow/cli | head -20Length of output: 115
Let me check if the MCP server actually processes these CLI commands with string defaults:
🏁 Script executed:
#!/bin/bash
# Search for imports of traces or scorers CLI in mcp/server.py
rg -n "traces|scorers" mlflow/mcp/server.pyLength of output: 277
🏁 Script executed:
#!/bin/bash
# Check what functions/modules are imported from mlflow.cli in mcp/server.py
rg -n -A 5 "from mlflow.cli|import.*cli" mlflow/mcp/server.pyLength of output: 379
Critical logic error confirmed: legitimate string defaults are excluded from MCP tool schemas.
The condition not isinstance(p.default, str) and repr(p.default) != "Sentinel.UNSET" excludes all string defaults because if p.default is a string, not isinstance(p.default, str) is False, making the entire expression False.
This affects real commands: mlflow/cli/traces.py and mlflow/cli/scorers.py both have string defaults ("table", "main", "local") that are processed by cmd_to_function_tool() at lines 121-122 of mlflow/mcp/server.py. These defaults will be omitted from the generated MCP tool schemas.
Apply this diff:
- if p.default is not None and (
- # In click >= 8.3.0, the default value is set to `Sentinel.UNSET` when no default is
- # provided. Skip setting the default in this case.
- # See https://github.com/pallets/click/pull/3030 for more details.
- not isinstance(p.default, str) and repr(p.default) != "Sentinel.UNSET"
- ):
+ if p.default is not None and (
+ # In click >= 8.3.0, the default value is set to `Sentinel.UNSET` when no default is
+ # provided. Skip setting the default in this case.
+ # See https://github.com/pallets/click/pull/3030 for more details.
+ isinstance(p.default, str) or repr(p.default) != "Sentinel.UNSET"
+ ):
schema["default"] = p.default🤖 Prompt for AI Agents
In mlflow/mcp/server.py around lines 44-49, the conditional that checks whether
to include parameter defaults wrongly excludes legitimate string defaults due to
the `not isinstance(p.default, str) and` clause; remove that clause so the check
only skips Click's Sentinel.UNSET. Change the condition to: if p.default is not
None and repr(p.default) != "Sentinel.UNSET": so string defaults like "table",
"main", "local" are preserved in generated MCP tool schemas.
| def llm(): | ||
| # NB: CrewAI >= 1.0.0 introduced native LLM connectors that don't rely on LiteLLM. To use | ||
| # consistent mock between 1.x and 0.x, we opt-in to use LiteLLM for 1.x. | ||
| if _IS_CREWAI_V1: | ||
| from crewai import LLM | ||
|
|
||
| return LLM(model="openai/gpt-4o-mini", is_litellm=True) | ||
| else: | ||
| return "openai/gpt-4o-mini" | ||
|
|
There was a problem hiding this comment.
🧩 Analysis chain
Verify CrewAI v1 LLM API compatibility.
The version-aware LLM helper correctly differentiates between CrewAI v1 and earlier versions. However, ensure the is_litellm=True parameter is the correct way to opt into LiteLLM for CrewAI v1+.
🌐 Web query:
CrewAI version 1 LLM class is_litellm parameter documentation
💡 Result:
Short answer: There is no public, formal “is_litellm” parameter documented for CrewAI v1’s LLM class — it appears to be an internal flag injected when CrewAI calls litellm and is not part of the documented public API. Users have reported it showing up in litellm requests (and causing errors with some providers such as Groq). Workarounds: strip the key from litellm calls (monkey‑patch litellm.completion) or implement a custom BaseLLM to avoid the litellm path. References: community report showing the Groq error and a sample monkey‑patch, plus CrewAI docs on custom LLMs and LLM connections. [1][2][3]
Sources
- Community report and example monkey‑patch for removing "is_litellm". [1]
- CrewAI docs: Custom LLM implementation (BaseLLM). [2]
- CrewAI docs: Connect to any LLM / LiteLLM integration. [3]
Do not use undocumented is_litellm parameter; implement custom BaseLLM instead.
The is_litellm=True parameter is not part of CrewAI v1's documented public API and is an internal flag. It's known to cause errors with certain providers (e.g., Groq). Per CrewAI docs, use a custom BaseLLM implementation or the recommended LiteLLM integration approach for reliable, supported behavior.
🤖 Prompt for AI Agents
In tests/crewai/test_crewai_autolog.py around lines 30 to 39, replace the
undocumented use of the internal is_litellm=True flag with a small custom
BaseLLM subclass that implements the minimal public BaseLLM API and delegates to
the desired model ("openai/gpt-4o-mini") or integrates with LiteLLM via the
documented extension points; remove the is_litellm parameter entirely,
instantiate and return your custom BaseLLM implementation when _IS_CREWAI_V1 is
true, and for older versions keep returning the string model name so tests use a
consistent, supported mock.
10278ca to
b933f52
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
tests/db/test_schema.py (1)
2-2: Dialect‑aware unique‑constraint reattachment looks correct; consider sharing helper logicThe
_reattach_missing_unique_constraints()/_get_unique_constraints()combination correctly compensates for MySQL and MSSQL reflection quirks and is scoped to the expected constraint name set, so it shouldn’t disturb other tables. The call fromdump_schema()is also in the right place (after reflection, before DDL generation). One improvement you might consider is de‑duplicating this reflection logic with the similar helpers in the Alembic migration module so future constraint additions stay in sync.Also applies to: 8-16, 32-32, 46-88, 90-107
tests/db/check_migration.py (1)
52-61: Workspace backfill checks are solid; table-name list could be centralizedThe added
WORKSPACE_TABLESplus theSELECT DISTINCT workspace ...assertions inpost_migration()give good coverage that all affected tables are fully backfilled to the"default"workspace. To avoid drift with future workspace-enabled tables, you might later derive this set from a shared constant (e.g., the migration module or ORM models) instead of hard-coding the names here.Also applies to: 69-96, 129-132
mlflow/store/model_registry/sqlalchemy_store.py (1)
64-64: Registry now writes workspace correctly; future multi-workspace support will need query changesThe changes to use
DEFAULT_WORKSPACE_NAMEwhen creating registered models, model versions, tags, and aliases (including thesession.merge(...)upserts) are necessary and sufficient to satisfy the new composite PK/FK schemas while you still operate in a single default workspace. All write paths now populate the workspace column consistently.Read/lookup paths (
_get_registered_model*,_get_model_version*, tag and alias getters, search queries) still filter only onname/version, which is fine as long as the only workspace in play is the default one. Once multi-workspace behavior is surfaced at the store or API level, you’ll need to thread workspace context through these queries and joins to avoid ambiguity across workspaces.Also applies to: 172-207, 669-707, 730-821, 1186-1212, 1237-1270
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (17)
docs/api_reference/api_inventory.txt(2 hunks)mlflow/entities/__init__.py(2 hunks)mlflow/entities/workspace.py(1 hunks)mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py(1 hunks)mlflow/store/model_registry/dbmodels/models.py(11 hunks)mlflow/store/model_registry/sqlalchemy_store.py(6 hunks)mlflow/store/tracking/dbmodels/models.py(5 hunks)mlflow/store/workspace/dbmodels/__init__.py(1 hunks)mlflow/store/workspace/dbmodels/models.py(1 hunks)mlflow/utils/workspace_utils.py(1 hunks)tests/db/check_migration.py(3 hunks)tests/db/schemas/mssql.sql(6 hunks)tests/db/schemas/mysql.sql(6 hunks)tests/db/schemas/postgresql.sql(6 hunks)tests/db/schemas/sqlite.sql(6 hunks)tests/db/test_schema.py(3 hunks)tests/db/test_workspace_migration.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (5)
- mlflow/entities/init.py
- mlflow/utils/workspace_utils.py
- mlflow/entities/workspace.py
- mlflow/store/workspace/dbmodels/init.py
- tests/db/schemas/mysql.sql
🧰 Additional context used
🧬 Code graph analysis (7)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
tests/db/test_schema.py (1)
_get_unique_constraints(90-107)
tests/db/check_migration.py (1)
mlflow/store/model_registry/sqlalchemy_store.py (1)
create_registered_model(172-214)
mlflow/store/workspace/dbmodels/models.py (1)
mlflow/entities/workspace.py (1)
Workspace(9-13)
mlflow/store/model_registry/sqlalchemy_store.py (1)
mlflow/store/model_registry/dbmodels/models.py (3)
SqlRegisteredModelTag(155-190)SqlModelVersionTag(193-230)SqlRegisteredModelAlias(233-268)
tests/db/test_schema.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
_get_unique_constraints(144-159)
tests/db/test_workspace_migration.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (2)
upgrade(138-421)downgrade(424-678)
mlflow/store/model_registry/dbmodels/models.py (1)
mlflow/entities/model_registry/registered_model.py (4)
name(44-46)name(49-50)description(69-71)description(74-75)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (34)
- GitHub Check: lint
- GitHub Check: typescript-sdk (20)
- GitHub Check: typescript-sdk (24)
- GitHub Check: typescript-sdk (22)
- GitHub Check: js (ubuntu-latest, --testPathPattern)
- GitHub Check: js (windows-latest, --testPathPattern)
- GitHub Check: js (windows-latest, --testPathIgnorePatterns)
- GitHub Check: js (ubuntu-latest, --testPathIgnorePatterns)
- GitHub Check: windows (3)
- GitHub Check: pyfunc (3)
- GitHub Check: windows (4)
- GitHub Check: pyfunc (4)
- GitHub Check: database
- GitHub Check: pyfunc (2)
- GitHub Check: pyfunc (1)
- GitHub Check: windows (1)
- GitHub Check: windows (2)
- GitHub Check: java
- GitHub Check: python (3)
- GitHub Check: python-skinny
- GitHub Check: evaluate (1)
- GitHub Check: python (2)
- GitHub Check: evaluate (2)
- GitHub Check: python (1)
- GitHub Check: flavors
- GitHub Check: models (1)
- GitHub Check: models (2)
- GitHub Check: genai
- GitHub Check: core
- GitHub Check: build (tracing)
- GitHub Check: build (skinny)
- GitHub Check: build (dev)
- GitHub Check: r
- GitHub Check: build
🔇 Additional comments (23)
docs/api_reference/api_inventory.txt (1)
531-531: Workspace entries correctly added to API inventoryBoth
mlflow.entities.Workspaceandmlflow.entities.workspace.Workspaceare listed consistently with other entities; this keeps the public surface in sync with the new model.Also applies to: 631-631
tests/db/schemas/mssql.sql (1)
29-31: MSSQL schema workspace changes align with the migrationThe added
workspacecolumns, composite PKs/FKs,uq_experiments_workspace_name, and the newworkspacestable all match the shapes introduced in the Alembic migration and ORM models. This should keeptest_schema_is_up_to_dateconsistent for MSSQL after running the migration.Also applies to: 41-44, 85-87, 100-101, 105-109, 195-198, 205-208, 215-218, 364-367
mlflow/store/workspace/dbmodels/models.py (1)
1-22: SqlWorkspace ORM mapping matches the new catalog tableThe
SqlWorkspacemodel correctly mirrors theworkspacestable (PK onname, optionaldescription) and cleanly converts to theWorkspaceentity viato_mlflow_entity. Looks good.tests/db/schemas/postgresql.sql (1)
29-31: PostgreSQL schema workspace DDL matches the migration designWorkspace columns, composite primary keys / foreign keys,
uq_experiments_workspace_name, and theworkspacestable definition are all consistent with the Alembic upgrade logic and ORM models, so the reflected schema should continue to match this file.Also applies to: 41-44, 86-88, 101-102, 106-110, 198-201, 208-211, 218-221, 370-373
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (2)
20-25: Upgrade path and MSSQL unique-introspection helpers look soundThe migration’s helper stack (
_NAMING_CONVENTION,_WORKSPACE_TABLES,_SQLITE_LEGACY_FKS,_fetch_mssql_unique_metadata, and the_get_unique_constraints/_get_unique_indexes/_detect_unique_on_name/_collect_foreign_keystrio) is thoughtfully put together: it handles the lack ofget_unique_constraints()on MSSQL, the MySQL “unique-as-index” pattern, and unnamed SQLite FKs without overreaching into other dialects. The upgrade steps then cleanly:
- Drop pre-existing uniqueness on
nameforexperiments/registered_models,- Add non-null
workspacewith default to the 8 tables,- Rebuild PKs/FKs to use
(workspace, …)keys, and- Create the required workspace indexes and the seeded
workspacescatalog.Given the amount of dialect-specific introspection, it’s worth running this migration against all four supported backends to confirm inspector behavior matches assumptions.
Also applies to: 26-61, 64-134, 136-238
424-678: Downgrade safety checks are careful; evaluation_datasets conflict check is slightly conservativeThe downgrade logic is well-structured: it (1) reflects tables once and caches them, (2) asserts there are no conflicting rows when collapsing workspaces back to a single global space, (3) bulk-updates all workspace columns to
"default", (4) tears down workspace-specific FKs/PKs/indexes/columns in a dialect-aware way (including batched SQLite ops), and (5) restores pre-workspace PKs/FKs plus auq_experiments_nameconstraint before droppingworkspaces. The only mildly conservative bit is treating collisions onevaluation_datasets.nameas a hard blocker even though that table’s PK never changes, but that only affects downgrade scenarios and errs on the side of safety, so it’s acceptable.Please sanity-check downgrade by applying upgrade + downgrade on a test DB with some synthetic data across the impacted tables for all dialects you care about.
mlflow/store/tracking/dbmodels/models.py (4)
64-64: LGTM: Workspace column properly configured.The workspace column is correctly defined with:
- Appropriate length (63 characters)
- Non-nullable constraint with proper defaults
- Server-side default for database-level consistency
Also applies to: 103-111
98-102: LGTM: Uniqueness correctly scoped to workspace.The unique constraint has been properly updated from a single-column constraint on
nameto a composite constraint on(workspace, name), ensuring experiment names are unique per workspace.Also applies to: 137-137
1315-1323: LGTM: Workspace support added consistently.The workspace column and index for
SqlEvaluationDatasetfollow the same pattern asSqlExperiment, maintaining consistency across the codebase.Also applies to: 1379-1379
143-158: Verify workspace handling in entity layer.The
to_mlflow_entity()methods don't propagate theworkspacefield to the returnedExperimentandEvaluationDatasetentities. Please confirm this is intentional and that workspace context is managed at the store or API layer rather than in the entity objects themselves.Also applies to: 1382-1419
tests/db/test_workspace_migration.py (5)
10-21: LGTM: Test setup properly isolated.The database preparation creates a fresh SQLite database for each test using
tmp_path, ensuring test isolation. Migration toPREVIOUS_REVISIONcorrectly establishes the pre-workspace state.
24-232: LGTM: Pre-migration seeding correctly uses legacy schema.The deliberate use of raw SQL to seed data without workspace columns is the right approach for testing the migration. The comment clearly explains this design decision.
242-579: LGTM: Post-migration helpers properly workspace-aware.The helper functions correctly require
workspaceparameters and use the post-migration schema, providing good test utilities for scenarios after the migration is applied.
590-719: LGTM: Comprehensive upgrade validation.The upgrade test thoroughly validates:
- Workspace column presence and defaults
- Data backfill to 'default' workspace
- Updated primary keys and constraints
- New indexes
This provides strong confidence in the migration's correctness.
722-941: LGTM: Thorough downgrade and conflict testing.The downgrade validation ensures schema correctly reverts to pre-workspace state. The parameterized conflict detection tests comprehensively cover different resource types that could have name collisions when merging workspaces, with proper error message validation.
tests/db/schemas/sqlite.sql (4)
29-29: LGTM: Workspaces catalog and dataset scoping properly defined.The new
workspacestable provides a central catalog, and the workspace column inevaluation_datasetsfollows the established pattern.Also applies to: 106-110
41-44: LGTM: Experiment uniqueness correctly scoped.The unique constraint change from
namealone to(workspace, name)enables workspace-scoped experiment names.
86-87: LGTM: Model registry tables consistently workspace-scoped.All model registry tables have been properly updated with:
- Workspace columns with appropriate defaults
- Composite primary keys with workspace as the leading column
- Foreign keys correctly referencing the new composite keys
- Appropriate CASCADE options for referential integrity
Also applies to: 198-200, 208-210, 218-220, 372-374
101-102: Verify webhook primary key design.The
webhookstable adds a workspace column but retainswebhook_idas the sole primary key (unlikeregistered_modelswhich uses a composite PK). Confirm this is intentional—likely becausewebhook_idis globally unique across all workspaces.mlflow/store/model_registry/dbmodels/models.py (4)
1-1: LGTM: Workspace support properly integrated into SqlRegisteredModel.The workspace column and composite primary key changes are consistent with the tracking models, establishing workspace-scoped registered models.
Also applies to: 36-36, 42-57
88-93: LGTM: Model version properly linked to workspace-scoped registered models.The foreign key constraint correctly references the composite key
(workspace, name)inregistered_modelswithonupdate="cascade"to maintain referential integrity.Also applies to: 127-133
196-223: LGTM: Model version tags properly scoped and linked.The workspace column, composite foreign key, and bidirectional relationship are all correctly configured. The relationship provides convenient navigation from model versions to their tags.
158-183: LGTM: Remaining models consistently updated with workspace support.All models follow the established pattern with workspace columns, appropriate indexes, and correctly configured foreign key relationships. The webhook design (workspace column with non-composite PK) is consistent with the schema.
Also applies to: 236-261, 301-321
mprahl
force-pushed
the
workspaces-db-model
branch
from
b933f52 to
f8322c0
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (4)
tests/db/test_schema.py (1)
2-16: Dialect-specific reattachment of unique constraints looks solid; keep the mapping in sync over timeThe
_reattach_missing_unique_constraints/_get_unique_constraintspairing is a reasonable way to normalize reflection quirks (MySQL unique-as-index, MSSQL lackingget_unique_constraints) before generating schema strings. Usingduplicates_indexto gate MySQL recreation and warning on missing columns are good guardrails.One thing to watch:
_DIALECT_REFLECTED_UNIQUE_CONSTRAINTSis currently hard-coded foruq_experiments_workspace_nameonly. If future migrations introduce additional unique constraints that suffer from the same reflection issues on MySQL or MSSQL, this map will need to be extended to keepdump_schemaoutput aligned with the real schema.Also applies to: 32-43, 46-88, 90-107
mlflow/store/tracking/dbmodels/models.py (1)
64-65: Tracking models’ workspace columns match the migration; entity mapping can stay “default-only” for nowThe additions of
workspacetoSqlExperiment(withuq_experiments_workspace_name) and toSqlEvaluationDataset(plusidx_evaluation_datasets_workspace) line up with the Alembic migration and the updated Postgres schema: same length, non-nullability, and'default'server default viaDEFAULT_WORKSPACE_NAME. That gives you a consistent workspace dimension in the tracking schema.Right now
to_mlflow_entity/from_mlflow_entityfor these models ignoreworkspace, which keeps the runtime effectively “single default workspace” even though the schema is multi-tenant-ready. That’s fine if the intention of this PR is to get the database into a workspace-capable shape first, but once workspace becomes part of the user-visible API you’ll likely want to thread it through the corresponding entities and conversions.Also applies to: 94-111, 131-138, 1309-1323, 1375-1380
tests/db/test_workspace_migration.py (1)
1-21: Workspace migration tests give good end-to-end coverage of upgrade, downgrade, and conflict scenariosThe SQLite-based harness that seeds a pre-workspace schema, upgrades to
REVISION, and then exercises both upgrade and downgrade paths looks solid:
_seed_pre_workspace_entitiescorrectly targets the legacy (no-workspace) layout, ensuring the migration itself is responsible for backfilling.- Post-upgrade assertions validate:
- workspace columns (non-null +
'default'default),- data backfill for experiments, registry objects, and evaluation datasets,
- new PK shapes and FKs, and
- the expected workspace indexes.
- Post-downgrade assertions confirm:
- removal of workspace columns and the
workspacestable,- restoration of legacy PK/FK shapes, and
- that data survive a downgrade unchanged.
- The conflict tests neatly mirror
_assert_no_workspace_conflictsand verify that both the downgrade guards and the underlying grouping queries behave as intended.The only minor note is that
_get_workspace_defaultcurrently assumes fairly simple default representations, which is fine for SQLite in this test but might need tightening if you ever extend these checks to other dialects.Also applies to: 24-239, 242-247, 530-587, 590-718, 722-805, 807-859, 861-895, 898-941
mlflow/store/model_registry/sqlalchemy_store.py (1)
64-65: Registry store now writes workspace-aware rows while remaining effectively “default workspace only”Plumbing
DEFAULT_WORKSPACE_NAMEinto:
SqlRegisteredModel/SqlRegisteredModelTagcreation,SqlModelVersion/SqlModelVersionTagcreation,set_registered_model_tag,set_model_version_tag, andset_registered_model_alias,keeps all registry writes compatible with the new composite primary keys and non-null
workspacecolumns in the ORM models.At the same time, lookup paths like
_get_registered_model,_get_registered_model_tag,_get_model_version_tag, and_get_registered_model_aliasstill filter only onname(and version/key), notworkspace, so the effective behavior remains “single default workspace” for now. That’s a reasonable interim state for a schema-first migration; when you later introduce true multi-workspace behavior, those queries and the tag/alias search joins will need to be extended to include a workspace dimension to avoid ambiguity.Also applies to: 172-206, 670-685, 686-707, 730-821, 1186-1212, 1237-1245, 1247-1270
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (18)
docs/api_reference/api_inventory.txt(2 hunks)mlflow/entities/__init__.py(2 hunks)mlflow/entities/workspace.py(1 hunks)mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py(1 hunks)mlflow/store/model_registry/dbmodels/models.py(11 hunks)mlflow/store/model_registry/sqlalchemy_store.py(6 hunks)mlflow/store/tracking/dbmodels/models.py(5 hunks)mlflow/store/workspace/dbmodels/__init__.py(1 hunks)mlflow/store/workspace/dbmodels/models.py(1 hunks)mlflow/utils/workspace_utils.py(1 hunks)tests/db/check_migration.py(3 hunks)tests/db/schemas/mssql.sql(6 hunks)tests/db/schemas/mysql.sql(6 hunks)tests/db/schemas/postgresql.sql(6 hunks)tests/db/schemas/sqlite.sql(6 hunks)tests/db/test_schema.py(3 hunks)tests/db/test_workspace_migration.py(1 hunks)tests/resources/db/latest_schema.sql(6 hunks)
🚧 Files skipped from review as they are similar to previous changes (6)
- mlflow/entities/workspace.py
- tests/db/check_migration.py
- mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py
- mlflow/entities/init.py
- mlflow/utils/workspace_utils.py
- mlflow/store/workspace/dbmodels/models.py
🧰 Additional context used
🧬 Code graph analysis (5)
tests/db/test_workspace_migration.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (2)
upgrade(138-615)downgrade(618-873)
mlflow/store/workspace/dbmodels/__init__.py (1)
mlflow/store/workspace/dbmodels/models.py (1)
SqlWorkspace(10-22)
tests/db/test_schema.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
_get_unique_constraints(144-159)
mlflow/store/model_registry/dbmodels/models.py (1)
mlflow/entities/model_registry/registered_model.py (4)
name(44-46)name(49-50)description(69-71)description(74-75)
mlflow/store/model_registry/sqlalchemy_store.py (1)
mlflow/store/model_registry/dbmodels/models.py (3)
SqlRegisteredModelTag(155-190)SqlModelVersionTag(193-230)SqlRegisteredModelAlias(233-268)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (34)
- GitHub Check: js (windows-latest, --testPathIgnorePatterns)
- GitHub Check: js (ubuntu-latest, --testPathPattern)
- GitHub Check: js (ubuntu-latest, --testPathIgnorePatterns)
- GitHub Check: js (windows-latest, --testPathPattern)
- GitHub Check: typescript-sdk (24)
- GitHub Check: typescript-sdk (20)
- GitHub Check: flavors
- GitHub Check: typescript-sdk (22)
- GitHub Check: core
- GitHub Check: lint-macos
- GitHub Check: windows (4)
- GitHub Check: windows (2)
- GitHub Check: models (1)
- GitHub Check: models (2)
- GitHub Check: pyfunc (4)
- GitHub Check: pyfunc (3)
- GitHub Check: database
- GitHub Check: genai
- GitHub Check: windows (3)
- GitHub Check: java
- GitHub Check: evaluate (2)
- GitHub Check: windows (1)
- GitHub Check: pyfunc (2)
- GitHub Check: evaluate (1)
- GitHub Check: pyfunc (1)
- GitHub Check: python (3)
- GitHub Check: build (tracing)
- GitHub Check: build (skinny)
- GitHub Check: r
- GitHub Check: python (1)
- GitHub Check: build (dev)
- GitHub Check: python (2)
- GitHub Check: python-skinny
- GitHub Check: build
🔇 Additional comments (15)
mlflow/store/workspace/dbmodels/__init__.py (1)
1-3: ExposeSqlWorkspacevia package root looks correctRe-exporting
SqlWorkspacefrommlflow.store.workspace.dbmodelsvia__all__matches existing dbmodels patterns and provides the expected import surface for callers.docs/api_reference/api_inventory.txt (1)
531-531: Workspace entity entries align with existing inventory conventionsListing both
mlflow.entities.Workspaceandmlflow.entities.workspace.Workspaceis consistent with how other entities are represented here (top-level alias + module-qualified class). No further changes needed.Also applies to: 631-631
tests/db/schemas/postgresql.sql (1)
19-31: PostgreSQL schema updates correctly reflect workspace-scoped registry layoutThe added
workspacecolumns, updated composite primary keys, and adjusted foreign keys (e.g.,model_versions→registered_models, tag/alias tables) are consistent with the workspace-aware migration and ORM models. The newworkspacescatalog table also matches the expected shape. I don’t see schema-level issues here.Also applies to: 34-45, 81-88, 91-103, 106-110, 184-201, 204-211, 214-221, 365-373
tests/db/schemas/mysql.sql (4)
29-29: LGTM: Workspace column additions are consistent.The workspace columns are uniformly defined across all affected tables with
VARCHAR(63) DEFAULT 'default' NOT NULL, which ensures backward compatibility with existing data.Also applies to: 41-41, 86-86, 101-101, 197-197, 207-207, 217-217, 372-372
106-110: LGTM: Workspaces catalog table.The new
workspacestable properly establishes the foundation for multi-tenancy with a simple structure (name as PK, optional description).
43-44: LGTM: Workspace-scoped experiment name uniqueness.The unique constraint
uq_experiments_workspace_nameon(workspace, name)correctly ensures experiments are unique within each workspace while allowing the same name across different workspaces.
87-87: Review the cascade constraint asymmetry and confirm design intent.Your review comment correctly identifies an inconsistent cascade pattern in the schema. The verification confirms:
model_versions,registered_model_tags, andmodel_version_tagsuse ON UPDATE CASCADE only (no ON DELETE CASCADE)registered_model_aliasesuses ON DELETE CASCADE ON UPDATE CASCADE- No tests validate cascade behavior for these constraints
- No documentation explains why the design differs
The asymmetry between
registered_model_aliasesand the other workspace-aware model tables warrants clarification. Confirm whether:
- Deletion of a
registered_modelshould cascade tomodel_versions,registered_model_tags, andmodel_version_tags(matching business requirements)- The inconsistency with
registered_model_aliasesis intentional, or ifmodel_versionsshould also include ON DELETE CASCADE- Application code handles cascading deletes manually where foreign keys restrict them
tests/db/schemas/sqlite.sql (1)
29-29: LGTM: SQLite schema consistent with MySQL.The workspace additions, composite keys, and foreign key constraints in the SQLite schema are consistent with the MySQL schema, ensuring cross-database compatibility.
Also applies to: 41-41, 86-86, 101-101, 106-110, 198-200, 208-210, 218-220, 372-374
tests/resources/db/latest_schema.sql (1)
29-29: LGTM: Latest schema reference is consistent.The latest schema file correctly reflects the workspace-aware multi-tenancy changes and matches the dialect-specific schema files (MySQL, SQLite, MSSQL).
Also applies to: 41-44, 86-87, 101-101, 106-110, 198-200, 208-210, 218-220, 372-374
tests/db/schemas/mssql.sql (1)
29-29: LGTM: MSSQL schema with appropriate dialect-specific syntax.The MSSQL schema correctly uses:
COLLATE "SQL_Latin1_General_CP1_CI_AS"for VARCHAR columns- Parentheses around DEFAULT values:
DEFAULT ('default')- Workspace-aware composite keys and foreign key constraints
The workspace additions are structurally consistent with other dialects.
Also applies to: 41-41, 85-86, 100-100, 105-109, 195-197, 205-207, 215-217, 364-366
mlflow/store/model_registry/dbmodels/models.py (5)
57-57: LGTM: Composite primary keys correctly defined.All workspace-aware tables have properly updated their primary keys to include
workspaceas the first component of composite keys:
SqlRegisteredModel: (workspace, name)SqlModelVersion: (workspace, name, version)SqlRegisteredModelTag: (workspace, key, name)SqlModelVersionTag: (workspace, key, name, version)SqlRegisteredModelAlias: (workspace, name, alias)Also applies to: 132-132, 182-183, 217-218, 260-260
126-133: LGTM: Foreign key constraints properly reference workspace-qualified keys.The
ForeignKeyConstraintdefinitions correctly reference composite keys that include workspace:
SqlModelVersion→registered_models(workspace, name)with ON UPDATE CASCADESqlRegisteredModelTag→registered_models(workspace, name)with ON UPDATE CASCADESqlModelVersionTag→model_versions(workspace, name, version)with ON UPDATE CASCADESqlRegisteredModelAlias→registered_models(workspace, name)with ON DELETE CASCADE and ON UPDATE CASCADEThis ensures referential integrity across the workspace-scoped multi-tenancy model.
Also applies to: 176-183, 211-218, 252-261
36-36: No issues found. The import and constant usage are correct.Verification confirms that
DEFAULT_WORKSPACE_NAME = "default", which matches the expected value in the SQL schema. The import at line 36 is correct, and its usage inserver_defaultparameters throughout the SQLAlchemy models aligns with the schema definitions.
65-82: The review comment is incorrect and should be dismissed.The workspace field in
SqlRegisteredModelis an internal database/storage concern for multi-tenancy and is deliberately not included in theRegisteredModelentity or API contract. This is confirmed by:
- The protobuf schema (
mlflow/protos/model_registry.protoline 382+) explicitly excludes workspace from theRegisteredModelmessage definition- The entity's
from_proto()andto_proto()methods make no reference to workspace- Similar multi-tenant models like
SqlRegisteredModelTagandSqlRegisteredModelAliasalso exclude workspace from their entity/proto representationsWorkspace is an architectural boundary: it exists at the backend storage layer to support multi-tenancy but does not cross into the entity or API layers. Adding it to
RegisteredModelwould require modifying the proto schema and could break the API contract.Likely an incorrect or invalid review comment.
317-322: Review comment is incorrect - workspace is not enforced via FK to workspaces table in this codebase.The review's assumption that "other workspace-aware tables" have FK constraints to the workspaces table is false. Investigation shows SqlRegisteredModelTag, SqlModelVersion, and other workspace-aware tables in model_registry all store workspace as plain String(63) columns without FK constraints. This is the established pattern: workspace functions as a multi-tenant identifier string, not a foreign key reference. No table in mlflow/store/model_registry has a FK to the workspaces table.
However, there is a separate schema issue: the workspace column exists in the SqlWebhook model definition but is missing from the migration file (1a0cddfcaa16), which could cause migration failures or schema sync problems.
Likely an incorrect or invalid review comment.
mprahl
force-pushed
the
workspaces-db-model
branch
from
f8322c0 to
0446441
Compare
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
mlflow/store/model_registry/dbmodels/models.py (1)
271-295: Critical: Fallback key generation causes data loss on restart.Line 284 generates a random encryption key if
MLFLOW_WEBHOOK_SECRET_ENCRYPTION_KEYis not set. This key is ephemeral and not persisted, meaning:
- Each application restart generates a different key
- Previously encrypted secrets become permanently unreadable
- Webhook secrets are effectively lost
This creates a critical data integrity issue in production.
Apply this diff to fail fast when the encryption key is missing:
def __init__(self): super().__init__() - # Get encryption key from environment variable or generate one - # In production, this should come from a secure key management service - encryption_key = MLFLOW_WEBHOOK_SECRET_ENCRYPTION_KEY.get() or Fernet.generate_key() + # Get encryption key from environment variable + # In production, this should come from a secure key management service + encryption_key = MLFLOW_WEBHOOK_SECRET_ENCRYPTION_KEY.get() + if encryption_key is None: + raise ValueError( + "MLFLOW_WEBHOOK_SECRET_ENCRYPTION_KEY must be set to encrypt webhook secrets" + ) self.cipher = Fernet(encryption_key)
🧹 Nitpick comments (14)
mlflow/utils/workspace_utils.py (1)
1-5: DEFAULT_WORKSPACE_NAME helper looks good; consider centralizing usage in tests.The constant and export are clean and match the rest of the PR’s workspace defaults. As a follow‑up, consider importing and using
DEFAULT_WORKSPACE_NAMEin places liketests/db/check_migration.py(rather than hard‑coding"default") to reduce drift if the default ever changes.tests/db/schemas/postgresql.sql (1)
19-31: Workspace-aware DDL looks coherent; consider optional FK toworkspaces.The PostgreSQL schema updates form a consistent workspace model:
- All relevant tables get
workspace VARCHAR(63) DEFAULT 'default'::character varying NOT NULL.- Experiments have
uq_experiments_workspace_namefor(workspace, name).- Registered models, versions, aliases, and tags move to composite PKs/FKs including
workspace.- A dedicated
workspacestable is added withnameas PK.This all lines up with the multi‑tenancy changes elsewhere in the PR and should work fine as-is.
If, in the future, you want stronger integrity guarantees, you could optionally add
FOREIGN KEY (workspace) REFERENCES workspaces (name)to the workspace-backed tables to enforce that all workspace values are registered, but that’s a design/compat decision and not required for this change.Also applies to: 34-45, 81-88, 91-103, 106-110, 184-201, 204-211, 214-221, 365-373
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (2)
64-130: MSSQL unique-metadata helper looks solid; keep it aligned with schema testsThe
_fetch_mssql_unique_metadata()helper plus the_get_unique_constraints()/_get_unique_indexes()wrappers give you a robust MSSQL fallback, mirroring whattests/db/test_schema.pydoes via inspector/indexes. The catalog queries and the returned shape ({"name", "column_names"}) look correct.The only thing to watch is long‑term drift: if you ever add new unique constraints or change naming conventions for MSSQL, this helper and the schema tests (_get_unique_constraints in
tests/db/test_schema.py) must be updated together so migration detection onnamekeeps working consistently.
618-751: Downgrade conflict detection is thorough; keep specs in sync with schema evolutionThe
_assert_no_workspace_conflicts()+conflict_specs+legacy_*_defssetup indowngrade()is a good safety net to prevent lossy merges when dropping the workspace dimension, and the messages align with what the tests expect.Going forward, any new workspace‑scoped tables or changes to key uniqueness (e.g., adding a unique on
(workspace, something_else)) will need matching updates in:
conflict_specs_WORKSPACE_TABLESlegacy_pk_defs/legacy_fk_defsto avoid silent mismatches during downgrade. It could be worth adding a brief comment near these lists stating that they must be kept in lockstep with the “latest schema” SQL files and workspace‑migration tests.
mlflow/store/tracking/dbmodels/models.py (2)
92-138: Experiment workspace column and uniqueness look consistent with the migrationThe
SqlExperiment.workspacecolumn plusUniqueConstraint("workspace", "name", name="uq_experiments_workspace_name")line up with the new migration and schema snapshots. This preserves existing behavior for a single default workspace while giving room to scope experiment names per workspace.Note that the migration hard‑codes
'default'as the server default while the ORM usesDEFAULT_WORKSPACE_NAME; if you ever change that constant, you’ll need a follow‑up migration rather than just tweaking this model, otherwise existing rows and new inserts could diverge.
1303-1380: Evaluation dataset workspace support is wired at the DB level but not surfaced in entitiesAdding
workspacewith default/server_default and theidx_evaluation_datasets_workspaceindex matches the migration and dialect schemas and should give good performance for workspace‑scoped queries.Right now
SqlEvaluationDataset.to_mlflow_entity()/from_mlflow_entity()don’t expose or accept workspace, so everything implicitly lives in the default workspace. That’s fine for this step, but once you introduce a workspace‑aware public API, you’ll need to thread workspace through these mappers (and probably theEvaluationDatasetentity) to avoid silent misrouting of data.tests/db/test_workspace_migration.py (3)
235-240:_get_workspace_defaultis fine for SQLite; keep scope limited
_get_workspace_default()is tailored to how SQLite reports column defaults (string expressions) and works for the engine created in_prepare_database. If you ever extend these tests to run against other dialects, you’ll likely need a more dialect‑aware default parser (or to reuse SQLAlchemy’s compiled default inspection helpers) to avoid brittle string handling.
590-659: Upgrade test gives strong coverage; consider also asserting onwebhooks.workspace
test_workspace_migration_upgrade_adds_columns_and_backfills()does a nice job verifying:
- workspace columns exist and are non‑nullable,
- legacy rows are backfilled to
"default",- PKs/UKs/FKs and indexes match the new schema,
- the
workspacestable is created and seeded.Since the migration also adds a non‑nullable
workspacetowebhooksand createsidx_webhooks_workspace, it’d be a small but useful extension to:
- include
"webhooks"in theassert_workspace_column(...)loop, and- assert that
has_index("webhooks", "idx_webhooks_workspace", ["workspace"])is true.Not critical, but it would fully exercise all the workspace‑aware tables in the upgrade path.
807-941: Downgrade conflict tests align well with the downgrade safety checksThe downgrade conflict tests (
test_workspace_migration_downgrade_detects_conflictsandtest_workspace_migration_conflict_detection_queries) mirror the logic in_assert_no_workspace_conflictsand_fetch_conflicts, and the expected fragments match theresource_descriptionstrings in the migration.This is exactly the kind of safety net you want before allowing a downgrade that merges workspaces; as you add more workspace‑scoped tables or adjust key uniqueness, remember to add corresponding
_setup_*_conflicthelpers and param cases here.mlflow/store/model_registry/sqlalchemy_store.py (4)
172-207: Creating registry objects in the default workspace is correct but implicitUsing
workspace=DEFAULT_WORKSPACE_NAMEwhen creatingSqlRegisteredModeland its tags ensures new rows satisfy the composite PK/constraints introduced in the models and migrations, and keeps behavior equivalent to the pre‑workspace world (everything in a single logical tenant).This is a good transitional step; just keep in mind that if/when you support non‑default workspaces through this store, you’ll probably want to make workspace an explicit parameter rather than relying on a hard‑coded default here.
670-707: Tag helpers should become workspace‑aware before multi‑workspace support goes live
set_registered_model_tag()now merges aSqlRegisteredModelTagwithworkspace=DEFAULT_WORKSPACE_NAME, but the corresponding getter_get_registered_model_tag()still filters only on(name, key)and assumes a single row.That’s fine as long as every registered model lives in the default workspace, but once you allow the same model name in multiple workspaces you’ll need to:
- include
workspacein the query predicates (or in a uniqueness guarantee), and- decide how workspace is passed through the public API.
Otherwise you’ll either hit the “Found N tags” INVALID_STATE branch or mix tags across workspaces.
1186-1212: Model version tag helpers have the same single‑workspace assumption
set_model_version_tag()now persists tags withworkspace=DEFAULT_WORKSPACE_NAME, but_get_model_version_tag()still queries only by(name, version, key). Similar to the registered‑model tag case, this implicitly assumes there is at most one workspace in play for a given(name, version).Before multi‑workspace usage of the registry store, it would be safer either to:
- add
SqlModelVersionTag.workspace == DEFAULT_WORKSPACE_NAMEto the query, or- make workspace an explicit parameter on the public tag APIs.
That will avoid ambiguous results or cross‑workspace interference later.
1247-1270: Registered model aliases also need a workspace story
set_registered_model_alias()now merges aSqlRegisteredModelAliaswithworkspace=DEFAULT_WORKSPACE_NAME, but_get_registered_model_alias()still filters by(name, alias)only. As with tags, this is fine while there’s only the default workspace, but it’s ambiguous as soon as you allow the same alias/name pair in multiple workspaces.Consider either:
- constraining this store to the default workspace and documenting that assumption, or
- threading workspace into alias lookups (and making it part of the public API) so that PK and query conditions stay aligned.
mlflow/store/model_registry/dbmodels/models.py (1)
42-57: Consider adding an index on workspace for query performance.While the composite primary key on
(workspace, name)exists, queries that filter by workspace alone (e.g., listing all models in a workspace) would benefit from a dedicated index.Apply this diff to add the index:
- __table_args__ = (PrimaryKeyConstraint("workspace", "name", name="registered_model_pk"),) + __table_args__ = ( + PrimaryKeyConstraint("workspace", "name", name="registered_model_pk"), + Index("idx_registered_models_workspace", "workspace"), + )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (18)
docs/api_reference/api_inventory.txt(2 hunks)mlflow/entities/__init__.py(2 hunks)mlflow/entities/workspace.py(1 hunks)mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py(1 hunks)mlflow/store/model_registry/dbmodels/models.py(11 hunks)mlflow/store/model_registry/sqlalchemy_store.py(6 hunks)mlflow/store/tracking/dbmodels/models.py(5 hunks)mlflow/store/workspace/dbmodels/__init__.py(1 hunks)mlflow/store/workspace/dbmodels/models.py(1 hunks)mlflow/utils/workspace_utils.py(1 hunks)tests/db/check_migration.py(3 hunks)tests/db/schemas/mssql.sql(6 hunks)tests/db/schemas/mysql.sql(6 hunks)tests/db/schemas/postgresql.sql(6 hunks)tests/db/schemas/sqlite.sql(6 hunks)tests/db/test_schema.py(3 hunks)tests/db/test_workspace_migration.py(1 hunks)tests/resources/db/latest_schema.sql(6 hunks)
🚧 Files skipped from review as they are similar to previous changes (4)
- mlflow/entities/workspace.py
- mlflow/store/workspace/dbmodels/init.py
- mlflow/store/workspace/dbmodels/models.py
- docs/api_reference/api_inventory.txt
🧰 Additional context used
🧬 Code graph analysis (6)
mlflow/entities/__init__.py (1)
mlflow/entities/workspace.py (1)
Workspace(9-13)
tests/db/test_schema.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
_get_unique_constraints(144-159)
tests/db/test_workspace_migration.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (2)
upgrade(138-615)downgrade(618-901)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
tests/db/test_schema.py (1)
_get_unique_constraints(90-107)
mlflow/store/model_registry/sqlalchemy_store.py (1)
mlflow/store/model_registry/dbmodels/models.py (3)
SqlRegisteredModelTag(155-190)SqlModelVersionTag(193-230)SqlRegisteredModelAlias(233-268)
tests/db/check_migration.py (1)
mlflow/store/model_registry/sqlalchemy_store.py (1)
create_registered_model(172-214)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (34)
- GitHub Check: pyfunc (1)
- GitHub Check: pyfunc (4)
- GitHub Check: pyfunc (2)
- GitHub Check: pyfunc (3)
- GitHub Check: evaluate (2)
- GitHub Check: evaluate (1)
- GitHub Check: genai
- GitHub Check: models (2)
- GitHub Check: windows (2)
- GitHub Check: windows (3)
- GitHub Check: flavors
- GitHub Check: models (1)
- GitHub Check: windows (1)
- GitHub Check: windows (4)
- GitHub Check: python (3)
- GitHub Check: python (2)
- GitHub Check: python (1)
- GitHub Check: java
- GitHub Check: database
- GitHub Check: python-skinny
- GitHub Check: typescript-sdk (22)
- GitHub Check: typescript-sdk (24)
- GitHub Check: typescript-sdk (20)
- GitHub Check: build
- GitHub Check: r
- GitHub Check: js (ubuntu-latest, --testPathIgnorePatterns)
- GitHub Check: js (windows-latest, --testPathPattern)
- GitHub Check: js (windows-latest, --testPathIgnorePatterns)
- GitHub Check: js (ubuntu-latest, --testPathPattern)
- GitHub Check: build (tracing)
- GitHub Check: build (skinny)
- GitHub Check: build (dev)
- GitHub Check: lint
- GitHub Check: core
🔇 Additional comments (16)
mlflow/entities/__init__.py (1)
65-65: Workspace re-export is consistent with existing entities API.Importing
Workspaceand adding it to__all__nicely aligns with the other entity exports and enablesfrom mlflow.entities import Workspaceas expected. No issues from this change.Also applies to: 129-129
tests/db/test_schema.py (1)
2-3: Implementation logic is sound; verify test coverage on mysql/mssql in CI.The code changes implement proper dialect-specific handling:
_DIALECT_REFLECTED_UNIQUE_CONSTRAINTScorrectly limits reattachment touq_experiments_workspace_nameon mysql/mssql._get_unique_constraints()properly handlesNotImplementedErrorfrominspector.get_unique_constraints()and falls back to filteringget_indexes()for unique constraints.- MySQL-specific guard (line 67) checks
duplicates_indexto distinguish between unique constraints and unique indexes.- Guards against missing columns (lines 75-79) and duplicate constraints (line 73).
- Constraint name matches the model definition.
Existing tests in
test_workspace_migration.py(lines 692–695, 778–780) already validateinspector.get_unique_constraints("experiments")behavior on mysql/mssql dialects, confirming the inspector returns the expected constraint metadata.However,
test_schema_is_up_to_date()is not parameterized by dialect in the current codebase—it runs only on the default database connection. To fully verify the requested concern, manually confirm in CI that either:
test_schema_is_up_to_dateexecutes on mysql/mssql database containers and successfully restores the constraint, or- Add dialect parameterization to the test to ensure all three dialects are covered.
tests/db/schemas/mysql.sql (2)
19-45: MySQL schema matches workspace changes; confirm index expectationsThe MySQL DDL for
evaluation_datasets,experiments(includinguq_experiments_workspace_name),registered_models, and the model‑registry tables reflects the new workspace columns, composite PKs, and workspace‑qualified FKs introduced in the migration.The migration also creates several workspace indexes (
idx_experiments_workspace,idx_registered_models_workspace,idx_evaluation_datasets_workspace,idx_webhooks_workspace). If yourcheck_migration/ schema‑parity tests assert on indexes using this file as ground truth, it would be good to ensure those indexes are represented here as well (either inline viaKEYclauses or asCREATE INDEXstatements), so the snapshot matches a migrated database.
81-88: Workspace‑qualified PK/FKs for registry tables look consistentFor
registered_models,model_versions,registered_model_aliases,registered_model_tags, andmodel_version_tags, the MySQL schema now:
- adds
workspace VARCHAR(63) DEFAULT 'default' NOT NULL,- promotes
(workspace, name[ , version ])to the primary key, and- updates FKs to reference
registered_models/model_versionsvia their workspace‑qualified keys.This lines up with the ORM models and the Alembic migration. No issues spotted here; just the same note as above about keeping any index definitions in sync if schema tests rely on them.
Also applies to: 183-200, 203-220, 367-375
tests/db/schemas/sqlite.sql (2)
19-45: SQLite schema’s workspace additions and constraints align with the migrationThe SQLite schema correctly introduces
workspaceonevaluation_datasetsandexperiments, with:
workspace VARCHAR(63) DEFAULT 'default' NOT NULL, anduq_experiments_workspace_nameenforcing per‑workspace uniqueness of experiment names.This matches the ORM definition and the upgrade migration behavior.
81-88: Registry workspace PK/FKs in SQLite look correct; consider index coverage in snapshotsFor the registry tables, the schema updates:
- add non‑nullable
workspacewith default'default',- switch PKs to
(workspace, name)/(workspace, name, version)/(workspace, key, name[ , version ]), and- wire FKs through those workspace‑qualified keys using the same constraint names as the migration.
That’s exactly what the downgrade/upgrade tests expect. As with MySQL, if your schema‑parity tests ever start asserting on indexes, you may want to ensure the workspace indexes created by the migration are also reflected in the SQLite snapshot.
Also applies to: 184-201, 204-221, 367-375
tests/resources/db/latest_schema.sql (2)
19-45: “Latest schema” snapshot matches the new workspace-aware logical schemaThe generic
latest_schema.sqlcorrectly reflects:
workspaceonevaluation_datasetsandexperiments,- the
uq_experiments_workspace_nameunique constraint, and- the non‑nullable default
'default'behavior.This keeps the logical schema in sync with what the migration produces and what the ORM models express.
81-88: Registry workspace keys andworkspacescatalog are consistent across snapshotsThe additions of:
- workspace‑qualified PKs/FKs for
registered_models,model_versions,registered_model_aliases,registered_model_tags, andmodel_version_tags, and- the
workspacestable withworkspaces_pkonnameclosely mirror the dialect‑specific schemas and the Alembic migration. This looks internally consistent and should be a reliable target for schema‑parity checks.
As with the dialect schemas, if your tests ever verify indexes as part of “latest schema” parity (e.g.,
idx_experiments_workspace), be sure to add those here too to avoid divergence between migrated DBs and this snapshot.Also applies to: 184-201, 204-221, 367-375, 106-111
tests/db/schemas/mssql.sql (2)
19-45: MSSQL workspace columns and experiment uniqueness match the migration logicFor MSSQL, the additions of:
workspace VARCHAR(63) ... DEFAULT ('default') NOT NULLonevaluation_datasetsandexperiments, anduq_experiments_workspace_nameunique on(workspace, name)look correct and compatible with the catalog queries in
_fetch_mssql_unique_metadata()that detect unique constraints onname. This should keep the migration’s constraint‑detection logic and the expected schema in sync.
80-87: MSSQL registry workspace PK/FKs and workspaces table are consistentThe MSSQL schema updates:
- add
workspacewith default'default'acrossregistered_models,model_versions,registered_model_aliases,registered_model_tags, andmodel_version_tags,- convert PKs to workspace‑qualified forms (e.g.,
(workspace, name, version)), and- update FKs to reference the new composite keys with constraint names used in the migration (e.g.,
fk_model_versions_registered_models,fk_model_version_tags_model_versions),plus the new
workspacestable withworkspaces_pkonname. This matches both the Alembic migration and the ORM definitions and should behave correctly under the MSSQL‑specific migration paths.Also applies to: 181-198, 201-218, 359-367, 105-109
mlflow/store/model_registry/dbmodels/models.py (6)
1-1: LGTM!The new imports are correctly used throughout the file to support workspace multi-tenancy.
Also applies to: 36-36
220-223: LGTM!The explicit
model_versionrelationship improves code clarity and enables bidirectional navigation between model versions and their tags.
252-261: LGTM!The foreign key constraint correctly includes both
onupdate="cascade"andondelete="cascade", ensuring proper referential integrity at the database level. This is the correct pattern that should be applied to other tables as well.
301-321: Verify the primary key strategy for workspace isolation.Unlike other tables in this file,
SqlWebhookuses a single-column primary key (webhook_id) rather than a composite key includingworkspace. This means:
webhook_idmust be globally unique across all workspaces- The design pattern is inconsistent with
SqlRegisteredModel,SqlModelVersion, etc.If
webhook_idis generated as a UUID or similar globally-unique identifier, this is acceptable. However, if there's any possibility of workspace-scoped collision, or if consistency with the rest of the schema is preferred, consider making the primary key(workspace, webhook_id).Please verify whether
webhook_idis guaranteed to be globally unique or whether the primary key should be updated to(workspace, webhook_id)for consistency.
344-367: SqlWebhookEvent design depends on webhook_id uniqueness.
SqlWebhookEventdoes not have aworkspacecolumn and referenceswebhooks.webhook_iddirectly. This design is acceptable ifwebhook_idis globally unique across workspaces. However, if the primary key ofSqlWebhookis changed to a composite key(workspace, webhook_id), thenSqlWebhookEventwould need:
- A
workspacecolumn- A composite foreign key constraint on
(workspace, webhook_id)This is directly related to the design decision for
SqlWebhookflagged above.
65-82: Code is correct as-is; no changes needed.The
RegisteredModelentity class does not have aworkspaceparameter in its__init__method, and this is intentional. Workspace is a database-level tenancy concern managed by the ORM models (e.g.,SqlRegisteredModel), while the entity layer represents domain objects independent of infrastructure details. This pattern is consistent across all entity conversions in the codebase.
| __table_args__ = ( | ||
| ForeignKeyConstraint( | ||
| ["workspace", "name"], | ||
| ["registered_models.workspace", "registered_models.name"], | ||
| onupdate="cascade", | ||
| ), | ||
| PrimaryKeyConstraint("workspace", "name", "version", name="model_version_pk"), | ||
| ) |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Add ondelete="cascade" to the foreign key constraint for consistency.
The ForeignKeyConstraint includes onupdate="cascade" but lacks ondelete="cascade". While the ORM-level relationship (line 123) specifies cascade="all", having database-level cascade deletion ensures referential integrity even for direct SQL operations.
Apply this diff:
__table_args__ = (
ForeignKeyConstraint(
["workspace", "name"],
["registered_models.workspace", "registered_models.name"],
onupdate="cascade",
+ ondelete="cascade",
),
PrimaryKeyConstraint("workspace", "name", "version", name="model_version_pk"),
)🤖 Prompt for AI Agents
In mlflow/store/model_registry/dbmodels/models.py around lines 126 to 133, the
ForeignKeyConstraint sets onupdate="cascade" but omits ondelete="cascade";
update the ForeignKeyConstraint to include ondelete="cascade" so the DB enforces
cascading deletes too (keep onupdate="cascade" as-is) and ensure the constraint
signature remains valid with both options specified.
| __table_args__ = ( | ||
| ForeignKeyConstraint( | ||
| ["workspace", "name"], | ||
| ["registered_models.workspace", "registered_models.name"], | ||
| onupdate="cascade", | ||
| ), | ||
| PrimaryKeyConstraint("workspace", "key", "name", name="registered_model_tag_pk"), | ||
| ) |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Add ondelete="cascade" to maintain consistency.
Similar to SqlModelVersion, this foreign key constraint should include ondelete="cascade" for database-level referential integrity.
Apply this diff:
__table_args__ = (
ForeignKeyConstraint(
["workspace", "name"],
["registered_models.workspace", "registered_models.name"],
onupdate="cascade",
+ ondelete="cascade",
),
PrimaryKeyConstraint("workspace", "key", "name", name="registered_model_tag_pk"),
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| __table_args__ = ( | |
| ForeignKeyConstraint( | |
| ["workspace", "name"], | |
| ["registered_models.workspace", "registered_models.name"], | |
| onupdate="cascade", | |
| ), | |
| PrimaryKeyConstraint("workspace", "key", "name", name="registered_model_tag_pk"), | |
| ) | |
| __table_args__ = ( | |
| ForeignKeyConstraint( | |
| ["workspace", "name"], | |
| ["registered_models.workspace", "registered_models.name"], | |
| onupdate="cascade", | |
| ondelete="cascade", | |
| ), | |
| PrimaryKeyConstraint("workspace", "key", "name", name="registered_model_tag_pk"), | |
| ) |
🤖 Prompt for AI Agents
In mlflow/store/model_registry/dbmodels/models.py around lines 176 to 183, the
ForeignKeyConstraint for (workspace, name) referencing registered_models lacks
ondelete="cascade"; update that ForeignKeyConstraint to include
ondelete="cascade" (matching SqlModelVersion) so deletions of registered_models
cascade to this table, leaving the PrimaryKeyConstraint unchanged.
| __table_args__ = ( | ||
| PrimaryKeyConstraint("key", "name", "version", name="model_version_tag_pk"), | ||
| ForeignKeyConstraint( | ||
| ("name", "version"), | ||
| ("model_versions.name", "model_versions.version"), | ||
| ["workspace", "name", "version"], | ||
| ["model_versions.workspace", "model_versions.name", "model_versions.version"], | ||
| onupdate="cascade", | ||
| ), | ||
| PrimaryKeyConstraint("workspace", "key", "name", "version", name="model_version_tag_pk"), | ||
| ) |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Add ondelete="cascade" for referential integrity.
The foreign key constraint should include ondelete="cascade" to ensure tags are removed when the parent model version is deleted, maintaining consistency with the ORM-level cascade behavior.
Apply this diff:
__table_args__ = (
ForeignKeyConstraint(
["workspace", "name", "version"],
["model_versions.workspace", "model_versions.name", "model_versions.version"],
onupdate="cascade",
+ ondelete="cascade",
),
PrimaryKeyConstraint("workspace", "key", "name", "version", name="model_version_tag_pk"),
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| __table_args__ = ( | |
| PrimaryKeyConstraint("key", "name", "version", name="model_version_tag_pk"), | |
| ForeignKeyConstraint( | |
| ("name", "version"), | |
| ("model_versions.name", "model_versions.version"), | |
| ["workspace", "name", "version"], | |
| ["model_versions.workspace", "model_versions.name", "model_versions.version"], | |
| onupdate="cascade", | |
| ), | |
| PrimaryKeyConstraint("workspace", "key", "name", "version", name="model_version_tag_pk"), | |
| ) | |
| __table_args__ = ( | |
| ForeignKeyConstraint( | |
| ["workspace", "name", "version"], | |
| ["model_versions.workspace", "model_versions.name", "model_versions.version"], | |
| onupdate="cascade", | |
| ondelete="cascade", | |
| ), | |
| PrimaryKeyConstraint("workspace", "key", "name", "version", name="model_version_tag_pk"), | |
| ) |
🤖 Prompt for AI Agents
In mlflow/store/model_registry/dbmodels/models.py around lines 211 to 218, the
ForeignKeyConstraint linking tags to model_versions lacks ondelete="cascade", so
database-level deletions of a model_version won't cascade to its tags; update
the ForeignKeyConstraint to include ondelete="cascade" (i.e., add
ondelete="cascade" alongside the existing onupdate="cascade") so that DB-level
referential integrity matches the ORM cascade behavior and tags are removed when
their parent model_version is deleted.
0446441 to
72bfda2
Compare
There was a problem hiding this comment.
Actionable comments posted: 5
♻️ Duplicate comments (3)
mlflow/store/model_registry/dbmodels/models.py (3)
126-133: Past review comment not yet addressed: Missingondelete="cascade"in ForeignKeyConstraint.The past review comment correctly identified that this foreign key constraint lacks
ondelete="cascade". While the ORM relationship (line 123) hascascade="all", database-level cascade ensures referential integrity for direct SQL operations.Apply this diff:
__table_args__ = ( ForeignKeyConstraint( ["workspace", "name"], ["registered_models.workspace", "registered_models.name"], onupdate="cascade", + ondelete="cascade", ), PrimaryKeyConstraint("workspace", "name", "version", name="model_version_pk"), )Based on learnings from past review comments.
176-183: Past review comment not yet addressed: Missingondelete="cascade"in ForeignKeyConstraint.The past review comment correctly identified this missing constraint.
SqlRegisteredModelAlias(lines 252-261) shows the correct pattern with bothonupdate="cascade"andondelete="cascade".Apply this diff:
__table_args__ = ( ForeignKeyConstraint( ["workspace", "name"], ["registered_models.workspace", "registered_models.name"], onupdate="cascade", + ondelete="cascade", ), PrimaryKeyConstraint("workspace", "key", "name", name="registered_model_tag_pk"), )Based on learnings from past review comments.
211-218: Past review comment not yet addressed: Missingondelete="cascade"in ForeignKeyConstraint.The past review comment correctly identified this issue. Tags should be removed when the parent model version is deleted.
Apply this diff:
__table_args__ = ( ForeignKeyConstraint( ["workspace", "name", "version"], ["model_versions.workspace", "model_versions.name", "model_versions.version"], onupdate="cascade", + ondelete="cascade", ), PrimaryKeyConstraint("workspace", "key", "name", "version", name="model_version_tag_pk"), )Based on learnings from past review comments.
🧹 Nitpick comments (4)
mlflow/store/tracking/dbmodels/models.py (1)
1315-1323: Clarify workspace handling for evaluation datasets and consider composite indexThe added
workspacecolumn plus a dedicated workspace index are a good starting point for workspace‑scoped evaluation datasets, and the default/server_default pattern mirrors the experiment change nicely.A couple of follow‑ups to consider:
Entity visibility:
SqlEvaluationDataset.to_mlflow_entity()/from_mlflow_entity()currently ignoreworkspace. If evaluation datasets are meant to be workspace‑aware at the entity/API level, we should threadworkspacethrough there; otherwise, it’s implicitly a store‑only dimension, similar to the current experiment behavior.Index/uniqueness shape: If typical queries look like
WHERE workspace = :ws AND name = :name(or includedigest), a composite index such as(workspace, name)could be more efficient than separate indexes onworkspaceandname. And if we logically expect a dataset name (or name+digest) to be unique per workspace, a correspondingUniqueConstraintwould help enforce that invariant and avoid race conditions.None of this blocks the current change but is worth aligning with expected access patterns and invariants.
Also applies to: 1375-1380
mlflow/entities/workspace.py (1)
8-13: Consider adding validation for the workspace name.The
Workspacedataclass lacks validation for thenamefield. Based on the database schema (VARCHAR(63)), the name should be validated for:
- Non-empty value
- Maximum length of 63 characters
- Potentially valid characters/format
Consider adding a
__post_init__method for validation:@dataclass(frozen=True, slots=True) class Workspace: """Minimal metadata describing a workspace.""" name: str description: str | None = None + + def __post_init__(self): + if not self.name or not self.name.strip(): + raise ValueError("Workspace name cannot be empty") + if len(self.name) > 63: + raise ValueError("Workspace name cannot exceed 63 characters")tests/db/test_schema.py (1)
90-107: Potential code duplication with migration logic.The
_get_unique_constraintshelper closely resembles the logic in the migration file atmlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py(lines 141-156). Both handle the same MSSQL fallback and index-based constraint detection.Consider extracting this shared logic to a common utility module (e.g.,
mlflow/store/db/utils.py) to avoid duplication and ensure consistent behavior across migration and test code.mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
133-134: TODO comment about removing server_default - track this technical debt.The comment indicates that
server_defaultshould be removed once the tracking and model registry stores properly set the workspace. This is important to ensure data integrity and prevent silent fallback to 'default' workspace.Please create a follow-up task or issue to:
- Update the stores to explicitly set workspace values
- Remove the
server_defaultfrom the migration- Verify all code paths properly handle workspace assignment
Would you like me to open an issue to track this work?
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (18)
docs/api_reference/api_inventory.txt(2 hunks)mlflow/entities/__init__.py(2 hunks)mlflow/entities/workspace.py(1 hunks)mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py(1 hunks)mlflow/store/model_registry/dbmodels/models.py(11 hunks)mlflow/store/model_registry/sqlalchemy_store.py(6 hunks)mlflow/store/tracking/dbmodels/models.py(5 hunks)mlflow/store/workspace/dbmodels/__init__.py(1 hunks)mlflow/store/workspace/dbmodels/models.py(1 hunks)mlflow/utils/workspace_utils.py(1 hunks)tests/db/check_migration.py(3 hunks)tests/db/schemas/mssql.sql(6 hunks)tests/db/schemas/mysql.sql(6 hunks)tests/db/schemas/postgresql.sql(6 hunks)tests/db/schemas/sqlite.sql(6 hunks)tests/db/test_schema.py(3 hunks)tests/db/test_workspace_migration.py(1 hunks)tests/resources/db/latest_schema.sql(6 hunks)
🚧 Files skipped from review as they are similar to previous changes (6)
- mlflow/utils/workspace_utils.py
- tests/db/check_migration.py
- mlflow/entities/init.py
- mlflow/store/workspace/dbmodels/init.py
- docs/api_reference/api_inventory.txt
- mlflow/store/workspace/dbmodels/models.py
🧰 Additional context used
🧬 Code graph analysis (4)
tests/db/test_workspace_migration.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (2)
upgrade(135-612)downgrade(615-945)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
tests/db/test_schema.py (1)
_get_unique_constraints(90-107)
tests/db/test_schema.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
_get_unique_constraints(141-156)
mlflow/store/model_registry/sqlalchemy_store.py (1)
mlflow/store/model_registry/dbmodels/models.py (3)
SqlRegisteredModelTag(155-190)SqlModelVersionTag(193-230)SqlRegisteredModelAlias(233-268)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (34)
- GitHub Check: flavors
- GitHub Check: evaluate (2)
- GitHub Check: models (1)
- GitHub Check: genai
- GitHub Check: evaluate (1)
- GitHub Check: models (2)
- GitHub Check: java
- GitHub Check: pyfunc (4)
- GitHub Check: windows (1)
- GitHub Check: windows (2)
- GitHub Check: windows (4)
- GitHub Check: windows (3)
- GitHub Check: pyfunc (3)
- GitHub Check: database
- GitHub Check: pyfunc (1)
- GitHub Check: pyfunc (2)
- GitHub Check: python (1)
- GitHub Check: python (2)
- GitHub Check: python-skinny
- GitHub Check: python (3)
- GitHub Check: typescript-sdk (22)
- GitHub Check: typescript-sdk (24)
- GitHub Check: typescript-sdk (20)
- GitHub Check: build (skinny)
- GitHub Check: build (tracing)
- GitHub Check: build (dev)
- GitHub Check: js (windows-latest, --testPathPattern)
- GitHub Check: js (ubuntu-latest, --testPathPattern)
- GitHub Check: core
- GitHub Check: js (ubuntu-latest, --testPathIgnorePatterns)
- GitHub Check: js (windows-latest, --testPathIgnorePatterns)
- GitHub Check: lint-macos
- GitHub Check: build
- GitHub Check: r
🔇 Additional comments (12)
mlflow/store/tracking/dbmodels/models.py (2)
64-64: ConfirmDEFAULT_WORKSPACE_NAMEimport doesn’t create circular dependenciesImporting
DEFAULT_WORKSPACE_NAMEhere is reasonable, but please double-check thatmlflow.utils.workspace_utilsdoes not, directly or indirectly, importmlflow.store.tracking.dbmodels.modelsto avoid circular import issues.
98-111: Workspace‑scoped experiments look consistent; verify entity & migration alignmentThe new
workspacecolumn plusUniqueConstraint("workspace", "name", name="uq_experiments_workspace_name")is a clean way to move from globally unique names to workspace‑scoped experiments. The combination ofdefault+server_defaultonworkspaceshould also make migrations/backfills for existing rows straightforward.Two things to confirm:
Experiment entity surface:
SqlExperiment.to_mlflow_entity()still doesn’t exposeworkspace. If theExperimententity class (or surrounding APIs) is intended to carry workspace information, we’ll need to plumbworkspacethrough here; if workspace is intentionally a store‑only dimension, this is fine but worth making explicit.Migration/constraint naming: Make sure the Alembic (or equivalent) migration that adds
workspaceand drops the old global unique constraint uses the same constraint nameuq_experiments_workspace_name, so upgrade/downgrade paths and CI checks don’t drift from the model definition.Also applies to: 131-138
tests/db/schemas/postgresql.sql (1)
29-29: LGTM! Workspace schema changes are well-structured.The workspace column additions and constraint updates are consistent across all affected tables:
- Proper NOT NULL constraints with sensible defaults
- Composite primary keys correctly include workspace
- Foreign key constraints properly reference workspace-qualified parent keys
- Appropriate CASCADE behavior for updates and deletes
Also applies to: 41-41, 86-86, 101-101, 106-110, 198-200, 208-210, 218-220, 370-372
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (2)
626-646: LGTM! Conflict detection prevents data loss during downgrade.The
_assert_no_workspace_conflictsfunction properly validates that downgrading (which merges all workspaces into 'default') won't create duplicate keys. The error messages are clear and actionable.
206-412: I need to gather more information about the actual test coverage for SQLite migrations.Now let me search more specifically for test files related to this migration:Now let me search the codebase for actual test files related to this migration:
Based on my verification through documentation and code analysis, I've gathered sufficient information. The original review comment asks to verify SQLite batch operations but the requested verification commands had limited utility. Let me provide the rewritten review comment:
SQLite batch operations appear correctly implemented but verification incomplete.
The migration properly uses batch operations with constraint handling, which is the standard pattern for SQLite. With batch operations set to "auto", SQLite will recreate tables when operations other than add_column(), create_index(), or drop_index() are present, and this migration performs multiple constraint modifications requiring recreation.
However, the initial verification search (
rg -n "dialect.*sqlite" tests/db/test_workspace_migration.py) found no results, andUSE_EXTERNAL_DBindicates tests prefer external databases over SQLite. Without access to comprehensive SQLite-specific test coverage verification, manual testing of this migration against SQLite is recommended to confirm all constraint transitions (PK/FK drops, recreations, and the new workspace column additions) execute successfully.tests/db/test_workspace_migration.py (3)
160-342: Excellent use of raw SQL for pre-migration seeding.Using raw SQL to insert pre-migration data (without workspace columns) ensures the migration is truly responsible for adding and backfilling these fields. This is a best practice for testing migrations.
704-829: Comprehensive upgrade test coverage.The test verifies:
- Workspace column presence and defaults
- Data backfilling to 'default' workspace
- Primary key updates
- Foreign key relationships
- Unique constraints
- Index creation
- Workspace catalog seeding
795-806: MSSQL constraint reflection workaround is appropriate.The test correctly handles MSSQL's
NotImplementedErrorforget_unique_constraintsby catching the exception and settingunique_experiments = None. This aligns with the migration's MSSQL-specific handling.mlflow/store/model_registry/sqlalchemy_store.py (1)
198-198: Workspace is hardcoded to DEFAULT_WORKSPACE_NAME.All entity creation operations use
DEFAULT_WORKSPACE_NAMEas a hardcoded value. This is appropriate for the current phase of the migration but indicates that multi-workspace support is not yet implemented at the API level.Verify that:
- There's a plan to add workspace parameter to the public API methods
- The hardcoded workspace aligns with the migration's server_default behavior
- Future PRs will enable dynamic workspace selection
This is noted in the migration file's TODO comment about removing server_default once stores support workspaces properly.
Based on learnings
Also applies to: 204-204, 704-704, 813-813, 819-819, 1206-1206, 1268-1268
tests/db/schemas/mysql.sql (1)
29-29: LGTM! MySQL schema matches PostgreSQL structure.The workspace changes are consistent across dialects:
- Same column definitions (VARCHAR(63), DEFAULT 'default', NOT NULL)
- Same constraint structure (composite PKs, workspace-qualified FKs)
- Proper MySQL-specific syntax (
_utf8mb4charset, CHECK constraint format)Also applies to: 41-44, 86-87, 106-110, 198-200, 208-210, 218-220, 372-374
tests/db/schemas/sqlite.sql (1)
29-29: LGTM! SQLite schema is consistent with other dialects.The workspace schema changes match PostgreSQL and MySQL:
- Identical column definitions and defaults
- Same composite key structure
- SQLite-appropriate syntax for constraints
All three dialect schemas (PostgreSQL, MySQL, SQLite) are now consistent with the workspace multi-tenancy model.
Also applies to: 41-43, 86-87, 106-110, 199-200, 209-210, 219-220, 373-374
mlflow/store/model_registry/dbmodels/models.py (1)
252-261: LGTM: Correct FK constraint with bothonupdateandondeletecascade.This
SqlRegisteredModelAliasFK constraint correctly includes bothonupdate="cascade"andondelete="cascade", ensuring proper referential integrity at the database level. This is the pattern that should be followed by the other models.
| workspace = Column( | ||
| String(63), | ||
| nullable=False, | ||
| default=DEFAULT_WORKSPACE_NAME, | ||
| server_default=sa.text(f"'{DEFAULT_WORKSPACE_NAME}'"), | ||
| ) |
There was a problem hiding this comment.
🧩 Analysis chain
Critical: Add foreign key constraints from workspace columns to workspaces table.
All workspace columns across SqlRegisteredModel, SqlModelVersion, SqlRegisteredModelTag, SqlModelVersionTag, SqlRegisteredModelAlias, and SqlWebhook lack foreign key constraints to a workspaces table. This mirrors the schema file issues and allows invalid workspace values.
Verify that a SqlWorkspace model exists and add FK constraints:
If SqlWorkspace exists, add FK constraints. For example, in SqlRegisteredModel.__table_args__:
__table_args__ = (
ForeignKeyConstraint(
["workspace"],
["workspaces.name"],
name="fk_registered_models_workspace",
),
PrimaryKeyConstraint("workspace", "name", name="registered_model_pk"),
)Apply similar changes to SqlWebhook and any other root-level models with workspace columns. Child models like SqlModelVersion inherit workspace validation through their composite FKs to parent tables.
Also applies to: 88-93, 158-163, 196-201, 236-241, 301-306
🏁 Script executed:
#!/bin/bash
# Search for SqlWorkspace model definition
ast-grep --pattern 'class SqlWorkspace($$$)'
# Search for workspace-related models and imports
rg -n "SqlWorkspace|class.*Workspace.*Base" --type=pyLength of output: 1377
I found that SqlWorkspace exists. Now I need to verify whether FK constraints have already been added to the workspace columns in the models mentioned in the review comment.
Now I need to directly inspect the models file to check if FK constraints have been added:
🏁 Script executed:
#!/bin/bash
# Extract workspace column and __table_args__ definitions from all relevant models
# Check SqlRegisteredModel, SqlModelVersion, SqlRegisteredModelTag, SqlModelVersionTag,
# SqlRegisteredModelAlias, and SqlWebhook
cat mlflow/store/model_registry/dbmodels/models.py | head -350 | tail -100
echo "---"
# Search for __table_args__ definitions containing ForeignKeyConstraint
rg -n "ForeignKeyConstraint|__table_args__" mlflow/store/model_registry/dbmodels/models.py -A 5 -B 2Length of output: 6137
Critical: FK constraints to workspaces table are missing from SqlRegisteredModel and SqlWebhook.
Verification confirms that SqlWorkspace exists (mlflow/store/workspace/dbmodels/models.py:10) with the workspaces table keyed on "name". However, FK constraints are absent from root-level models:
SqlRegisteredModel(line 57):__table_args__contains only PrimaryKeyConstraint, no FK to workspaces tableSqlWebhook(lines 317-322):__table_args__contains only PrimaryKeyConstraint and Indexes, no FK to workspaces table
While child models have composite FKs to their parents, these do not enforce that workspace values exist in the workspaces table. Add direct FK constraints in:
- SqlRegisteredModel.table_args (line 57): Add constraint
ForeignKeyConstraint(["workspace"], ["workspaces.name"], name="fk_registered_models_workspace") - SqlWebhook.table_args (lines 317-322): Add constraint
ForeignKeyConstraint(["workspace"], ["workspaces.name"], name="fk_webhooks_workspace")
Child models' workspace validation will cascade through their composite FKs to parents once root models enforce workspace existence.
🤖 Prompt for AI Agents
In mlflow/store/model_registry/dbmodels/models.py around lines 42–47 and
specifically update SqlRegisteredModel.__table_args__ at ~line 57 and
SqlWebhook.__table_args__ at ~lines 317–322: add a ForeignKeyConstraint on the
"workspace" column referencing workspaces.name for each root model (name them
fk_registered_models_workspace and fk_webhooks_workspace respectively) so the
root models enforce that workspace values exist in the workspaces table; ensure
you insert the ForeignKeyConstraint entries into the existing __table_args__
tuple alongside the current PrimaryKeyConstraint/Indexes without removing them.
| CREATE TABLE workspaces ( | ||
| name VARCHAR(63) COLLATE "SQL_Latin1_General_CP1_CI_AS" NOT NULL, | ||
| description VARCHAR COLLATE "SQL_Latin1_General_CP1_CI_AS", | ||
| CONSTRAINT workspaces_pk PRIMARY KEY (name) | ||
| ) |
There was a problem hiding this comment.
Critical: Add foreign key constraints from workspace columns to the workspaces table.
Same issue as in latest_schema.sql: the MSSQL schema lacks foreign key constraints from workspace columns to the workspaces catalog table. This allows invalid workspace references.
Add FK constraints for MSSQL dialect:
-- Add after evaluation_datasets table definition
ALTER TABLE evaluation_datasets ADD CONSTRAINT fk_evaluation_datasets_workspace
FOREIGN KEY(workspace) REFERENCES workspaces(name);
-- Add after experiments table definition
ALTER TABLE experiments ADD CONSTRAINT fk_experiments_workspace
FOREIGN KEY(workspace) REFERENCES workspaces(name);
-- Add after registered_models table definition
ALTER TABLE registered_models ADD CONSTRAINT fk_registered_models_workspace
FOREIGN KEY(workspace) REFERENCES workspaces(name);
-- Add after webhooks table definition
ALTER TABLE webhooks ADD CONSTRAINT fk_webhooks_workspace
FOREIGN KEY(workspace) REFERENCES workspaces(name);🤖 Prompt for AI Agents
In tests/db/schemas/mssql.sql around lines 105 to 109 the workspaces table is
defined but there are no foreign key constraints from other tables' workspace
columns to workspaces(name); add ALTER TABLE statements after each relevant
table definition (evaluation_datasets, experiments, registered_models, webhooks)
to create FK constraints (e.g. fk_evaluation_datasets_workspace,
fk_experiments_workspace, fk_registered_models_workspace, fk_webhooks_workspace)
that reference workspaces(name) so workspace columns cannot contain invalid
values.
| CONSTRAINT model_version_pk PRIMARY KEY (workspace, name, version), | ||
| CONSTRAINT fk_model_versions_registered_models FOREIGN KEY(workspace, name) REFERENCES registered_models (workspace, name) ON UPDATE CASCADE | ||
| ) |
There was a problem hiding this comment.
Major: Inconsistent ON DELETE CASCADE in MSSQL schema.
The MSSQL schema has the same inconsistent cascade behavior as latest_schema.sql:
registered_model_aliases(line 207) hasON DELETE CASCADEmodel_versions(line 197),registered_model_tags(line 217), andmodel_version_tags(line 366) lackON DELETE CASCADE
Apply the same fix as in latest_schema.sql:
workspace VARCHAR(63) COLLATE "SQL_Latin1_General_CP1_CI_AS" DEFAULT ('default') NOT NULL,
CONSTRAINT model_version_pk PRIMARY KEY (workspace, name, version),
- CONSTRAINT fk_model_versions_registered_models FOREIGN KEY(workspace, name) REFERENCES registered_models (workspace, name) ON UPDATE CASCADE
+ CONSTRAINT fk_model_versions_registered_models FOREIGN KEY(workspace, name) REFERENCES registered_models (workspace, name) ON DELETE CASCADE ON UPDATE CASCADE
)
...
workspace VARCHAR(63) COLLATE "SQL_Latin1_General_CP1_CI_AS" DEFAULT ('default') NOT NULL,
CONSTRAINT registered_model_tag_pk PRIMARY KEY (workspace, key, name),
- CONSTRAINT fk_registered_model_tags_registered_models FOREIGN KEY(workspace, name) REFERENCES registered_models (workspace, name) ON UPDATE CASCADE
+ CONSTRAINT fk_registered_model_tags_registered_models FOREIGN KEY(workspace, name) REFERENCES registered_models (workspace, name) ON DELETE CASCADE ON UPDATE CASCADE
)
...
workspace VARCHAR(63) COLLATE "SQL_Latin1_General_CP1_CI_AS" DEFAULT ('default') NOT NULL,
CONSTRAINT model_version_tag_pk PRIMARY KEY (workspace, key, name, version),
- CONSTRAINT fk_model_version_tags_model_versions FOREIGN KEY(workspace, name, version) REFERENCES model_versions (workspace, name, version) ON UPDATE CASCADE
+ CONSTRAINT fk_model_version_tags_model_versions FOREIGN KEY(workspace, name, version) REFERENCES model_versions (workspace, name, version) ON DELETE CASCADE ON UPDATE CASCADE
)Also applies to: 206-208, 216-218, 365-367
🤖 Prompt for AI Agents
In tests/db/schemas/mssql.sql around lines 196-198 (and also update 206-208,
216-218, 365-367), the FOREIGN KEY constraints for model_versions,
registered_model_tags, and model_version_tags are missing ON DELETE CASCADE
causing inconsistent cascade behavior vs registered_model_aliases; modify each
referenced FOREIGN KEY ... REFERENCES ... clauses to append ON DELETE CASCADE
(matching latest_schema.sql) so deletions of registered_models or model_versions
cascade to their dependent rows.
| CREATE TABLE workspaces ( | ||
| name VARCHAR(63) NOT NULL, | ||
| description TEXT, | ||
| CONSTRAINT workspaces_pk PRIMARY KEY (name) | ||
| ) |
There was a problem hiding this comment.
Critical: Add foreign key constraints from workspace columns to the workspaces table.
The new workspaces catalog table is defined, but none of the tables with workspace columns have foreign key constraints referencing it. This allows insertion of invalid workspace values (e.g., workspace='nonexistent') that don't exist in the workspaces table, defeating the purpose of the catalog.
Add FK constraints for these root tables:
evaluation_datasets.workspace → workspaces.nameexperiments.workspace → workspaces.nameregistered_models.workspace → workspaces.namewebhooks.workspace → workspaces.name
Child tables (model_versions, registered_model_aliases, etc.) inherit workspace validation through their composite FK constraints to parent tables.
-- Add after evaluation_datasets table definition
ALTER TABLE evaluation_datasets ADD CONSTRAINT fk_evaluation_datasets_workspace
FOREIGN KEY(workspace) REFERENCES workspaces(name);
-- Add after experiments table definition
ALTER TABLE experiments ADD CONSTRAINT fk_experiments_workspace
FOREIGN KEY(workspace) REFERENCES workspaces(name);
-- Add after registered_models table definition
ALTER TABLE registered_models ADD CONSTRAINT fk_registered_models_workspace
FOREIGN KEY(workspace) REFERENCES workspaces(name);
-- Add after webhooks table definition
ALTER TABLE webhooks ADD CONSTRAINT fk_webhooks_workspace
FOREIGN KEY(workspace) REFERENCES workspaces(name);🤖 Prompt for AI Agents
In tests/resources/db/latest_schema.sql around lines 106 to 110, the new
workspaces table is created but other root tables with a workspace column lack
foreign key constraints, allowing invalid workspace values; add ALTER TABLE
statements after the respective table definitions to add FK constraints: ALTER
TABLE evaluation_datasets ADD CONSTRAINT fk_evaluation_datasets_workspace
FOREIGN KEY(workspace) REFERENCES workspaces(name); ALTER TABLE experiments ADD
CONSTRAINT fk_experiments_workspace FOREIGN KEY(workspace) REFERENCES
workspaces(name); ALTER TABLE registered_models ADD CONSTRAINT
fk_registered_models_workspace FOREIGN KEY(workspace) REFERENCES
workspaces(name); ALTER TABLE webhooks ADD CONSTRAINT fk_webhooks_workspace
FOREIGN KEY(workspace) REFERENCES workspaces(name).
| CONSTRAINT model_version_pk PRIMARY KEY (workspace, name, version), | ||
| CONSTRAINT fk_model_versions_registered_models FOREIGN KEY(workspace, name) REFERENCES registered_models (workspace, name) ON UPDATE CASCADE | ||
| ) |
There was a problem hiding this comment.
Major: Inconsistent ON DELETE CASCADE behavior in foreign key constraints.
The foreign key constraints have inconsistent cascade deletion behavior:
- Line 210:
registered_model_aliaseshas bothON DELETE CASCADE ON UPDATE CASCADE✓ - Line 200:
model_versionshas onlyON UPDATE CASCADE(missing ON DELETE) - Line 220:
registered_model_tagshas onlyON UPDATE CASCADE(missing ON DELETE) - Line 374:
model_version_tagshas onlyON UPDATE CASCADE(missing ON DELETE)
This inconsistency means deleting a registered_model would cascade to aliases but leave orphaned tags and model versions, requiring manual cleanup. Either:
- Add
ON DELETE CASCADEto all these constraints for automatic cleanup, or - Document why some relationships should preserve child records on parent deletion (though this seems unlikely for model metadata)
Apply this diff for consistency:
workspace VARCHAR(63) DEFAULT 'default' NOT NULL,
CONSTRAINT model_version_pk PRIMARY KEY (workspace, name, version),
- CONSTRAINT fk_model_versions_registered_models FOREIGN KEY(workspace, name) REFERENCES registered_models (workspace, name) ON UPDATE CASCADE
+ CONSTRAINT fk_model_versions_registered_models FOREIGN KEY(workspace, name) REFERENCES registered_models (workspace, name) ON DELETE CASCADE ON UPDATE CASCADE
)
...
workspace VARCHAR(63) DEFAULT 'default' NOT NULL,
CONSTRAINT registered_model_tag_pk PRIMARY KEY (workspace, key, name),
- CONSTRAINT fk_registered_model_tags_registered_models FOREIGN KEY(workspace, name) REFERENCES registered_models (workspace, name) ON UPDATE CASCADE
+ CONSTRAINT fk_registered_model_tags_registered_models FOREIGN KEY(workspace, name) REFERENCES registered_models (workspace, name) ON DELETE CASCADE ON UPDATE CASCADE
)
...
workspace VARCHAR(63) DEFAULT 'default' NOT NULL,
CONSTRAINT model_version_tag_pk PRIMARY KEY (workspace, key, name, version),
- CONSTRAINT fk_model_version_tags_model_versions FOREIGN KEY(workspace, name, version) REFERENCES model_versions (workspace, name, version) ON UPDATE CASCADE
+ CONSTRAINT fk_model_version_tags_model_versions FOREIGN KEY(workspace, name, version) REFERENCES model_versions (workspace, name, version) ON DELETE CASCADE ON UPDATE CASCADE
)Also applies to: 209-211, 219-221, 373-375
🤖 Prompt for AI Agents
In tests/resources/db/latest_schema.sql around lines 199-201 (and also update
the related FK definitions at 209-211, 219-221, and 373-375), the foreign key
constraints lack ON DELETE CASCADE causing inconsistent cascade behavior; update
each referenced FOREIGN KEY(...) REFERENCES ... ON UPDATE CASCADE to include ON
DELETE CASCADE as well so that deleting a registered_models row will
automatically delete related model_versions, registered_model_tags,
model_version_tags (and any other child tables listed), ensuring consistent
cleanup across these relationships.
72bfda2 to
fea3308
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
mlflow/store/model_registry/sqlalchemy_store.py (1)
226-231: Critical: Query must filter by workspace to prevent ambiguity.The query filters by
nameonly (line 229) but the new schema has a composite primary key(workspace, name). Without filtering by workspace, this query can return multiple registered models with the same name from different workspaces, causing the error at lines 237-241.Add workspace filtering to the query:
query_options = cls._get_eager_registered_model_query_options() if eager else [] rms = ( session.query(SqlRegisteredModel) .options(*query_options) .filter(SqlRegisteredModel.name == name) + .filter(SqlRegisteredModel.workspace == DEFAULT_WORKSPACE_NAME) .all() )Note: If future multi-tenant support requires passing workspace as a parameter, update the method signature accordingly.
♻️ Duplicate comments (9)
tests/resources/db/latest_schema.sql (4)
29-43: Foreign key constraint to workspaces table is missing.The
evaluation_datasetsandexperimentstables now have workspace columns but lack foreign key constraints to theworkspacestable. This allows invalid workspace values to be inserted.
86-87: Foreign key constraint to workspaces table is missing.The
registered_modelstable now has a workspace column but lacks a foreign key constraint to theworkspacestable.
101-101: Foreign key constraint to workspaces table is missing.The
webhookstable now has a workspace column but lacks a foreign key constraint to theworkspacestable.
198-201: Inconsistent ON DELETE CASCADE behavior in foreign key constraints.The foreign key constraints have inconsistent cascade deletion behavior across model registry tables. This can lead to orphaned records when registered models are deleted.
tests/db/check_migration.py (1)
129-132: Workspace validation will fail for empty tables; use shared constant.The post-migration assertions have two issues:
- For empty tables,
set(df["workspace"]) == {"default"}becomesset() == {"default"}→False, causing spurious failures- The hardcoded
"default"string should use the shared constantmlflow/store/model_registry/sqlalchemy_store.py (4)
702-706: Verify: Hardcoded workspace limits multi-tenancy.Similar to model creation, the workspace parameter is hardcoded to
DEFAULT_WORKSPACE_NAMEwhen setting tags. This prevents workspace-scoped tag management.Verify whether the workspace should be derived from the parent registered model:
#!/bin/bash # Check if workspace should be obtained from the registered model rg -n "set_registered_model_tag.*workspace" --type=py -A5 -B5
813-813: Verify: Hardcoded workspace limits multi-tenancy.Model version creation also hardcodes workspace to
DEFAULT_WORKSPACE_NAME(lines 813, 819). Since model versions reference registered models via composite FK(workspace, name), the workspace must match the parent model's workspace for referential integrity.Verify that workspace is correctly inherited from the parent registered model:
#!/bin/bash # Check if workspace should be obtained from sql_registered_model ast-grep --pattern 'sql_registered_model = $$$ $$$ model_version = SqlModelVersion($$$)'Also applies to: 819-819
1204-1212: Verify: Hardcoded workspace limits multi-tenancy.Model version tag creation hardcodes workspace to
DEFAULT_WORKSPACE_NAME. Similar to model versions, tags must match their parent model version's workspace for FK integrity.
1266-1270: Verify: Hardcoded workspace limits multi-tenancy.Alias creation also hardcodes workspace to
DEFAULT_WORKSPACE_NAME. The workspace must match the parent registered model for FK integrity.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (18)
docs/api_reference/api_inventory.txt(2 hunks)mlflow/entities/__init__.py(2 hunks)mlflow/entities/workspace.py(1 hunks)mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py(1 hunks)mlflow/store/model_registry/dbmodels/models.py(11 hunks)mlflow/store/model_registry/sqlalchemy_store.py(6 hunks)mlflow/store/tracking/dbmodels/models.py(5 hunks)mlflow/store/workspace/dbmodels/__init__.py(1 hunks)mlflow/store/workspace/dbmodels/models.py(1 hunks)mlflow/utils/workspace_utils.py(1 hunks)tests/db/check_migration.py(3 hunks)tests/db/schemas/mssql.sql(6 hunks)tests/db/schemas/mysql.sql(6 hunks)tests/db/schemas/postgresql.sql(6 hunks)tests/db/schemas/sqlite.sql(6 hunks)tests/db/test_schema.py(3 hunks)tests/db/test_workspace_migration.py(1 hunks)tests/resources/db/latest_schema.sql(6 hunks)
🚧 Files skipped from review as they are similar to previous changes (4)
- mlflow/entities/init.py
- mlflow/utils/workspace_utils.py
- mlflow/store/workspace/dbmodels/init.py
- docs/api_reference/api_inventory.txt
🧰 Additional context used
🧬 Code graph analysis (7)
tests/db/check_migration.py (1)
mlflow/store/model_registry/sqlalchemy_store.py (1)
create_registered_model(172-214)
tests/db/test_schema.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
_get_unique_constraints(141-156)
tests/db/test_workspace_migration.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (2)
upgrade(135-612)downgrade(615-910)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
tests/db/test_schema.py (1)
_get_unique_constraints(90-107)
mlflow/store/model_registry/sqlalchemy_store.py (1)
mlflow/store/model_registry/dbmodels/models.py (3)
SqlRegisteredModelTag(155-190)SqlModelVersionTag(193-230)SqlRegisteredModelAlias(233-268)
mlflow/store/workspace/dbmodels/models.py (1)
mlflow/entities/workspace.py (1)
Workspace(9-13)
mlflow/store/model_registry/dbmodels/models.py (1)
mlflow/entities/model_registry/registered_model.py (4)
name(44-46)name(49-50)description(69-71)description(74-75)
🪛 GitHub Actions: Lint
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py
[error] 684-684: F841 Local variable 'fk_constraints' is assigned to but never used. Remove assignment to unused variable.
[error] 691-691: F841 Local variable 'pk_constraints' is assigned to but never used. Remove assignment to unused variable.
[error] 699-699: F841 Local variable 'indexes_to_drop' is assigned to but never used. Remove assignment to unused variable.
[error] 707-707: F841 Local variable 'legacy_pk_defs' is assigned to but never used. Remove assignment to unused variable.
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (33)
- GitHub Check: windows (4)
- GitHub Check: windows (1)
- GitHub Check: pyfunc (2)
- GitHub Check: genai
- GitHub Check: java
- GitHub Check: windows (2)
- GitHub Check: windows (3)
- GitHub Check: pyfunc (4)
- GitHub Check: database
- GitHub Check: pyfunc (3)
- GitHub Check: pyfunc (1)
- GitHub Check: evaluate (2)
- GitHub Check: evaluate (1)
- GitHub Check: models (2)
- GitHub Check: models (1)
- GitHub Check: flavors
- GitHub Check: python (1)
- GitHub Check: python (2)
- GitHub Check: python (3)
- GitHub Check: core
- GitHub Check: build (tracing)
- GitHub Check: python-skinny
- GitHub Check: r
- GitHub Check: build (skinny)
- GitHub Check: build (dev)
- GitHub Check: typescript-sdk (22)
- GitHub Check: js (windows-latest, --testPathIgnorePatterns)
- GitHub Check: js (windows-latest, --testPathPattern)
- GitHub Check: js (ubuntu-latest, --testPathIgnorePatterns)
- GitHub Check: js (ubuntu-latest, --testPathPattern)
- GitHub Check: typescript-sdk (24)
- GitHub Check: typescript-sdk (20)
- GitHub Check: build
🔇 Additional comments (23)
mlflow/store/tracking/dbmodels/models.py (5)
64-64: LGTM! Import added for workspace support.The import of
DEFAULT_WORKSPACE_NAMEis necessary for the workspace column defaults defined below.
98-111: LGTM! Workspace column and constraint implementation.The changes properly implement workspace-scoped experiments:
- Name column is now nullable=False without the unique constraint (moved to composite)
- Workspace column with proper defaults ensures backward compatibility
- Documentation clearly explains the purpose
137-137: LGTM! Composite unique constraint ensures per-workspace name uniqueness.The
UniqueConstrainton(workspace, name)correctly replaces the previous simple unique constraint on name alone, enabling multi-tenant workspace isolation.
1315-1323: LGTM! Workspace column added to SqlEvaluationDataset.The workspace column definition matches the pattern used for SqlExperiment, with proper defaults for legacy compatibility.
1379-1379: LGTM! Index added for workspace filtering.The index on the workspace column will optimize workspace-scoped queries for evaluation datasets.
tests/db/check_migration.py (2)
52-61: LGTM! Workspace tables constant defined.The
WORKSPACE_TABLESset correctly identifies all tables that should have workspace columns after migration.
93-95: LGTM! Additional resources created for comprehensive workspace testing.Creating extra experiments and registered models ensures the workspace backfill is properly tested across multiple resources.
mlflow/entities/workspace.py (1)
1-13: LGTM! Clean and efficient Workspace entity definition.The dataclass is well-designed:
frozen=Trueensures immutabilityslots=Trueoptimizes memory usage- Minimal fields (name and optional description) keep the interface simple
mlflow/store/workspace/dbmodels/models.py (1)
10-22: LGTM! SqlWorkspace ORM model correctly defined.The implementation properly:
- Defines the workspaces table structure
- Sets name as the primary key
- Provides conversion to the Workspace entity
tests/db/test_schema.py (4)
2-16: LGTM! Proper setup for dialect-aware constraint handling.The imports and dialect mapping are well-organized to handle database-specific unique constraint reflection behavior.
36-36: LGTM! Integration of missing constraint reattachment.Calling
_reattach_missing_unique_constraintsbefore dumping the schema ensures the in-memory metadata correctly represents all unique constraints across dialects.
46-87: LGTM! Robust dialect-aware constraint reattachment logic.The function properly handles:
- Dialect-specific constraint reflection differences
- MySQL's unique indexes vs. unique constraints
- MSSQL's lack of
get_unique_constraintsimplementation- Logging of skipped constraints due to missing columns
90-107: LGTM! Fallback logic for unique constraint retrieval.The helper function provides a unified interface for getting unique constraints, with proper fallback to indexes for dialects that don't expose constraints directly.
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (2)
135-613: LGTM! Comprehensive upgrade logic with proper dialect handling.The upgrade function properly:
- Detects and handles dialect-specific behaviors (SQLite, MSSQL, MySQL)
- Adds workspace columns with appropriate defaults
- Updates primary keys and foreign keys to include workspace
- Creates necessary indexes
- Seeds the default workspace
626-646: LGTM! Robust conflict detection prevents data loss on downgrade.The
_assert_no_workspace_conflictsfunction properly:
- Detects duplicate resources that would be created when merging workspaces
- Provides clear error messages with examples
- Aborts the downgrade to prevent data corruption
tests/db/test_workspace_migration.py (7)
47-55: LGTM! Workspace tables constant for test validation.The
_WORKSPACE_TABLEStuple correctly identifies all tables that receive workspace columns, enabling systematic validation across the test suite.
103-114: LGTM! MSSQL-specific identity insert handling.The context manager properly handles MSSQL's
IDENTITY_INSERTrequirement for explicit ID insertion during testing.
160-341: LGTM! Comprehensive legacy data seeding for migration testing.The
_seed_pre_workspace_entitiesfunction properly:
- Uses raw SQL matching the pre-migration schema (no workspace columns)
- Seeds a representative set of entities across all workspace-backed tables
- Ensures the migration is fully responsible for adding/backfilling workspace columns
704-829: LGTM! Thorough upgrade validation.The test comprehensively validates:
- Workspace columns are added with correct defaults
- Legacy data is backfilled with "default" workspace
- Primary keys and foreign keys are updated to include workspace
- Required indexes are created
- Workspaces table is created and seeded
832-922: LGTM! Downgrade validation ensures clean rollback.The test verifies:
- Workspace columns are removed
- Workspaces table is dropped
- Primary keys and foreign keys revert to legacy structure
- Data from multiple workspaces is merged into default
979-1013: LGTM! Conflict detection prevents data loss.The parameterized test ensures the downgrade properly detects and aborts when merging workspaces would create duplicate resources (experiments, models, datasets).
1016-1059: LGTM! Granular conflict query validation.The test validates that the conflict detection queries work correctly for child tables (model_versions, tags, aliases) that require checking multiple columns.
mlflow/store/model_registry/sqlalchemy_store.py (1)
198-198: Hardcoded workspace confirmed—blocks multi-tenancy.The workspace parameter is hardcoded to
DEFAULT_WORKSPACE_NAMEat line 198 (model creation) and line 204 (tag creation). Thecreate_registered_modelmethod signature includes no workspace parameter, so callers cannot specify a target workspace. All registered models will be created in the "default" workspace, preventing actual workspace isolation despite the database schema supporting it.Also applies to: 204-204
mprahl
force-pushed
the
workspaces-db-model
branch
from
fea3308 to
bd0a834
Compare
There was a problem hiding this comment.
Actionable comments posted: 6
♻️ Duplicate comments (6)
mlflow/store/model_registry/dbmodels/models.py (2)
126-133: Consider addingondelete="cascade"to composite FKs for versions and tagsThe composite foreign keys:
SqlModelVersion.__table_args__(["workspace", "name"]→registered_models)SqlRegisteredModelTag.__table_args__(["workspace", "name"]→registered_models)SqlModelVersionTag.__table_args__(["workspace", "name", "version"]→model_versions)all specify
onupdate="cascade"but noondelete, whereasSqlRegisteredModelAliasdoes useondelete="cascade". Without DB‑level cascading, deleting a registered model or model version directly at the SQL layer can leave dangling versions/tags until the application cleans them up.For consistency with
SqlRegisteredModelAliasand to harden referential integrity, it’s worth consideringondelete="cascade"on these FKs as well.Also applies to: 176-183, 211-218
42-49: Add FKs fromworkspaceto theworkspacescatalog for registered models and webhooksThe new
workspacecolumns onSqlRegisteredModelandSqlWebhookare not backed byForeignKeyConstraints to theworkspacestable (seeworkspaces.nameintests/db/schemas/postgresql.sql), so the DB will happily accept non‑existent or mistyped workspace names even though you now maintain a catalog.Consider adding explicit FK constraints, for example:
SqlRegisteredModel.__table_args__: FK["workspace"] -> ["workspaces.name"]SqlWebhook.__table_args__: FK["workspace"] -> ["workspaces.name"]to enforce integrity at the database level in addition to any application‑level checks.
Also applies to: 57-58, 301-306, 317-322
tests/resources/db/latest_schema.sql (2)
106-110: Critical: Missing foreign key constraints to workspaces table.Past reviews have already identified this critical issue. The workspaces catalog table exists but root tables with workspace columns lack FK constraints to enforce referential integrity.
199-201: Major: Inconsistent ON DELETE CASCADE in foreign keys.Past reviews have already flagged the inconsistent cascade deletion behavior. The registered_model_aliases constraint includes
ON DELETE CASCADEbut model_versions, registered_model_tags, and model_version_tags do not, creating orphaned records on deletion.Also applies to: 209-211, 219-221, 373-375
tests/db/schemas/mssql.sql (2)
105-109: Critical: MSSQL schema missing foreign key constraints to workspaces table.Past reviews have already identified this critical issue for the MSSQL schema. The workspaces table is defined but workspace columns lack FK constraints to enforce referential integrity.
196-198: Major: MSSQL schema has inconsistent ON DELETE CASCADE.Past reviews have already flagged the inconsistent cascade deletion behavior in the MSSQL schema, where registered_model_aliases cascades on delete but model_versions and related tables do not.
Also applies to: 206-208, 216-218, 365-367
🧹 Nitpick comments (2)
mlflow/store/model_registry/dbmodels/models.py (1)
271-295: Tighten webhook secret encryption ergonomics (key source and column length)The
EncryptedStringtype works, but two edge cases are worth considering:
- If
MLFLOW_WEBHOOK_SECRET_ENCRYPTION_KEYis unset, you generate a fresh random key per process. Any secrets written in one process become undecryptable after a restart, which can be surprising even in non‑prod environments. You might prefer to fail fast (or log loudly) when the env var is missing outside tests/dev.impl = String(1000)was fine for plaintext, but Fernet ciphertext is longer than the input. Very long secrets (near 1000 chars) risk truncation and decryption failure. Bumping the underlying length (or switching to an unbounded Text type for DB storage) would make this safer.Neither is a blocker, but adjusting these now will avoid hard‑to‑debug webhook issues later.
Also applies to: 312-312
tests/db/test_schema.py (1)
90-107: Minor:_get_unique_constraints’sdialectargument is unused and logic duplicates migration helper
_get_unique_constraintscurrently ignores itsdialectparameter and re‑implements logic that’s very similar to the helper in1b5f0d9ad7c1_add_workspace_columns_and_catalog.py. As a small cleanup, you could either:
- drop the unused
dialectargument here, or- refactor both callers to share a single helper (so any future dialect fixes only need to be made in one place).
This is purely a maintainability nit; behavior is fine as‑is.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (18)
docs/api_reference/api_inventory.txt(2 hunks)mlflow/entities/__init__.py(2 hunks)mlflow/entities/workspace.py(1 hunks)mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py(1 hunks)mlflow/store/model_registry/dbmodels/models.py(11 hunks)mlflow/store/model_registry/sqlalchemy_store.py(6 hunks)mlflow/store/tracking/dbmodels/models.py(5 hunks)mlflow/store/workspace/dbmodels/__init__.py(1 hunks)mlflow/store/workspace/dbmodels/models.py(1 hunks)mlflow/utils/workspace_utils.py(1 hunks)tests/db/check_migration.py(3 hunks)tests/db/schemas/mssql.sql(13 hunks)tests/db/schemas/mysql.sql(6 hunks)tests/db/schemas/postgresql.sql(6 hunks)tests/db/schemas/sqlite.sql(6 hunks)tests/db/test_schema.py(3 hunks)tests/db/test_workspace_migration.py(1 hunks)tests/resources/db/latest_schema.sql(6 hunks)
🚧 Files skipped from review as they are similar to previous changes (8)
- mlflow/entities/workspace.py
- mlflow/store/workspace/dbmodels/models.py
- tests/db/check_migration.py
- mlflow/store/model_registry/sqlalchemy_store.py
- docs/api_reference/api_inventory.txt
- mlflow/entities/init.py
- tests/db/test_workspace_migration.py
- mlflow/store/workspace/dbmodels/init.py
🧰 Additional context used
🧬 Code graph analysis (3)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
tests/db/test_schema.py (1)
_get_unique_constraints(90-107)
tests/db/test_schema.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
_get_unique_constraints(141-155)
mlflow/store/model_registry/dbmodels/models.py (1)
mlflow/entities/model_registry/registered_model.py (4)
name(44-46)name(49-50)description(69-71)description(74-75)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (35)
- GitHub Check: pyfunc (2)
- GitHub Check: pyfunc (1)
- GitHub Check: pyfunc (4)
- GitHub Check: pyfunc (3)
- GitHub Check: windows (4)
- GitHub Check: windows (3)
- GitHub Check: windows (2)
- GitHub Check: windows (1)
- GitHub Check: database
- GitHub Check: python (3)
- GitHub Check: python (2)
- GitHub Check: python (1)
- GitHub Check: python-skinny
- GitHub Check: genai
- GitHub Check: evaluate (1)
- GitHub Check: evaluate (2)
- GitHub Check: models (2)
- GitHub Check: models (1)
- GitHub Check: java
- GitHub Check: flavors
- GitHub Check: js (windows-latest, --testPathIgnorePatterns)
- GitHub Check: js (windows-latest, --testPathPattern)
- GitHub Check: js (ubuntu-latest, --testPathPattern)
- GitHub Check: js (ubuntu-latest, --testPathIgnorePatterns)
- GitHub Check: typescript-sdk (22)
- GitHub Check: typescript-sdk (24)
- GitHub Check: typescript-sdk (20)
- GitHub Check: protect
- GitHub Check: build (skinny)
- GitHub Check: build (dev)
- GitHub Check: r
- GitHub Check: build (tracing)
- GitHub Check: lint-macos
- GitHub Check: build
- GitHub Check: core
🔇 Additional comments (6)
mlflow/utils/workspace_utils.py (1)
1-5: DEFAULT_WORKSPACE_NAME helper looks good and keeps workspace defaults centralizedThe module-level constant and explicit
__all__export are clear and align with how the rest of the PR uses workspace defaults.mlflow/store/tracking/dbmodels/models.py (1)
64-64: Workspace scoping for experiments and evaluation datasets looks consistent; double‑check callers that assume globally unique names
SqlExperimentnow enforces uniqueness on(workspace, name)viauq_experiments_workspace_nameand adds a non‑nullableworkspacecolumn with a default of"default".SqlEvaluationDatasetsimilarly gains a non‑nullableworkspacewith default plus a workspace index.This lines up with the new Postgres schema and the multi‑workspace model, but it does relax the old invariant that experiment names are globally unique. Please confirm higher‑level code paths (e.g. “get by name” lookups, migrations, and uniqueness checks) have been updated to supply or infer the correct workspace rather than assuming a single global namespace.
Also applies to: 98-111, 131-138, 1315-1323, 1375-1380
tests/db/schemas/postgresql.sql (1)
19-31: Postgres schema matches the new workspace‑aware models; keep constraint names in sync with testsThe updated Postgres DDL for:
experiments(workspace column +uq_experiments_workspace_name),registered_models/model_versions/registered_model_aliases/registered_model_tags/model_version_tags,evaluation_datasetsandwebhooks,- and the new
workspacestable,is consistent with the corresponding SQLAlchemy models and composite keys in this PR. Note that
tests/db/test_schema.pynow hard‑codes"uq_experiments_workspace_name"for dialect‑specific reflection fixes, so if you ever rename that constraint you’ll need to update both this schema file and the_DIALECT_REFLECTED_UNIQUE_CONSTRAINTSmapping together.Also applies to: 34-45, 81-88, 91-103, 106-110, 184-201, 204-221, 365-373
tests/db/test_schema.py (1)
2-3: Dialect‑aware reattachment of missingUniqueConstraints is sensible; please validate on MySQL and MSSQLThe
_reattach_missing_unique_constraints+_get_unique_constraintsflow is a pragmatic way to ensureuq_experiments_workspace_nameis present in reflected metadata for MySQL and MSSQL before generating the schema string. Using inspector APIs first and falling back toget_indexeswhen needed should handle the major dialect quirks.Given how finicky SQLAlchemy’s dialects can be here, it’s worth explicitly verifying on real MySQL and MSSQL instances that:
uq_experiments_workspace_nameis reattached exactly once,- no spurious extra
UniqueConstraints are added, anddump_schemaoutput now matches the updatedschemas/*.sqlacross all supported DBs.Also applies to: 8-16, 32-37, 46-88, 90-107
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (2)
133-134: Note: server_default is temporary scaffolding.The comment indicates that the
server_default='default'is temporary until tracking and model registry stores properly set workspace values. Once store-level support is complete, the server_default should be removed to ensure explicit workspace assignment.Track this as a follow-up task to remove the server_default once stores are updated.
621-641: LGTM: Conflict detection prevents unsafe downgrade.The
_assert_no_workspace_conflictsfunction correctly prevents downgrade when merging workspaces would create duplicate resources. The error message provides clear guidance with specific conflicting values (up to 5 examples), making it actionable for operators.
| with _with_batch("experiments") as batch_op: | ||
| if experiments_unique_constraint: | ||
| batch_op.drop_constraint(experiments_unique_constraint, type_="unique") | ||
| elif experiments_unique_index: | ||
| batch_op.drop_index(experiments_unique_index) | ||
| batch_op.add_column( | ||
| sa.Column( | ||
| "workspace", | ||
| sa.String(length=63), | ||
| nullable=False, | ||
| server_default=sa.text("'default'"), | ||
| ) | ||
| ) | ||
| batch_op.create_unique_constraint( | ||
| "uq_experiments_workspace_name", | ||
| ["workspace", "name"], | ||
| ) | ||
|
|
||
| with _with_batch("registered_models") as batch_op: | ||
| if registered_models_unique_constraint: | ||
| batch_op.drop_constraint(registered_models_unique_constraint, type_="unique") | ||
| elif registered_models_unique_index: | ||
| batch_op.drop_index(registered_models_unique_index) | ||
| batch_op.add_column( | ||
| sa.Column( | ||
| "workspace", | ||
| sa.String(length=63), | ||
| nullable=False, | ||
| server_default=sa.text("'default'"), | ||
| ) | ||
| ) | ||
| batch_op.drop_constraint("registered_model_pk", type_="primary") | ||
| batch_op.create_primary_key("registered_model_pk", ["workspace", "name"]) | ||
|
|
||
| with _with_batch("model_versions") as batch_op: | ||
| batch_op.add_column( | ||
| sa.Column( | ||
| "workspace", | ||
| sa.String(length=63), | ||
| nullable=False, | ||
| server_default=sa.text("'default'"), | ||
| ) | ||
| ) | ||
| for fk_name in fk_model_versions: | ||
| batch_op.drop_constraint(fk_name, type_="foreignkey") | ||
| batch_op.drop_constraint("model_version_pk", type_="primary") | ||
| batch_op.create_primary_key("model_version_pk", ["workspace", "name", "version"]) | ||
| batch_op.create_foreign_key( | ||
| "fk_model_versions_registered_models", | ||
| "registered_models", | ||
| ["workspace", "name"], | ||
| ["workspace", "name"], | ||
| onupdate="CASCADE", | ||
| ) | ||
|
|
||
| with _with_batch("registered_model_tags") as batch_op: | ||
| batch_op.add_column( | ||
| sa.Column( | ||
| "workspace", | ||
| sa.String(length=63), | ||
| nullable=False, | ||
| server_default=sa.text("'default'"), | ||
| ) | ||
| ) | ||
| for fk_name in fk_registered_model_tags: | ||
| batch_op.drop_constraint(fk_name, type_="foreignkey") | ||
| batch_op.drop_constraint("registered_model_tag_pk", type_="primary") | ||
| batch_op.create_primary_key("registered_model_tag_pk", ["workspace", "key", "name"]) | ||
| batch_op.create_foreign_key( | ||
| "fk_registered_model_tags_registered_models", | ||
| "registered_models", | ||
| ["workspace", "name"], | ||
| ["workspace", "name"], | ||
| onupdate="CASCADE", | ||
| ) | ||
|
|
||
| with _with_batch("model_version_tags") as batch_op: | ||
| batch_op.add_column( | ||
| sa.Column( | ||
| "workspace", | ||
| sa.String(length=63), | ||
| nullable=False, | ||
| server_default=sa.text("'default'"), | ||
| ) | ||
| ) | ||
| for fk_name in fk_model_version_tags: | ||
| batch_op.drop_constraint(fk_name, type_="foreignkey") | ||
| batch_op.drop_constraint("model_version_tag_pk", type_="primary") | ||
| batch_op.create_primary_key( | ||
| "model_version_tag_pk", | ||
| ["workspace", "key", "name", "version"], | ||
| ) | ||
| batch_op.create_foreign_key( | ||
| "fk_model_version_tags_model_versions", | ||
| "model_versions", | ||
| ["workspace", "name", "version"], | ||
| ["workspace", "name", "version"], | ||
| onupdate="CASCADE", | ||
| ) | ||
|
|
||
| with _with_batch("registered_model_aliases") as batch_op: | ||
| batch_op.add_column( | ||
| sa.Column( | ||
| "workspace", | ||
| sa.String(length=63), | ||
| nullable=False, | ||
| server_default=sa.text("'default'"), | ||
| ) | ||
| ) | ||
| for fk_name in fk_registered_model_aliases: | ||
| batch_op.drop_constraint(fk_name, type_="foreignkey") | ||
| batch_op.drop_constraint("registered_model_alias_pk", type_="primary") | ||
| batch_op.create_primary_key( | ||
| "registered_model_alias_pk", | ||
| ["workspace", "name", "alias"], | ||
| ) | ||
| batch_op.create_foreign_key( | ||
| "fk_registered_model_aliases_registered_models", | ||
| "registered_models", | ||
| ["workspace", "name"], | ||
| ["workspace", "name"], | ||
| onupdate="CASCADE", | ||
| ondelete="CASCADE", | ||
| ) | ||
|
|
||
| with _with_batch("evaluation_datasets") as batch_op: | ||
| batch_op.add_column( | ||
| sa.Column( | ||
| "workspace", | ||
| sa.String(length=63), | ||
| nullable=False, | ||
| server_default=sa.text("'default'"), | ||
| ) | ||
| ) | ||
|
|
||
| with _with_batch("webhooks") as batch_op: | ||
| batch_op.add_column( | ||
| sa.Column( | ||
| "workspace", | ||
| sa.String(length=63), | ||
| nullable=False, | ||
| server_default=sa.text("'default'"), | ||
| ) | ||
| ) | ||
|
|
||
| op.create_index("idx_experiments_workspace", "experiments", ["workspace"]) | ||
| op.create_index("idx_registered_models_workspace", "registered_models", ["workspace"]) | ||
| op.create_index( | ||
| "idx_experiments_workspace_creation_time", | ||
| "experiments", | ||
| ["workspace", "creation_time"], | ||
| unique=False, | ||
| ) | ||
| op.create_index("idx_evaluation_datasets_workspace", "evaluation_datasets", ["workspace"]) | ||
| op.create_index("idx_webhooks_workspace", "webhooks", ["workspace"]) | ||
|
|
||
| op.create_table( | ||
| "workspaces", | ||
| sa.Column("name", sa.String(length=63), nullable=False), | ||
| sa.Column("description", sa.Text(), nullable=True), | ||
| sa.PrimaryKeyConstraint("name", name="workspaces_pk"), | ||
| ) | ||
|
|
||
| metadata = sa.MetaData() | ||
| workspaces_table = sa.Table( | ||
| "workspaces", | ||
| metadata, | ||
| sa.Column("name", sa.String(length=63)), | ||
| sa.Column("description", sa.Text()), | ||
| ) | ||
|
|
||
| conn.execute( | ||
| workspaces_table.insert().values( | ||
| name="default", | ||
| description="Default workspace for legacy resources", | ||
| ) | ||
| ) |
There was a problem hiding this comment.
Critical: Add foreign key constraints from workspace columns to the workspaces table.
The migration creates the workspaces catalog table and adds workspace columns to eight tables, but does not establish foreign key constraints from those workspace columns to workspaces.name. This allows insertion of invalid workspace values (e.g., workspace='nonexistent') that don't exist in the workspaces table.
Root tables that need FK constraints:
evaluation_datasets.workspace → workspaces.nameexperiments.workspace → workspaces.nameregistered_models.workspace → workspaces.namewebhooks.workspace → workspaces.name
Child tables (model_versions, registered_model_aliases, registered_model_tags, model_version_tags) inherit workspace validation through their composite FK constraints to parent tables, so they don't need direct FKs to workspaces.
For SQLite (batch operations), add the FK constraints within the batch contexts:
with _with_batch("experiments") as batch_op:
if experiments_unique_constraint:
batch_op.drop_constraint(experiments_unique_constraint, type_="unique")
elif experiments_unique_index:
batch_op.drop_index(experiments_unique_index)
batch_op.add_column(
sa.Column(
"workspace",
sa.String(length=63),
nullable=False,
server_default=sa.text("'default'"),
)
)
batch_op.create_unique_constraint(
"uq_experiments_workspace_name",
["workspace", "name"],
)
+ batch_op.create_foreign_key(
+ "fk_experiments_workspace",
+ "workspaces",
+ ["workspace"],
+ ["name"],
+ )Apply similar changes for registered_models, evaluation_datasets, and webhooks batch operations, then add the corresponding non-batch FK constraints for non-SQLite dialects after line 574.
🤖 Prompt for AI Agents
In
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py
around lines 229-405, add foreign key constraints from the newly added workspace
columns to workspaces.name for the root tables that currently lack them: inside
the SQLite batch contexts for experiments, registered_models,
evaluation_datasets, and webhooks use batch_op.create_foreign_key(...) to
reference workspaces(name) (e.g., fk_experiments_workspaces -> "workspaces",
["workspace"], ["name"]); do not add direct workspace FKs for child tables
(model_versions, registered_model_aliases, registered_model_tags,
model_version_tags) because they are validated via their composite FKs to parent
tables; additionally, for non-SQLite dialects add corresponding
op.create_foreign_key(...) statements after line 574 so experiments,
registered_models, evaluation_datasets, and webhooks enforce workspace ->
workspaces.name at the non-batch level as well.
| batch_op.create_foreign_key( | ||
| "fk_model_versions_registered_models", | ||
| "registered_models", | ||
| ["workspace", "name"], | ||
| ["workspace", "name"], | ||
| onupdate="CASCADE", | ||
| ) |
There was a problem hiding this comment.
Major: Inconsistent ON DELETE CASCADE behavior creates orphaned records.
Foreign key constraints have inconsistent cascade deletion behavior:
- Line 345-351 (SQLite) & 555-562 (non-SQLite):
registered_model_aliaseshasON DELETE CASCADE✓ - Lines 276-282 (SQLite) & 483-490 (non-SQLite):
model_versionshas onlyON UPDATE CASCADE(missing ON DELETE) - Lines 297-303 (SQLite) & 507-514 (non-SQLite):
registered_model_tagshas onlyON UPDATE CASCADE(missing ON DELETE) - Lines 321-327 (SQLite) & 531-538 (non-SQLite):
model_version_tagshas onlyON UPDATE CASCADE(missing ON DELETE)
When a registered_model is deleted, aliases cascade correctly but model_versions, registered_model_tags, and model_version_tags remain orphaned, requiring manual cleanup.
Add ON DELETE CASCADE to these foreign keys:
batch_op.create_foreign_key(
"fk_model_versions_registered_models",
"registered_models",
["workspace", "name"],
["workspace", "name"],
onupdate="CASCADE",
+ ondelete="CASCADE",
)Apply to all four locations: SQLite batch operations for model_versions (line 276-282), registered_model_tags (297-303), model_version_tags (321-327), and the corresponding non-SQLite operations (lines 483-490, 507-514, 531-538).
Also applies to: 321-327, 483-490, 531-538
🤖 Prompt for AI Agents
In
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py
around lines 276-282, 297-303, 321-327 (SQLite batch_op blocks) and their
non-SQLite counterparts around lines 483-490, 507-514, 531-538, the foreign key
definitions currently only set ON UPDATE CASCADE; add ON DELETE CASCADE to each
of these foreign_key calls so that deletions of registered_models cascade to
model_versions, registered_model_tags, and model_version_tags; update the
batch_op.create_foreign_key(...) calls in each of the six locations to include
ondelete="CASCADE" (in addition to the existing onupdate="CASCADE") to ensure
consistent cascade delete behavior across SQLite and non-SQLite blocks.
| CREATE TABLE workspaces ( | ||
| name VARCHAR(63) NOT NULL, | ||
| description TEXT, | ||
| PRIMARY KEY (name) | ||
| ) |
There was a problem hiding this comment.
Critical: Schema should include foreign key constraints to workspaces table.
The workspaces table is defined but the schema lacks foreign key constraints from workspace columns to workspaces.name. This should match the corrected migration that adds these FK constraints.
Expected FK constraints:
evaluation_datasets.workspace → workspaces.nameexperiments.workspace → workspaces.nameregistered_models.workspace → workspaces.namewebhooks.workspace → workspaces.name
This schema file should reflect the corrected migration once FK constraints are added to the migration file.
🤖 Prompt for AI Agents
In tests/db/schemas/mysql.sql around lines 106 to 110, the schema lacks the
foreign key relationships from workspace columns in other tables to
workspaces.name; update the schema to add foreign key constraints so that
evaluation_datasets.workspace, experiments.workspace,
registered_models.workspace, and webhooks.workspace all reference
workspaces(name). Ensure the referenced column types and lengths match
(VARCHAR(63)), add explicit CONSTRAINT names for clarity, and include
appropriate ON UPDATE/ON DELETE behavior (e.g., ON DELETE CASCADE or RESTRICT
consistent with the migration) so the SQL mirrors the corrected migration.
| workspace VARCHAR(63) DEFAULT 'default' NOT NULL, | ||
| PRIMARY KEY (workspace, name, version), | ||
| CONSTRAINT fk_model_versions_registered_models FOREIGN KEY(workspace, name) REFERENCES registered_models (workspace, name) ON UPDATE CASCADE |
There was a problem hiding this comment.
Major: Schema has inconsistent ON DELETE CASCADE behavior.
The foreign key constraints show inconsistent cascade deletion:
- Line 209:
registered_model_aliaseshasON DELETE CASCADE ON UPDATE CASCADE✓ - Line 199:
model_versionslacksON DELETE CASCADE✗ - Line 219:
registered_model_tagslacksON DELETE CASCADE✗ - Line 374:
model_version_tagslacksON DELETE CASCADE✗
Once the migration is corrected to include ON DELETE CASCADE on all these constraints, regenerate this schema file to reflect the consistent behavior.
Also applies to: 207-209, 217-219, 372-374
🤖 Prompt for AI Agents
In tests/db/schemas/mysql.sql around lines 197-199 (and also ensure changes at
207-209, 217-219, 372-374), the foreign key constraints referencing
registered_models and model_versions are missing ON DELETE CASCADE causing
inconsistent cascade behavior; update each FOREIGN KEY clause to include ON
DELETE CASCADE ON UPDATE CASCADE (or at minimum ON DELETE CASCADE) where
appropriate, then run the migration/regeneration process to produce an updated
schema file so the file reflects the corrected constraints consistently across
all listed lines.
| CREATE TABLE workspaces ( | ||
| name VARCHAR(63) NOT NULL, | ||
| description TEXT, | ||
| CONSTRAINT workspaces_pk PRIMARY KEY (name) | ||
| ) |
There was a problem hiding this comment.
Critical: SQLite schema missing foreign key constraints to workspaces table.
The SQLite schema lacks foreign key constraints from workspace columns to workspaces.name, allowing invalid workspace references. This should be corrected once the migration adds these FK constraints.
Expected FK constraints for SQLite:
evaluation_datasets.workspace → workspaces.nameexperiments.workspace → workspaces.nameregistered_models.workspace → workspaces.namewebhooks.workspace → workspaces.name
🤖 Prompt for AI Agents
In tests/db/schemas/sqlite.sql around lines 106 to 110, the workspaces table
exists but other tables referencing a workspace lack foreign key constraints to
workspaces.name; update the SQLite schema by adding FOREIGN KEY constraints from
evaluation_datasets.workspace, experiments.workspace,
registered_models.workspace, and webhooks.workspace to workspaces(name) (either
inline in each CREATE TABLE definition or via ALTER TABLE ADD CONSTRAINT
statements) and ensure PRAGMA foreign_keys=ON is considered for tests so invalid
workspace references are rejected.
| workspace VARCHAR(63) DEFAULT 'default' NOT NULL, | ||
| CONSTRAINT model_version_pk PRIMARY KEY (workspace, name, version), | ||
| CONSTRAINT fk_model_versions_registered_models FOREIGN KEY(workspace, name) REFERENCES registered_models (workspace, name) ON UPDATE CASCADE |
There was a problem hiding this comment.
Major: SQLite schema has inconsistent ON DELETE CASCADE.
The SQLite schema shows the same inconsistent cascade deletion behavior as other dialects:
- Line 210:
registered_model_aliasesincludesON DELETE CASCADE✓ - Lines 200, 220, 374:
model_versions,registered_model_tags,model_version_tagslackON DELETE CASCADE✗
This schema should be regenerated after the migration is corrected to include consistent ON DELETE CASCADE behavior across all model registry foreign keys.
Also applies to: 208-210, 218-220, 372-374
🤖 Prompt for AI Agents
In tests/db/schemas/sqlite.sql around lines 198-200 (and also update 208-210,
218-220, 372-374), the model registry foreign keys for model_versions,
registered_model_tags and model_version_tags are missing ON DELETE CASCADE;
update the migration that generated this schema so those FOREIGN KEY clauses
include ON DELETE CASCADE for consistency with registered_model_aliases, then
regenerate the sqlite.sql schema file so the referenced lines include the ON
DELETE CASCADE option on the appropriate foreign keys (ensure references to
registered_models and model_versions use ON DELETE CASCADE and commit the
regenerated file).
bd0a834 to
b8fb77e
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tests/db/schemas/postgresql.sql (1)
19-31: Missing FKs from workspace columns toworkspacestable (PostgreSQL schema).The new
workspacescatalog exists, butevaluation_datasets.workspace,experiments.workspace,registered_models.workspace, andwebhooks.workspaceare not constrained toworkspaces(name). This allows invalid workspace values and breaks the catalog invariant.Add
FOREIGN KEY(workspace) REFERENCES workspaces(name)constraints (either inline with reordered table creation or viaALTER TABLEstatements) for these root tables and regenerate this schema from the corrected migration.Also applies to: 34-45, 81-88, 91-103, 106-110
mlflow/store/model_registry/dbmodels/models.py (1)
271-296: Address encryption key persistence and increase column size to prevent data loss and truncation errors.The
EncryptedStringimplementation has two confirmed issues:
Key regeneration on process restart: The
__init__method instantiates the cipher once at column-definition time (when the module is first imported). When a process restarts andMLFLOW_WEBHOOK_SECRET_ENCRYPTION_KEYis unset,__init__generates a new random key. Previously encrypted secrets become undecryptable on subsequent process starts, causing silent data loss until a decrypt attempt fails. This requires the environment variable to be set in production; the fallback should not be the default behavior.Insufficient column capacity: Fernet adds approximately 100 characters of overhead (57 bytes binary + base64 encoding) to encrypted output. For a 900-character plaintext secret, the ciphertext exceeds 1100 characters, truncating against the
String(1000)column. The test only validates small secrets ("my_secret" ≈ 100 chars encrypted), masking this boundary condition.Recommended fixes:
- Require
MLFLOW_WEBHOOK_SECRET_ENCRYPTION_KEYto be explicitly set; remove theor Fernet.generate_key()fallback (or gate it behind an explicit test-mode flag).- Increase the column to
Textor a largerString(e.g.,String(2000)) to safely accommodate encrypted tokens.
♻️ Duplicate comments (12)
tests/db/check_migration.py (1)
52-61: Use the shared default workspace constant instead of hard-coding"default"
WORKSPACE_TABLESand the post‑migration assertions assume"default"as the only workspace value:for table in WORKSPACE_TABLES: df = pd.read_sql(sa.text(f"SELECT DISTINCT workspace FROM {table}"), conn) assert not df["workspace"].isna().any() assert set(df["workspace"]) == {"default"}To keep this test aligned with the rest of the codebase and future‑proof against changes to the default workspace name, consider importing and using the shared constant (e.g.
DEFAULT_WORKSPACE_NAMEfrommlflow.utils.workspace_utils) instead of the string literal here.Also applies to: 161-164
tests/db/schemas/postgresql.sql (1)
184-201: Align model registry FKs with alias behavior by addingON DELETE CASCADE.
model_versions,registered_model_tags, andmodel_version_tagsFKs only specifyON UPDATE CASCADE, whereasregistered_model_aliasesalready usesON DELETE CASCADE ON UPDATE CASCADE. This inconsistency means deleting a registered model can leave orphaned versions and tags unless all deletions are carefully ordered in application code.Consider adding
ON DELETE CASCADEto:
fk_model_versions_registered_modelsfk_registered_model_tags_registered_modelsfk_model_version_tags_model_versionsand then regenerating this PostgreSQL schema to match.
Also applies to: 214-221, 365-373
tests/db/schemas/mysql.sql (2)
19-31: Still missing FKs from workspace columns toworkspaces(MySQL schema).
evaluation_datasets.workspace,experiments.workspace,registered_models.workspace, andwebhooks.workspaceare not constrained toworkspaces.name, so invalid workspace values are still possible.Update the migration to add these FKs and regenerate this MySQL schema so it includes
FOREIGN KEY(workspace) REFERENCES workspaces(name)for the four root tables.Also applies to: 34-45, 81-88, 91-103, 106-110
183-200: Model registry FKs remain inconsistent w.r.tON DELETE CASCADE(MySQL).
model_versions,registered_model_tags, andmodel_version_tagsFKs only haveON UPDATE CASCADE, whileregistered_model_aliasesincludesON DELETE CASCADE ON UPDATE CASCADE. This can leave orphan child rows when a registered model or model version is deleted.For consistency and easier cleanup, consider adding
ON DELETE CASCADEto these FKs in the migration and re‑exporting this schema.Also applies to: 213-220, 367-375
tests/resources/db/latest_schema.sql (2)
19-31: Root tables are not constrained to existing workspaces in latest_schema.sql.The
workspacescatalog exists, butevaluation_datasets.workspace,experiments.workspace,registered_models.workspace, andwebhooks.workspaceare not declared as FKs toworkspaces(name). This schema snapshot still permits invalid workspace values.Once the migration adds these FK constraints, please regenerate
latest_schema.sqlso it reflects them.Also applies to: 34-45, 81-88, 91-103, 106-110
184-201: Cascade behavior for model registry FKs is still inconsistent in latest_schema.sql.Here too,
model_versions,registered_model_tags, andmodel_version_tagsFKs only specifyON UPDATE CASCADE, whileregistered_model_aliasesusesON DELETE CASCADE ON UPDATE CASCADE. This makes deletion semantics depend on table, increasing the chance of orphaned metadata.After deciding the desired behavior (likely full cascade from registered_models/model_versions to all children), update the migration and regenerate this schema with consistent
ON DELETE CASCADEon the affected FKs.Also applies to: 204-221, 367-375
tests/db/schemas/sqlite.sql (2)
19-31: SQLite schema still lacks FKs from workspace columns toworkspaces.
evaluation_datasets.workspace,experiments.workspace,registered_models.workspace, andwebhooks.workspaceare unconstrained, so rows can reference non‑existent workspaces even though aworkspacescatalog exists.Once the migration enforces these relationships, regenerate
sqlite.sqlso each of these tables has aFOREIGN KEY(workspace) REFERENCES workspaces(name)(and ensure tests enablePRAGMA foreign_keys=ON).Also applies to: 34-45, 81-88, 91-103, 106-110
184-201: InconsistentON DELETE CASCADEon SQLite model registry FKs.As with other dialects,
model_versions,registered_model_tags, andmodel_version_tagsonly haveON UPDATE CASCADE, whereasregistered_model_aliasesalso hasON DELETE CASCADE. This inconsistency complicates cleanup when deleting registered models or model versions.Consider updating the migration so these FKs include
ON DELETE CASCADE, then regenerate the SQLite schema to match.Also applies to: 204-221, 367-375
mlflow/store/model_registry/dbmodels/models.py (2)
42-47: Add FK constraints fromworkspaceto theworkspacestable in ORM models.
SqlRegisteredModel.workspaceandSqlWebhook.workspaceare plain columns with noForeignKeyConstraintto theworkspacestable, so the ORM does not enforce that their workspace values exist in the catalog, mirroring the schema issues.Consider extending
__table_args__for both models to include FKs, e.g.:class SqlRegisteredModel(Base): @@ - __table_args__ = (PrimaryKeyConstraint("workspace", "name", name="registered_model_pk"),) + __table_args__ = ( + ForeignKeyConstraint( + ["workspace"], + ["workspaces.name"], + name="fk_registered_models_workspace", + ), + PrimaryKeyConstraint("workspace", "name", name="registered_model_pk"), + ) @@ class SqlWebhook(Base): @@ - __table_args__ = ( - PrimaryKeyConstraint("webhook_id", name="webhook_pk"), - Index("idx_webhooks_status", "status"), - Index("idx_webhooks_name", "name"), - Index("idx_webhooks_workspace", "workspace"), - ) + __table_args__ = ( + ForeignKeyConstraint( + ["workspace"], + ["workspaces.name"], + name="fk_webhooks_workspace", + ), + PrimaryKeyConstraint("webhook_id", name="webhook_pk"), + Index("idx_webhooks_status", "status"), + Index("idx_webhooks_name", "name"), + Index("idx_webhooks_workspace", "workspace"), + )This keeps child models validating workspace via their composite FKs to these roots.
Also applies to: 57-57, 301-306, 317-322
126-133: Addondelete="cascade"to model registry composite FKs for DB‑level integrity.
SqlModelVersion,SqlRegisteredModelTag, andSqlModelVersionTagdefine compositeForeignKeyConstraints withonupdate="cascade"only. Given thatSqlRegisteredModelAliasalready usesondelete="cascade", and ORM relationships usecascade="all", omitting DB‑level delete cascades can leave orphans when rows are deleted via SQL or if cascades are misordered in application code.Align these constraints with alias by adding
ondelete="cascade", for example:ForeignKeyConstraint( ["workspace", "name"], ["registered_models.workspace", "registered_models.name"], - onupdate="cascade", + onupdate="cascade", + ondelete="cascade", ), @@ ForeignKeyConstraint( ["workspace", "name"], ["registered_models.workspace", "registered_models.name"], - onupdate="cascade", + onupdate="cascade", + ondelete="cascade", ), @@ ForeignKeyConstraint( ["workspace", "name", "version"], ["model_versions.workspace", "model_versions.name", "model_versions.version"], - onupdate="cascade", + onupdate="cascade", + ondelete="cascade", ),and then ensure the migrations and schema snapshots are regenerated accordingly.
Also applies to: 176-183, 211-218
tests/db/schemas/mssql.sql (2)
105-109: [DUPLICATE] Add foreign key constraints from workspace columns to the workspaces table.This issue was flagged in the previous review. The MSSQL schema still lacks foreign key constraints from workspace columns to the
workspaces(name)table across multiple tables (evaluation_datasets, experiments, registered_models, webhooks, model_versions, registered_model_aliases, registered_model_tags, model_version_tags). Without these constraints, invalid workspace references are possible.Please reference the prior review comment for the suggested FK constraint additions. The same fix applies here.
Also applies to: 29-29, 41-41, 85-85, 100-100, 195-195, 205-205, 215-215, 364-364
197-197: [DUPLICATE] Inconsistent ON DELETE CASCADE in MSSQL schema.This issue was flagged in the previous review. Three FK constraints (lines 197, 217, 366) lack
ON DELETE CASCADEwhilefk_registered_model_aliases_registered_models(line 207) has it, causing inconsistent cascade behavior:
- Line 197:
fk_model_versions_registered_models— missingON DELETE CASCADE- Line 217:
fk_registered_model_tags_registered_models— missingON DELETE CASCADE- Line 366:
fk_model_version_tags_model_versions— missingON DELETE CASCADEPlease reference the prior review comment for the suggested diffs to add
ON DELETE CASCADEto these constraints.Also applies to: 217-217, 366-366
🧹 Nitpick comments (2)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
288-301: Confirm intended cascade semantics for registry FK relationships
model_versions,registered_model_tags, andmodel_version_tagsare now workspace-scoped and re-wired toregistered_models/model_versionswithonupdate="CASCADE"only, whileregistered_model_aliasesadds bothonupdate="CASCADE"andondelete="CASCADE". This means a direct delete of a registered model will DB‑cascade aliases but not versions/tags, which may or may not match your lifecycle expectations and ORM cascade configuration.Please double‑check that:
- ORM‑level cascades or business logic fully clean up dependent rows when a registered model is deleted, and
- you either explicitly rely on no DB‑level
ON DELETE CASCADEfor versions/tags (to prevent accidental hard deletes) or, if you do want DB‑level cascading, mirror the alias pattern and addondelete="CASCADE"here plus the corresponding downgrade changes.Also applies to: 302-315, 316-331, 405-419, 420-437, 439-456
tests/db/test_schema.py (1)
2-3: Schema reflection fix for MySQL/MSSQL unique constraints looks correctThe combination of
_DIALECT_REFLECTED_UNIQUE_CONSTRAINTS,_reattach_missing_unique_constraints, and_get_unique_constraintsgives a clear, dialect-aware way to ensureuq_experiments_workspace_nameis present in reflected metadata beforedump_schema()compares DDL. The filtering by constraint name and by existing column sets, plus the warning on missing columns, keeps the behavior safe and debuggable.The only minor nit is that
_get_unique_constraintstakes adialectparameter that is not currently used; if you don’t plan to branch on it later, you could drop or rename it to_for clarity, but it’s harmless as-is.Also applies to: 8-16, 32-37, 46-88, 90-107
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (18)
docs/api_reference/api_inventory.txt(2 hunks)mlflow/entities/__init__.py(2 hunks)mlflow/entities/workspace.py(1 hunks)mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py(1 hunks)mlflow/store/model_registry/dbmodels/models.py(11 hunks)mlflow/store/model_registry/sqlalchemy_store.py(6 hunks)mlflow/store/tracking/dbmodels/models.py(5 hunks)mlflow/store/workspace/dbmodels/__init__.py(1 hunks)mlflow/store/workspace/dbmodels/models.py(1 hunks)mlflow/utils/workspace_utils.py(1 hunks)tests/db/check_migration.py(3 hunks)tests/db/schemas/mssql.sql(13 hunks)tests/db/schemas/mysql.sql(6 hunks)tests/db/schemas/postgresql.sql(6 hunks)tests/db/schemas/sqlite.sql(6 hunks)tests/db/test_schema.py(3 hunks)tests/db/test_workspace_migration.py(1 hunks)tests/resources/db/latest_schema.sql(6 hunks)
🚧 Files skipped from review as they are similar to previous changes (7)
- mlflow/store/workspace/dbmodels/init.py
- mlflow/utils/workspace_utils.py
- mlflow/entities/workspace.py
- docs/api_reference/api_inventory.txt
- mlflow/entities/init.py
- tests/db/test_workspace_migration.py
- mlflow/store/workspace/dbmodels/models.py
🧰 Additional context used
🧬 Code graph analysis (5)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
tests/db/test_schema.py (1)
_get_unique_constraints(90-107)
mlflow/store/model_registry/sqlalchemy_store.py (1)
mlflow/store/model_registry/dbmodels/models.py (3)
SqlRegisteredModelTag(155-190)SqlModelVersionTag(193-230)SqlRegisteredModelAlias(233-268)
tests/db/check_migration.py (1)
mlflow/store/model_registry/sqlalchemy_store.py (4)
create_model_version(730-838)set_registered_model_alias(1247-1270)create_registered_model(172-214)create_webhook(1335-1367)
tests/db/test_schema.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
_get_unique_constraints(150-164)
mlflow/store/model_registry/dbmodels/models.py (1)
mlflow/entities/model_registry/registered_model.py (4)
name(44-46)name(49-50)description(69-71)description(74-75)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (35)
- GitHub Check: models (1)
- GitHub Check: windows (4)
- GitHub Check: windows (3)
- GitHub Check: models (2)
- GitHub Check: genai
- GitHub Check: windows (1)
- GitHub Check: flavors
- GitHub Check: windows (2)
- GitHub Check: java
- GitHub Check: pyfunc (1)
- GitHub Check: pyfunc (4)
- GitHub Check: pyfunc (2)
- GitHub Check: pyfunc (3)
- GitHub Check: evaluate (1)
- GitHub Check: evaluate (2)
- GitHub Check: database
- GitHub Check: python (1)
- GitHub Check: python (2)
- GitHub Check: python-skinny
- GitHub Check: python (3)
- GitHub Check: js (ubuntu-latest, --testPathIgnorePatterns)
- GitHub Check: js (windows-latest, --testPathIgnorePatterns)
- GitHub Check: js (windows-latest, --testPathPattern)
- GitHub Check: js (ubuntu-latest, --testPathPattern)
- GitHub Check: build (dev)
- GitHub Check: build (tracing)
- GitHub Check: build (skinny)
- GitHub Check: lint-macos
- GitHub Check: protect
- GitHub Check: build
- GitHub Check: core
- GitHub Check: r
- GitHub Check: typescript-sdk (24)
- GitHub Check: typescript-sdk (22)
- GitHub Check: typescript-sdk (20)
🔇 Additional comments (5)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
491-552: Downgrade safety and conflict checks look solidThe downgrade path validates per-table conflict specs before collapsing workspaces, rewrites all resources back to the
"default"workspace, then cleanly restores legacy PKs/FKs and drops workspace columns and indexes (including MSSQL default-constraint cleanup). This gives a reversible story with clear failure modes if multiple workspaces coexist.No changes requested here.
Also applies to: 555-587, 626-641, 645-666, 679-686, 688-721
tests/db/check_migration.py (1)
69-80: Good coverage of all workspace-backed tables in pre-migration dataThe extended
log_everything()now touches all entries inWORKSPACE_TABLES(model version + alias, extra experiment/model, webhook, and a directevaluation_datasetsrow), so the post‑migration workspace assertions will exercise real backfill behavior instead of failing on empty tables. This is a much more robust CI signal for the migration.Also applies to: 85-107, 109-127
mlflow/store/tracking/dbmodels/models.py (1)
64-65: Workspace columns and constraints for tracking models align with the migrationThe additions of:
workspaceonSqlExperimentwithUniqueConstraint("workspace", "name", name="uq_experiments_workspace_name"), andworkspaceplusIndex("idx_evaluation_datasets_workspace", "workspace")onSqlEvaluationDatasetare consistent with the Alembic migration’s DDL and the workspace backfill strategy. Using
DEFAULT_WORKSPACE_NAMEwith bothdefaultandserver_defaultgives a clean path for new rows while the migration seeds legacy rows.It’s also reasonable that the entity mappers don’t surface
workspaceyet, keeping behavior effectively single‑workspace until higher layers are ready to pass it through.Also applies to: 87-112, 131-138, 1315-1323, 1375-1380
mlflow/store/model_registry/sqlalchemy_store.py (1)
64-65: Consistent default-workspace wiring for registry entitiesThe registry store now consistently sets
workspace=DEFAULT_WORKSPACE_NAMEwhen:
- creating
SqlRegisteredModeland itsSqlRegisteredModelTags,- creating
SqlModelVersionand itsSqlModelVersionTags, and- creating/upserting
SqlRegisteredModelAliasand tags via theset_*APIs.Given that all reads and tag/alias lookups still filter by
name(andversion) only, constraining writes to the default workspace avoids ambiguity for now and keeps existing behavior intact. When you later introduce user-visible multi-workspace support, this will be the natural place to thread an explicit workspace parameter through and extend the query filters accordingly.Also applies to: 192-199, 200-206, 687-706, 800-814, 815-821, 1202-1212, 1264-1270
tests/db/schemas/mssql.sql (1)
42-43: Schema design for workspace scoping looks good.The workspace columns are consistently defined, composite primary keys correctly include workspace where needed, and the UNIQUE constraint on (workspace, name) for experiments appropriately prevents duplicates within workspaces. The collation and default values are uniform.
Also applies to: 86-86, 196-196, 206-206, 216-216, 365-365
mprahl
force-pushed
the
workspaces-db-model
branch
from
b8fb77e to
3fb5a9e
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (16)
tests/db/check_migration.py (1)
52-61: Consider de-duplicating workspace table metadata and avoiding hard-coded"default"
WORKSPACE_TABLEShere and in the migration/tests are hand-maintained lists, and the post-migration assertions hard-code"default"as the only valid workspace. To reduce drift when the workspace model evolves, consider:
- Either sharing a single source of truth for the workspace table set (or at least adding a brief comment tying this constant to the migration’s
_WORKSPACE_TABLES), and- Importing and using
DEFAULT_WORKSPACE_NAMEinstead of the string"default"in thepost_migrationchecks.This keeps the CI check aligned with the runtime defaults as workspaces mature.
Also applies to: 161-164
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (2)
219-252: Add foreign keys from workspace columns to theworkspacescatalog to enforce referential integrityThe migration introduces a
workspacestable and adds non-nullworkspacecolumns to 8 tables, but no FK ties those columns back toworkspaces.name. That means rows can be inserted with arbitraryworkspacevalues that don’t exist in the catalog, undermining the integrity of the new tenant boundary.At minimum, the “root” tables should validate their workspace against the catalog:
experiments.workspace → workspaces.nameregistered_models.workspace → workspaces.nameevaluation_datasets.workspace → workspaces.namewebhooks.workspace → workspaces.nameChild tables (
model_versions,registered_model_aliases,registered_model_tags,model_version_tags) already route validation through composite FKs to these parents, so they don’t need direct FKs toworkspaces.Concretely, you can:
- In the SQLite batches for these four tables, add
batch_op.create_foreign_key(...)calls after the workspace column is added and before exiting each batch.- In the non-SQLite path, add corresponding
op.create_foreign_key(...)definitions after the workspace columns and primary keys are created and before_create_workspace_indexes_and_catalog()(or immediately afterwards, as long as theworkspacestable exists first).Example for non-SQLite
experiments(illustrative):op.create_unique_constraint( "uq_experiments_workspace_name", "experiments", ["workspace", "name"], ) +op.create_foreign_key( + "fk_experiments_workspace", + "experiments", + "workspaces", + ["workspace"], + ["name"], +)And add analogous FKs for
registered_models,evaluation_datasets, andwebhooksin both the SQLite and non-SQLite branches, with correspondingdrop_constraintcalls indowngrade().Also applies to: 351-359, 378-485
288-301: Align cascade behavior: addON DELETE CASCADEto workspace-aware child FKsThe workspace migration updates FKs from:
model_versions → registered_modelsregistered_model_tags → registered_modelsmodel_version_tags → model_versionsbut only specifies
onupdate="CASCADE". In contrast,registered_model_aliasescorrectly uses bothonupdate="CASCADE"andondelete="CASCADE".If a registered model is deleted, you’ll currently:
- Cascade-delete aliases, but
- Leave orphaned
model_versions,registered_model_tags, andmodel_version_tagsrows.To keep behavior consistent and avoid orphans, extend those FKs to include
ondelete="CASCADE"in both SQLite batch ops and non-SQLiteop.create_foreign_keycalls. For example (pattern only):batch_op.create_foreign_key( "fk_model_versions_registered_models", "registered_models", ["workspace", "name"], ["workspace", "name"], - onupdate="CASCADE", + onupdate="CASCADE", + ondelete="CASCADE", )Apply the same to:
- SQLite:
fk_model_versions_registered_models,fk_registered_model_tags_registered_models,fk_model_version_tags_model_versions.- Non-SQLite: the corresponding
op.create_foreign_keydefinitions in the upgrade path.
downgrade()already recreates the legacy non-workspace FKs (with CASCADE semantics where applicable), so only the upgrade side needs adjustment for consistency.Also applies to: 302-315, 316-331, 405-418, 420-437, 439-456
tests/db/schemas/postgresql.sql (2)
106-110: Critical: FK constraints to workspaces table remain missing.This issue was previously identified. The workspaces table exists but root tables with workspace columns still lack foreign key constraints to enforce referential integrity.
198-200: Major: Inconsistent ON DELETE CASCADE behavior persists.This issue was previously identified. The FK constraint on model_versions lacks ON DELETE CASCADE, creating inconsistency with registered_model_aliases which includes it.
tests/db/schemas/sqlite.sql (2)
106-110: Critical: FK constraints to workspaces table remain missing.This issue was previously identified. The SQLite schema still lacks foreign key constraints from workspace columns to the workspaces table.
198-200: Major: Inconsistent ON DELETE CASCADE behavior persists.This issue was previously identified. The SQLite schema shows the same inconsistent cascade deletion behavior.
tests/resources/db/latest_schema.sql (2)
106-110: Critical: FK constraints to workspaces table remain missing.This issue was previously identified. The schema still needs foreign key constraints added for workspace columns referencing workspaces.name.
199-201: Major: Inconsistent ON DELETE CASCADE behavior persists.This issue was previously identified. Foreign key constraints still lack ON DELETE CASCADE for model_versions, registered_model_tags, and model_version_tags.
mlflow/store/model_registry/dbmodels/models.py (5)
42-47: Critical: FK constraint to workspaces table remains missing.This issue was previously identified. SqlRegisteredModel has a workspace column but no foreign key constraint to the workspaces table.
Also applies to: 57-57
126-133: Major: ondelete="cascade" still missing.This issue was previously identified. The ForeignKeyConstraint for SqlModelVersion lacks ondelete="cascade" for database-level referential integrity.
176-183: Major: ondelete="cascade" still missing.This issue was previously identified. The ForeignKeyConstraint for SqlRegisteredModelTag lacks ondelete="cascade".
211-218: Major: ondelete="cascade" still missing.This issue was previously identified. The ForeignKeyConstraint for SqlModelVersionTag lacks ondelete="cascade".
301-306: Critical: FK constraint to workspaces table remains missing.This issue was previously identified. SqlWebhook has a workspace column but no foreign key constraint to the workspaces table.
Also applies to: 317-322
tests/db/schemas/mssql.sql (2)
105-109: Critical: FK constraints to workspaces table remain missing.This issue was previously identified. The MSSQL schema still lacks foreign key constraints from workspace columns to the workspaces table.
196-198: Major: Inconsistent ON DELETE CASCADE behavior persists.This issue was previously identified. The MSSQL schema shows the same inconsistent cascade deletion behavior as other dialects.
🧹 Nitpick comments (5)
mlflow/store/__init__.py (1)
1-4: Temporary workspace models import looks fine; consider a stronger TODO/tracker referenceThis side‑effect import is consistent with the existing
_unity_catalogpattern and keeps tests happy without touching__all__, so no functional concerns. To reduce the risk of this “temporary” hook sticking around, consider turning the comment into an explicit TODO (optionally referencing an issue/PR) so it’s easier to find and remove once a workspace store is wired in.tests/db/test_schema.py (1)
2-16: Unique-constraint reattachment helper is well-scoped; just keep the dialect map in syncThe
_reattach_missing_unique_constraints/_get_unique_constraintspair gives you a clean, dialect-tolerant way to normalize missingUniqueConstraintobjects before dumping schema, without touching live DB state. The guards for MSSQL and MySQL look correct and idempotent.Going forward, just remember that
_DIALECT_REFLECTED_UNIQUE_CONSTRAINTSis now the choke point for any additional workspace-related uniques that are mis-reflected; new constraints with similar issues will need an entry here.Also applies to: 32-43, 46-107
tests/db/test_workspace_migration.py (2)
47-55: Includewebhooksin_WORKSPACE_TABLESto fully exercise the migration surfaceThe migration and
tests/db/check_migration.pyboth treatwebhooksas workspace-backed, but_WORKSPACE_TABLEShere omits it, so_assert_workspace_columns()never validates thewebhooks.workspacecolumn or its default.To keep this test module in sync with the migration’s contract, consider adding
"webhooks"to_WORKSPACE_TABLESand extending any related assertions as needed.Also applies to: 143-146
57-94: Avoid duplicatesa.tabledeclarations for tags/aliases/evaluation_datasets
_REGISTERED_MODEL_TAGS,_MODEL_VERSION_TAGS,_REGISTERED_MODEL_ALIASES, and_EVALUATION_DATASETSare each defined twice with identical schemas (once near the top and again around the helper insert functions). This works because the latter assignments just overwrite the former, but it makes the module harder to read and reason about.You can safely keep a single definition of each table (preferably near the top) and reuse it across helpers and tests.
Also applies to: 587-684
mlflow/store/model_registry/sqlalchemy_store.py (1)
64-65: Workspace defaulting in registry store is consistent for writes; capture the single-workspace assumptionThe changes correctly ensure that all new registry entities (registered models, model versions, their tags, and aliases) are created with
workspace=DEFAULT_WORKSPACE_NAME, matching the new composite PKs and FK layout introduced by the migration.Note, however, that read paths and tag/alias lookups (
_get_registered_model*,_get_model_version*,_get_registered_model_alias) still filter only onname/version/key/aliaswithout a workspace predicate. That’s fine as long as the runtime only ever uses a single workspace (the default), but it becomes ambiguous once non-default workspaces are introduced.It’d be helpful to:
- Explicitly document in this module (or in
DEFAULT_WORKSPACE_NAME’s definition) that the store currently assumes a single default workspace, and- Plan to extend these filters to include workspace once the public API grows a workspace dimension, so lookups stay unambiguous.
No change is strictly required for the current default-only behavior, but calling out the assumption here will make the future multi-workspace extension safer.
Also applies to: 193-206, 686-707, 803-821, 1202-1212, 1264-1270
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (19)
docs/api_reference/api_inventory.txt(2 hunks)mlflow/entities/__init__.py(2 hunks)mlflow/entities/workspace.py(1 hunks)mlflow/store/__init__.py(1 hunks)mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py(1 hunks)mlflow/store/model_registry/dbmodels/models.py(11 hunks)mlflow/store/model_registry/sqlalchemy_store.py(6 hunks)mlflow/store/tracking/dbmodels/models.py(5 hunks)mlflow/store/workspace/dbmodels/__init__.py(1 hunks)mlflow/store/workspace/dbmodels/models.py(1 hunks)mlflow/utils/workspace_utils.py(1 hunks)tests/db/check_migration.py(3 hunks)tests/db/schemas/mssql.sql(13 hunks)tests/db/schemas/mysql.sql(6 hunks)tests/db/schemas/postgresql.sql(6 hunks)tests/db/schemas/sqlite.sql(6 hunks)tests/db/test_schema.py(3 hunks)tests/db/test_workspace_migration.py(1 hunks)tests/resources/db/latest_schema.sql(6 hunks)
🚧 Files skipped from review as they are similar to previous changes (8)
- mlflow/utils/workspace_utils.py
- mlflow/store/tracking/dbmodels/models.py
- mlflow/entities/init.py
- tests/db/schemas/mysql.sql
- mlflow/store/workspace/dbmodels/models.py
- docs/api_reference/api_inventory.txt
- mlflow/entities/workspace.py
- mlflow/store/workspace/dbmodels/init.py
🧰 Additional context used
🧬 Code graph analysis (5)
tests/db/test_workspace_migration.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (2)
upgrade(144-484)downgrade(487-732)
tests/db/check_migration.py (1)
mlflow/store/model_registry/sqlalchemy_store.py (4)
create_model_version(730-838)set_registered_model_alias(1247-1270)create_registered_model(172-214)create_webhook(1335-1367)
tests/db/test_schema.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
_get_unique_constraints(150-164)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
tests/db/test_schema.py (1)
_get_unique_constraints(90-107)
mlflow/store/model_registry/sqlalchemy_store.py (1)
mlflow/store/model_registry/dbmodels/models.py (3)
SqlRegisteredModelTag(155-190)SqlModelVersionTag(193-230)SqlRegisteredModelAlias(233-268)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (34)
- GitHub Check: windows (3)
- GitHub Check: windows (4)
- GitHub Check: evaluate (1)
- GitHub Check: genai
- GitHub Check: windows (2)
- GitHub Check: evaluate (2)
- GitHub Check: windows (1)
- GitHub Check: pyfunc (1)
- GitHub Check: pyfunc (4)
- GitHub Check: models (2)
- GitHub Check: models (1)
- GitHub Check: pyfunc (2)
- GitHub Check: pyfunc (3)
- GitHub Check: flavors
- GitHub Check: python (3)
- GitHub Check: python (2)
- GitHub Check: python (1)
- GitHub Check: database
- GitHub Check: java
- GitHub Check: build (dev)
- GitHub Check: build (tracing)
- GitHub Check: build (skinny)
- GitHub Check: lint-macos
- GitHub Check: r
- GitHub Check: lint
- GitHub Check: build
- GitHub Check: typescript-sdk (24)
- GitHub Check: typescript-sdk (22)
- GitHub Check: js (ubuntu-latest, --testPathPattern)
- GitHub Check: typescript-sdk (20)
- GitHub Check: js (windows-latest, --testPathIgnorePatterns)
- GitHub Check: js (ubuntu-latest, --testPathIgnorePatterns)
- GitHub Check: js (windows-latest, --testPathPattern)
- GitHub Check: protect
🔇 Additional comments (1)
tests/db/check_migration.py (1)
69-127: Seeding all workspace-backed tables before snapshotting looks solidThe extended
log_everything()now exercises every entry inWORKSPACE_TABLES(including aliases, evaluation datasets, and webhooks), which avoids the previous empty-table issue inpost_migrationand gives the migration CI check real data to validate.
mprahl
force-pushed
the
workspaces-db-model
branch
from
3fb5a9e to
64ac13f
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (15)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (2)
219-231: Add FKs from workspace columns toworkspaces.namefor root tablesRight now
experiments,registered_models,evaluation_datasets, andwebhooksall gain a non-nullworkspacecolumn plus aworkspacescatalog, but there is no FK enforcing thatworkspacevalues actually exist inworkspaces.name. That allows invalid workspaces to be inserted and weakens the integrity guarantees of the new catalog.Consider adding FKs from each root table’s
workspacecolumn toworkspaces.namein both branches:
- SQLite: inside the relevant
batch_opblocks forexperiments,registered_models,evaluation_datasets, andwebhooks.- Non-SQLite: after the
op.add_column(... _workspace_column())calls for those tables.For example (SQLite branch):
with _with_batch("experiments") as batch_op: @@ batch_op.create_unique_constraint( "uq_experiments_workspace_name", ["workspace", "name"], ) + batch_op.create_foreign_key( + "fk_experiments_workspace", + "workspaces", + ["workspace"], + ["name"], + )And similarly for
registered_models,evaluation_datasets, andwebhooks(plus correspondingop.create_foreign_key(...)calls in the non-SQLite path). You’d also need to reflect these FKs into the schema snapshots (e.g.,tests/db/schemas/*.sql) to keeptest_schemagreen.Also applies to: 265-359, 377-485
288-315: Align cascade behavior: addON DELETE CASCADEto remaining workspace-aware FKsThe new workspace-aware FKs from:
model_versions.workspace,name → registered_models.workspace,nameregistered_model_tags.workspace,name → registered_models.workspace,namemodel_version_tags.workspace,name,version → model_versions.workspace,name,versiononly specify
onupdate="CASCADE", whileregistered_model_aliases.workspace,name → registered_models.workspace,nameuses bothonupdate="CASCADE"andondelete="CASCADE". This inconsistency means that deleting a registered model at the database level can leave orphaned versions and tags, but aliases are cleaned up.To keep referential behavior consistent and avoid orphaned rows when deletions occur outside SQLAlchemy’s ORM cascades, consider adding
ondelete="CASCADE"to these FKs in both branches, for example:SQLite:
batch_op.create_foreign_key( "fk_model_versions_registered_models", "registered_models", ["workspace", "name"], ["workspace", "name"], - onupdate="CASCADE", + onupdate="CASCADE", + ondelete="CASCADE", )and similarly for
fk_registered_model_tags_registered_modelsandfk_model_version_tags_model_versions.Non-SQLite:
op.create_foreign_key( "fk_model_versions_registered_models", "model_versions", "registered_models", ["workspace", "name"], ["workspace", "name"], - onupdate="CASCADE", + onupdate="CASCADE", + ondelete="CASCADE", )with analogous updates for the tag FKs. Schema snapshots and ORM
ForeignKeyConstraintdefinitions should be updated to match.Also applies to: 405-456
tests/db/schemas/mysql.sql (2)
106-110: Critical: Schema still lacks foreign key constraints from workspace columns to workspaces table.The workspaces catalog table exists but there are no FK constraints from workspace columns in root tables (evaluation_datasets, experiments, registered_models, webhooks) to
workspaces.name. This allows insertion of invalid workspace values that don't exist in the workspaces table.
197-199: Major: Inconsistent ON DELETE CASCADE behavior persists.The foreign key constraints continue to show inconsistent cascade deletion:
- Line 209:
registered_model_aliaseshasON DELETE CASCADE ON UPDATE CASCADE✓- Line 199:
model_versionslacksON DELETE CASCADE✗- Line 219:
registered_model_tagslacksON DELETE CASCADE✗- Line 374:
model_version_tagslacksON DELETE CASCADE✗This inconsistency means deleting a registered_model would cascade to aliases but orphan model_versions and tags.
Also applies to: 207-209, 217-219, 372-374
tests/resources/db/latest_schema.sql (2)
106-110: Critical: Foreign key constraints to workspaces table are still missing.The workspaces catalog table is defined but root tables with workspace columns (evaluation_datasets, experiments, registered_models, webhooks) lack foreign key constraints to
workspaces.name. This permits invalid workspace values.
199-201: Major: ON DELETE CASCADE inconsistency remains unresolved.Foreign key constraints continue to have inconsistent cascade deletion behavior:
- Line 210:
registered_model_aliasesincludes bothON DELETE CASCADE ON UPDATE CASCADE✓- Line 200:
model_versionshas onlyON UPDATE CASCADE✗- Line 220:
registered_model_tagshas onlyON UPDATE CASCADE✗- Line 374:
model_version_tagshas onlyON UPDATE CASCADE✗Also applies to: 209-211, 219-221, 373-375
tests/db/schemas/sqlite.sql (2)
106-110: Critical: SQLite schema still missing foreign key constraints to workspaces table.The workspaces table is defined but workspace columns in root tables (evaluation_datasets, experiments, registered_models, webhooks) lack FK constraints to
workspaces.name, allowing invalid workspace references.
198-200: Major: SQLite schema has persistent inconsistent ON DELETE CASCADE.The cascade deletion behavior remains inconsistent:
- Line 210:
registered_model_aliasesincludesON DELETE CASCADE✓- Lines 200, 220, 374:
model_versions,registered_model_tags,model_version_tagslackON DELETE CASCADE✗Also applies to: 208-210, 218-220, 372-374
mlflow/store/model_registry/dbmodels/models.py (5)
42-47: Critical: FK constraints to workspaces table remain missing from SqlRegisteredModel.The workspace column in
SqlRegisteredModellacks a foreign key constraint to the workspaces table. At line 57,__table_args__contains onlyPrimaryKeyConstraintwithout aForeignKeyConstrainttoworkspaces.name, allowing invalid workspace values.Also applies to: 57-57
126-133: Major: SqlModelVersion FK constraint still lacks ondelete="cascade".The
ForeignKeyConstraintat line 130 includesonupdate="cascade"but omitsondelete="cascade". While the ORM-level relationship specifiescascade="all", database-level cascade deletion is needed for referential integrity during direct SQL operations.
176-183: Major: SqlRegisteredModelTag FK constraint still lacks ondelete="cascade".The
ForeignKeyConstraintat line 180 includesonupdate="cascade"but omitsondelete="cascade", creating inconsistency withSqlRegisteredModelAliaswhich has both cascade options.
211-218: Major: SqlModelVersionTag FK constraint still lacks ondelete="cascade".The
ForeignKeyConstraintat line 215 includesonupdate="cascade"but omitsondelete="cascade". This means database-level deletions of model_versions won't cascade to their tags, requiring manual cleanup.
301-306: Critical: FK constraint to workspaces table missing from SqlWebhook.The workspace column in
SqlWebhooklacks a foreign key constraint to the workspaces table. At lines 317-322,__table_args__contains onlyPrimaryKeyConstraintand indexes without aForeignKeyConstrainttoworkspaces.name.Also applies to: 317-322
tests/db/schemas/mssql.sql (2)
105-109: Critical: MSSQL schema still lacks foreign key constraints to workspaces table.The workspaces table exists but workspace columns in root tables (evaluation_datasets, experiments, registered_models, webhooks) lack FK constraints to
workspaces.name, permitting invalid workspace references.
196-198: Major: MSSQL schema has persistent inconsistent ON DELETE CASCADE.The cascade deletion behavior remains inconsistent across foreign keys:
- Line 207:
registered_model_aliaseshasON DELETE CASCADE✓- Lines 197, 217, 366:
model_versions,registered_model_tags,model_version_tagslackON DELETE CASCADE✗Also applies to: 206-208, 216-218, 365-367
🧹 Nitpick comments (4)
tests/db/test_schema.py (1)
46-88: Minor: avoid duplication and unused parameter in unique-constraint helpersThe
_reattach_missing_unique_constraints/_get_unique_constraintspairing looks solid and should fix dialect-specific reflection issues.Two small maintainability nits:
_get_unique_constraints(inspector, dialect, table_name)never usesdialect; either drop the argument or use it for any future dialect branching to avoid confusion.- The implementation of
_get_unique_constraintshere closely mirrors the helper inside the migration module. If those need to stay behaviorally aligned over time (especially for MSSQL/MySQL edge cases), consider centralizing the logic in a shared utility to reduce drift risk.Also applies to: 90-107
mlflow/store/model_registry/sqlalchemy_store.py (1)
64-65: Explicitly document or enforce the “default workspace only” assumption in the storeAll write paths in this file now hard-code
workspace=DEFAULT_WORKSPACE_NAMEfor registered models, model versions, tags, and aliases, which is fine for a single-workspace deployment. However, the read paths and lookups (_get_registered_model,_get_registered_model_tag,_get_model_version_tag,_get_registered_model_alias, search queries) still filter only byname/versionand explicitly fail if more than one row matches.That means:
- If any other component starts inserting rows for additional workspaces, these helpers will begin raising
INVALID_STATEor returning ambiguous results.- The effective API contract of this store is “operates only on the default workspace”, but this isn’t surfaced anywhere.
Consider either:
- Making this explicit (e.g., via a short class-level comment and/or defensive checks that assert all rows are in
DEFAULT_WORKSPACE_NAME), or- Starting to thread a workspace identifier through these methods and include it in all filters so multi-workspace usage is safe when you decide to expose it.
This can be done incrementally, but clarifying the intended scope now will help avoid subtle breakage later.
Also applies to: 172-207, 686-707, 730-839, 1186-1212, 1247-1270
tests/db/test_workspace_migration.py (2)
47-55: Includewebhooksin workspace migration upgrade/downgrade assertionsThe migration adds a non-null
workspacecolumn (and an index) towebhooks, but the tests in this module don’t currently cover it:
_WORKSPACE_TABLESomits"webhooks", so_assert_workspace_columns()never checks thatwebhooks.workspaceexists and is backfilled.- The downgrade test’s loop over tables to check that
workspacewas removed also omits"webhooks".To fully exercise the migration behavior, consider adding
webhooksto both places, e.g.:_WORKSPACE_TABLES = ( "experiments", "registered_models", "model_versions", "registered_model_tags", "model_version_tags", "registered_model_aliases", "evaluation_datasets", + "webhooks", )and in
test_workspace_migration_downgrade_reverts_schema:- for table in ( + for table in ( "experiments", "registered_models", "model_versions", "registered_model_tags", "model_version_tags", "registered_model_aliases", "evaluation_datasets", + "webhooks", ):This keeps the e2e tests aligned with the actual set of workspace-enabled tables.
Also applies to: 143-146, 867-875
57-95: Remove duplicatesa.tabledefinitions for workspace-aware tables
_REGISTERED_MODEL_TAGS,_MODEL_VERSION_TAGS,_REGISTERED_MODEL_ALIASES, and_EVALUATION_DATASETSare each defined twice with identicalsa.table(...)declarations: once near the top of the file and again around lines 587–683.The duplicates are harmless at runtime (later assignments just overwrite the earlier ones) but add noise and can confuse future readers.
You can safely drop the second block of definitions and rely on the initial ones only, for example:
-_REGISTERED_MODEL_TAGS = sa.table( - "registered_model_tags", - sa.column("workspace"), - sa.column("key"), - sa.column("value"), - sa.column("name"), -) - - -_MODEL_VERSION_TAGS = sa.table( - "model_version_tags", - sa.column("workspace"), - sa.column("key"), - sa.column("value"), - sa.column("name"), - sa.column("version"), -) - -... - -_EVALUATION_DATASETS = sa.table( - "evaluation_datasets", - sa.column("dataset_id"), - sa.column("name"), - sa.column("schema"), - sa.column("profile"), - sa.column("digest"), - sa.column("created_time"), - sa.column("last_update_time"), - sa.column("created_by"), - sa.column("last_updated_by"), - sa.column("workspace"), -)This keeps the test helpers DRY and easier to maintain.
Also applies to: 587-683
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (19)
docs/api_reference/api_inventory.txt(2 hunks)mlflow/entities/__init__.py(2 hunks)mlflow/entities/workspace.py(1 hunks)mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py(1 hunks)mlflow/store/model_registry/dbmodels/models.py(11 hunks)mlflow/store/model_registry/sqlalchemy_store.py(6 hunks)mlflow/store/tracking/dbmodels/models.py(5 hunks)mlflow/store/workspace/dbmodels/__init__.py(1 hunks)mlflow/store/workspace/dbmodels/models.py(1 hunks)mlflow/utils/workspace_utils.py(1 hunks)tests/db/check_migration.py(3 hunks)tests/db/schemas/mssql.sql(13 hunks)tests/db/schemas/mysql.sql(6 hunks)tests/db/schemas/postgresql.sql(6 hunks)tests/db/schemas/sqlite.sql(6 hunks)tests/db/test_schema.py(3 hunks)tests/db/test_workspace_migration.py(1 hunks)tests/resources/db/latest_schema.sql(6 hunks)tests/store/tracking/test_sqlalchemy_store_schema.py(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- tests/store/tracking/test_sqlalchemy_store_schema.py
🚧 Files skipped from review as they are similar to previous changes (7)
- mlflow/utils/workspace_utils.py
- mlflow/store/workspace/dbmodels/init.py
- tests/db/check_migration.py
- mlflow/store/workspace/dbmodels/models.py
- mlflow/entities/init.py
- mlflow/store/tracking/dbmodels/models.py
- mlflow/entities/workspace.py
🧰 Additional context used
🧬 Code graph analysis (4)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
tests/db/test_schema.py (1)
_get_unique_constraints(90-107)
tests/db/test_schema.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
_get_unique_constraints(150-164)
mlflow/store/model_registry/sqlalchemy_store.py (1)
mlflow/store/model_registry/dbmodels/models.py (3)
SqlRegisteredModelTag(155-190)SqlModelVersionTag(193-230)SqlRegisteredModelAlias(233-268)
tests/db/test_workspace_migration.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (2)
upgrade(144-484)downgrade(487-732)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (34)
- GitHub Check: pyfunc (1)
- GitHub Check: pyfunc (4)
- GitHub Check: windows (4)
- GitHub Check: pyfunc (2)
- GitHub Check: pyfunc (3)
- GitHub Check: windows (1)
- GitHub Check: windows (2)
- GitHub Check: evaluate (2)
- GitHub Check: evaluate (1)
- GitHub Check: windows (3)
- GitHub Check: genai
- GitHub Check: python-skinny
- GitHub Check: java
- GitHub Check: models (2)
- GitHub Check: models (1)
- GitHub Check: flavors
- GitHub Check: database
- GitHub Check: python (1)
- GitHub Check: python (3)
- GitHub Check: python (2)
- GitHub Check: protect
- GitHub Check: build (skinny)
- GitHub Check: build (tracing)
- GitHub Check: build (dev)
- GitHub Check: r
- GitHub Check: build
- GitHub Check: js (windows-latest, --testPathIgnorePatterns)
- GitHub Check: js (windows-latest, --testPathPattern)
- GitHub Check: js (ubuntu-latest, --testPathPattern)
- GitHub Check: js (ubuntu-latest, --testPathIgnorePatterns)
- GitHub Check: core
- GitHub Check: typescript-sdk (22)
- GitHub Check: typescript-sdk (20)
- GitHub Check: typescript-sdk (24)
🔇 Additional comments (1)
docs/api_reference/api_inventory.txt (1)
531-531: API inventory entries are correct and properly synchronized with source exports.Verification confirms both entries document valid exports:
mlflow.entities.Workspaceis exported viamlflow/entities/__init__.py(line 65, included in__all__at line 129)mlflow.entities.workspace.Workspaceis accessible directly frommlflow/entities/workspace.py(class defined at line 9)
mprahl
force-pushed
the
workspaces-db-model
branch
from
64ac13f to
c0691c6
Compare
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
mlflow/R/mlflow/tests/testthat/test-run.R (1)
19-51: Now I need to check if this is an isolated issue or if other test files also usewith_mocked_bindings:Update
mlflow/R/mlflow/DESCRIPTIONto require testthat >= 3.1.7 (or >= 3.2.0 for stable).The
with_mocked_bindings()function was introduced as experimental in testthat 3.1.7, but the DESCRIPTION file specifiestestthat (>= 2.0.0). This dependency is incompatible with the test code, which will fail on any system with testthat 2.x or 3.0.x. Update DESCRIPTION to reflect the minimum required version:
- Use
testthat (>= 3.1.7)to allow experimental status, or- Use
testthat (>= 3.2.0)for the stable, released versionmlflow/R/mlflow/tests/testthat/test-client.R (1)
1-151: Update testthat version requirement in mlflow/R/mlflow/DESCRIPTION.The test file uses
with_mocked_bindings(), which was introduced experimentally in testthat 3.1.7 and became stable in testthat 3.2.0. The DESCRIPTION file currently requires onlytestthat (>= 2.0.0), which is incompatible. Update totestthat (>= 3.2.0)to ensure tests run correctly with the stable API.mlflow/cli/__init__.py (1)
631-671: Inconsistent default handling between server and gc commands.The
gccommand's--backend-store-urioption now defaults toNone(line 633), but unlike theservercommand (lines 552-553), there's no runtime default computation or user-facing message informing users which backend store URI is being used.This inconsistency may confuse users who expect the same transparent default behavior across commands. Consider adding similar logic after line 671 to compute and display the default backend store URI when not provided.
Apply this diff to add consistent default messaging:
def gc(older_than, backend_store_uri, artifacts_destination, run_ids, experiment_ids, tracking_uri): """ Permanently delete runs in the `deleted` lifecycle stage from the specified backend store. ... """ from mlflow.utils.time import get_current_time_millis + if not backend_store_uri: + backend_store_uri = _get_default_tracking_uri() + click.echo(f"Backend store URI not provided. Using {backend_store_uri}", err=True) + backend_store = _get_store(backend_store_uri, artifacts_destination)
♻️ Duplicate comments (13)
tests/db/check_migration.py (1)
52-61: Workspace backfill check is much stronger now; consider reusing the shared default constantThe extra resources created in
log_everything()(model version + alias, extra experiment / registered model, webhook, and a directevaluation_datasetsinsert) nicely exercise all entries inWORKSPACE_TABLES, so the post-migration workspace assertions should now see concrete data instead of empty tables.To keep things DRY and aligned with the runtime behavior, it would still be better to replace the hard-coded
"default"in:assert set(df["workspace"]) == {"default"}, ...with the shared constant (e.g.
DEFAULT_WORKSPACE_NAMEfrommlflow.utils.workspace_utils), so the CI check stays in sync if the default workspace name ever changes. This is essentially the same suggestion raised in a previous review.Also applies to: 85-127, 161-164
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (2)
219-230: Add foreign keys from workspace columns to theworkspacescatalog tableThe migration adds non‑nullable
workspacecolumns and creates theworkspacestable plus indexes, but there are no FKs enforcingworkspace→workspaces.namefor the root tables (experiments,registered_models,evaluation_datasets,webhooks). This allows invalid workspace values to be inserted and weakens referential integrity.Concretely:
- SQLite branch:
batch_op.add_column(_workspace_column())is used forexperiments,registered_models,evaluation_datasets, andwebhooks, but nobatch_op.create_foreign_key(..., "workspaces", ["workspace"], ["name"], ...)calls are added.- Non‑SQLite branch:
op.add_column(..., _workspace_column())is used for the same tables, and_create_workspace_indexes_and_catalog()createsworkspaces, but there are no correspondingop.create_foreign_keycalls.You should add FKs from each root table’s
workspacecolumn toworkspaces.namein both branches, and ensure the downgrade drops these FKs beforeop.drop_table("workspaces").Example for one table (SQLite branch):
with _with_batch("experiments") as batch_op: ... batch_op.add_column(_workspace_column()) batch_op.create_unique_constraint( "uq_experiments_workspace_name", ["workspace", "name"], ) + batch_op.create_foreign_key( + "fk_experiments_workspace", + "workspaces", + ["workspace"], + ["name"], + )And similarly for the non‑SQLite branch using
op.create_foreign_key(...)after adding theworkspacecolumns.Also applies to: 265-356, 378-485
288-315: Align cascade behavior on model registry FKs to avoid orphaned childrenFKs tying model registry child tables to
registered_models/model_versionsare inconsistent:
fk_registered_model_aliases_registered_modelshasonupdate="CASCADE", ondelete="CASCADE"(both SQLite and non‑SQLite), which is good.- But the FKs for:
model_versions.workspace,name→registered_models.workspace,nameregistered_model_tags.workspace,name→registered_models.workspace,namemodel_version_tags.workspace,name,version→model_versions.workspace,name,version
are created withonupdate="CASCADE"only (noondelete) in both SQLite (batch_op.create_foreign_key) and non‑SQLite (op.create_foreign_key) paths.- In downgrade, the legacy FKs
model_versions_name_fkey,registered_model_tags_name_fkey, andmodel_version_tags_mv_fkeyare recreated withonupdate="CASCADE"only as well.This means deleting a registered model will cascade to aliases but not to model versions or tags, leaving orphans and requiring manual cleanup.
Recommend adding
ondelete="CASCADE"to all of these FKs in both upgrade and downgrade definitions, e.g.:batch_op.create_foreign_key( "fk_model_versions_registered_models", "registered_models", ["workspace", "name"], ["workspace", "name"], - onupdate="CASCADE", + onupdate="CASCADE", + ondelete="CASCADE", )and the analogous
op.create_foreign_key(...)and legacy*_name_fkeydefinitions.Also applies to: 316-332, 405-419, 420-437, 449-456, 468-476, 568-581, 598-611, 613-624, 696-703, 704-711, 721-728
tests/resources/db/latest_schema.sql (2)
19-31: Workspace FKs to theworkspacescatalog are still missingAs noted in the earlier review, the root tables with
workspace(evaluation_datasets,experiments,registered_models,webhooks) are not constrained toworkspaces(name), so invalid workspace names can be inserted. If you decide to enforce catalog integrity, this file (and the migration + dialect-specific schemas) will need corresponding FKs.Also applies to: 34-45, 81-88, 91-103, 106-110
184-201: Cascade semantics remain inconsistent across workspace-related FKs
registered_model_aliasescascades on delete, butmodel_versions,registered_model_tags, andmodel_version_tagsonly specifyON UPDATE CASCADE. If the intent is to automatically clean up all child metadata when a registered model is deleted, the other FKs likely also needON DELETE CASCADE; otherwise, documenting the asymmetry would help future maintainers.Also applies to: 204-211, 214-221, 367-375
tests/db/schemas/sqlite.sql (2)
106-110: Critical: SQLite schema missing foreign key constraints to workspaces table.The SQLite schema lacks foreign key constraints from workspace columns to
workspaces.name, allowing invalid workspace references. This issue was flagged in previous reviews and remains unresolved.Expected FK constraints for SQLite:
evaluation_datasets.workspace → workspaces.nameexperiments.workspace → workspaces.nameregistered_models.workspace → workspaces.namewebhooks.workspace → workspaces.name
198-200: Inconsistent ON DELETE CASCADE across model registry foreign keys.The SQLite schema shows inconsistent cascade deletion behavior:
- Line 210:
registered_model_aliasesincludesON DELETE CASCADE✓- Lines 200, 220, 374:
model_versions,registered_model_tags,model_version_tagslackON DELETE CASCADE✗This issue was flagged in previous reviews and remains unresolved. The migration should be updated to include consistent
ON DELETE CASCADEbehavior across all model registry foreign keys for referential integrity.Also applies to: 208-210, 218-220, 372-374
mlflow/store/model_registry/dbmodels/models.py (4)
126-133: Addondelete="cascade"to the foreign key constraint for consistency.The
ForeignKeyConstraintincludesonupdate="cascade"but lacksondelete="cascade". While the ORM-level relationship (line 123) specifiescascade="all", having database-level cascade deletion ensures referential integrity even for direct SQL operations.This issue was flagged in previous reviews and remains unresolved.
Apply this diff:
__table_args__ = ( ForeignKeyConstraint( ["workspace", "name"], ["registered_models.workspace", "registered_models.name"], onupdate="cascade", + ondelete="cascade", ), PrimaryKeyConstraint("workspace", "name", "version", name="model_version_pk"), )
176-183: Addondelete="cascade"to maintain consistency.Similar to
SqlModelVersion, this foreign key constraint should includeondelete="cascade"for database-level referential integrity.This issue was flagged in previous reviews and remains unresolved.
Apply this diff:
__table_args__ = ( ForeignKeyConstraint( ["workspace", "name"], ["registered_models.workspace", "registered_models.name"], onupdate="cascade", + ondelete="cascade", ), PrimaryKeyConstraint("workspace", "key", "name", name="registered_model_tag_pk"), )
211-218: Addondelete="cascade"for referential integrity.The foreign key constraint should include
ondelete="cascade"to ensure tags are removed when the parent model version is deleted, maintaining consistency with the ORM-level cascade behavior.This issue was flagged in previous reviews and remains unresolved.
Apply this diff:
__table_args__ = ( ForeignKeyConstraint( ["workspace", "name", "version"], ["model_versions.workspace", "model_versions.name", "model_versions.version"], onupdate="cascade", + ondelete="cascade", ), PrimaryKeyConstraint("workspace", "key", "name", "version", name="model_version_tag_pk"), )
42-47: Critical: FK constraints to workspaces table are missing from root models.All workspace columns across
SqlRegisteredModel,SqlModelVersion,SqlRegisteredModelTag,SqlModelVersionTag,SqlRegisteredModelAlias, andSqlWebhooklack foreign key constraints to theworkspacestable. This allows invalid workspace values.This critical issue was flagged in previous reviews and remains unresolved.
Root-level models (
SqlRegisteredModelandSqlWebhook) must add FK constraints to enforce workspace existence:For SqlRegisteredModel (around line 57):
__table_args__ = ( + ForeignKeyConstraint( + ["workspace"], + ["workspaces.name"], + name="fk_registered_models_workspace", + ), PrimaryKeyConstraint("workspace", "name", name="registered_model_pk"), )For SqlWebhook (around lines 317-322):
__table_args__ = ( + ForeignKeyConstraint( + ["workspace"], + ["workspaces.name"], + name="fk_webhooks_workspace", + ), PrimaryKeyConstraint("webhook_id", name="webhook_pk"), Index("idx_webhooks_status", "status"), Index("idx_webhooks_name", "name"), Index("idx_webhooks_workspace", "workspace"), )Also applies to: 88-93, 158-163, 196-201, 236-241, 301-306
tests/db/schemas/mssql.sql (2)
105-109: Critical: Add foreign key constraints from workspace columns to the workspaces table.The MSSQL schema lacks foreign key constraints from workspace columns to the
workspacescatalog table, allowing invalid workspace references.This critical issue was flagged in previous reviews and remains unresolved.
Add FK constraints for MSSQL dialect:
-- Add after evaluation_datasets table definition ALTER TABLE evaluation_datasets ADD CONSTRAINT fk_evaluation_datasets_workspace FOREIGN KEY(workspace) REFERENCES workspaces(name); -- Add after experiments table definition ALTER TABLE experiments ADD CONSTRAINT fk_experiments_workspace FOREIGN KEY(workspace) REFERENCES workspaces(name); -- Add after registered_models table definition ALTER TABLE registered_models ADD CONSTRAINT fk_registered_models_workspace FOREIGN KEY(workspace) REFERENCES workspaces(name); -- Add after webhooks table definition ALTER TABLE webhooks ADD CONSTRAINT fk_webhooks_workspace FOREIGN KEY(workspace) REFERENCES workspaces(name);
196-198: Inconsistent ON DELETE CASCADE in MSSQL schema.The MSSQL schema has the same inconsistent cascade behavior:
registered_model_aliases(line 207) hasON DELETE CASCADE✓model_versions(line 197),registered_model_tags(line 217), andmodel_version_tags(line 366) lackON DELETE CASCADE✗This major issue was flagged in previous reviews and remains unresolved.
Apply the fix:
workspace VARCHAR(63) COLLATE "SQL_Latin1_General_CP1_CI_AS" DEFAULT ('default') NOT NULL, CONSTRAINT model_version_pk PRIMARY KEY (workspace, name, version), - CONSTRAINT fk_model_versions_registered_models FOREIGN KEY(workspace, name) REFERENCES registered_models (workspace, name) ON UPDATE CASCADE + CONSTRAINT fk_model_versions_registered_models FOREIGN KEY(workspace, name) REFERENCES registered_models (workspace, name) ON DELETE CASCADE ON UPDATE CASCADE ) ... workspace VARCHAR(63) COLLATE "SQL_Latin1_General_CP1_CI_AS" DEFAULT ('default') NOT NULL, CONSTRAINT registered_model_tag_pk PRIMARY KEY (workspace, key, name), - CONSTRAINT fk_registered_model_tags_registered_models FOREIGN KEY(workspace, name) REFERENCES registered_models (workspace, name) ON UPDATE CASCADE + CONSTRAINT fk_registered_model_tags_registered_models FOREIGN KEY(workspace, name) REFERENCES registered_models (workspace, name) ON DELETE CASCADE ON UPDATE CASCADE ) ... workspace VARCHAR(63) COLLATE "SQL_Latin1_General_CP1_CI_AS" DEFAULT ('default') NOT NULL, CONSTRAINT model_version_tag_pk PRIMARY KEY (workspace, key, name, version), - CONSTRAINT fk_model_version_tags_model_versions FOREIGN KEY(workspace, name, version) REFERENCES model_versions (workspace, name, version) ON UPDATE CASCADE + CONSTRAINT fk_model_version_tags_model_versions FOREIGN KEY(workspace, name, version) REFERENCES model_versions (workspace, name, version) ON DELETE CASCADE ON UPDATE CASCADE )Also applies to: 206-208, 216-218, 365-367
🧹 Nitpick comments (22)
mlflow/R/mlflow/tests/testthat/test-rest.R (1)
69-69: LGTM! Modern mocking API adopted correctly.The migration from the deprecated
with_mocktowith_mocked_bindingsis correct and properly targets the httr package namespace for mocking the GET function.Minor formatting nitpick: There's an extra space in
.package = "httr"(double space before=). Consider normalizing to single space for consistency.- with_mocked_bindings(.package = "httr", GET = function(...) { + with_mocked_bindings(.package = "httr", GET = function(...) {mlflow/R/mlflow/tests/testthat/test-model-registry.R (1)
14-14: LGTM! Consistent migration to modern mocking API.All test functions have been correctly updated to use
with_mocked_bindings(.package = "mlflow", ...)instead of the deprecatedwith_mock. The package namespace targeting is appropriate, and all test assertions remain intact.Minor style inconsistency: Lines 42-43 place the
.packageparameter on a separate line, while all other tests keep it on the same line aswith_mocked_bindings. Consider normalizing for consistency:- with_mocked_bindings( - .package = "mlflow", + with_mocked_bindings(.package = "mlflow",Also applies to: 42-43, 79-79, 96-96, 115-115, 135-135, 148-148, 166-166, 211-211, 226-226, 242-242, 262-262, 282-282, 303-303, 323-323, 344-344
mlflow/R/mlflow/tests/testthat/test-client.R (1)
96-151: Approve with optional formatting refinement.The migration to
with_mocked_bindingsfor httr functions is correct, properly specifying.package = "httr".Minor spacing inconsistency: lines 101, 120, and 136 use
.package =(double space) while the rest of the file uses.package =(single space).Apply this diff to normalize spacing:
- with_mocked_bindings(.package = "httr", POST = function(...) { + with_mocked_bindings(.package = "httr", POST = function(...) {- with_mocked_bindings(.package = "httr", GET = function(...) { + with_mocked_bindings(.package = "httr", GET = function(...) {- with_mocked_bindings(.package = "httr", POST = function(...) { + with_mocked_bindings(.package = "httr", POST = function(...) {mlflow/cli/__init__.py (1)
551-557: Good UX improvement with informative default messages.The runtime default computation with user-facing messages provides excellent transparency. Users will know exactly which URIs are being used, which helps with debugging and understanding the system behavior.
Consider whether these messages should go to stderr (using
click.echo(..., err=True)) instead of stdout to avoid interfering with command output parsing, similar to the security middleware messages at lines 572-596.Apply this diff to send messages to stderr:
- click.echo(f"Backend store URI not provided. Using {backend_store_uri}") + click.echo(f"Backend store URI not provided. Using {backend_store_uri}", err=True)And:
- click.echo(f"Registry store URI not provided. Using {registry_store_uri}") + click.echo(f"Registry store URI not provided. Using {registry_store_uri}", err=True)mlflow/server/otel_api.py (1)
129-139: Clarify what “completed trace” means and consider checking the root span’s end timeThe logic for counting “completed” traces currently treats “having a root span in this batch” as completion:
for span in trace_spans: if span.parent_id is None: # Only count traces with a root span as completed completed_trace_ids.add(trace_id) breakThis may overcount if the root span is emitted before it has actually ended (e.g., if an exporter batches or streams spans mid-trace).
If the intention is truly “root span has finished”, consider also checking the root span’s end time:
- for span in trace_spans: - if span.parent_id is None: - # Only count traces with a root span as completed - # (logging of the root span indicates a completed trace) - completed_trace_ids.add(trace_id) - break + for span in trace_spans: + if span.parent_id is None and span.end_time_ns is not None: + # Only count traces whose root span has finished as completed + completed_trace_ids.add(trace_id) + breakIf your exporters only send the root span once it’s finished, the current logic is functionally fine but the comment might be misleading; at minimum, aligning the comment to the actual behavior would help future readers.
It would be good to confirm how your OTLP exporters batch root spans (e.g., whether they’re ever sent before completion); if they always flush on end, this becomes purely a documentation choice.
tests/cli/test_eval.py (1)
176-224: Integration test forevaluate_traces+ custom scorer is solid, with minor robustness considerationsThe new
test_evaluate_traces_integrationwires together:
- Real traces created via
mlflow.start_span()with inputs/outputs.- An inline scorer defined with
@scorerthat parses a numeric score fromoutputs.- A patch of
mlflow.cli.eval.resolve_scorersto inject that scorer.- A post-check via
mlflow.search_traces(..., return_type="list")to assert that each trace has exactly onesimple_scorerassessment with the expected value.This gives good end-to-end coverage of the CLI, scorer resolution, and trace assessment persistence.
Two small things to keep in mind:
- The assertion
len(traces) == 3assumes that evaluation does not create additional traces in the same location. If future changes causeevaluate()to log extra traces, this may become brittle; alternatively, you could filter down to the originaltrace_idsor assert “≥ 3” and then restrict the subsequent checks to known traces.- Sorting by
t.data.spans[0].outputsis clever and should be stable given how you constructed the traces, but if the trace structure changes (e.g., multiple spans with outputs), basing the sort ontrace.info.trace_idor the parsed score fromsimple_scorermight be a bit more future-proof.As-is, the test is clear and valuable; these are just considerations if the tracing/eval plumbing evolves.
If you later extend
evaluate()to emit additional traces or alter how outputs are stored in spans, please revisit this test to ensure the assumptions aboutlen(traces)andt.data.spans[0].outputsstill hold.tests/tracing/test_otel_logging.py (1)
662-748: Consider verifying traces were persisted to the database.The test correctly verifies telemetry events but doesn't confirm that the 2 traces were actually logged to the database. While the test focus is on telemetry, adding a database verification (similar to other tests in this file) would provide more complete coverage.
To verify database persistence, you could add at the end:
traces = mlflow.search_traces( experiment_ids=[experiment_id], include_spans=False, return_type="list" ) assert len(traces) == 2tests/tracking/_tracking_service/test_utils.py (3)
48-56: Make default-URI tests robust to pre-setMLFLOW_TRACKING_URI.This test (and the ones below that call
_get_store()with nostore_uri) assumes the default tracking URI, but doesn't clearMLFLOW_TRACKING_URI. If a developer runs the suite with this env var already set, these assertions could fail for reasons unrelated to the behavior under test. Consider explicitly deleting it at the start, e.g.:def test_tracking_scheme_with_existing_mlruns(tmp_path, monkeypatch): - monkeypatch.chdir(tmp_path) + monkeypatch.chdir(tmp_path) + monkeypatch.delenv(MLFLOW_TRACKING_URI.name, raising=False)and similarly for the other default-URI tests in this block.
65-75: Consider factoring out a small helper to create a legacymlrunslayout.This test repeats the same
mlruns/0/meta.yamlsetup astest_tracking_scheme_with_existing_mlruns. A tiny helper likecreate_legacy_mlruns(tmp_path)would reduce duplication and keep the intent of “existing file-store data” in one place, making future changes to the legacy-detection heuristic easier.
78-85: Be aware these tests instantiate a realSqlAlchemyStore(potentially slower DB path).Unlike
test_get_store_sqlalchemy_store*further down, these two tests call_get_store()expecting aSqlAlchemyStorewithout patchingsqlalchemy.create_engineor schema helpers. That likely exercises real SQLite initialization againsttmp_path, which is fine functionally but may be slower and more brittle if schema upgrade logic changes.If CI time or flakiness becomes a concern, consider mirroring the lower tests and patching:
sqlalchemy.create_enginemlflow.store.db.utils._verify_schemamlflow.store.db.utils._initialize_tablesaround these
_get_store()calls while still asserting that the selection logic prefersSqlAlchemyStorewhenmlrunslacks validmeta.yaml.Also applies to: 87-94
tests/db/compose.yml (1)
14-22: Pinned DB images to digests for deterministic CI behaviorReferencing the database images by content digest improves reproducibility of migration tests without altering service wiring. If you ever need to track upstream versions, consider a brief comment noting the source tag, but this is fine as-is.
Also applies to: 40-48, 67-73
docs/docs/self-hosting/security/basic-http-auth.mdx (1)
7-14: Scorer permissions docs look good; consider minor polishThe new references to scorers and the dedicated permissions table are clear and consistent with the experiment and registered model sections. Two small nits you may want to address:
- Line 101: “an user” → “a user” for grammar.
- In the scorer table, the first two rows annotate permissions as “can_update (on experiment)” / “can_read (on experiment)” while later rows just say “can_read” / “can_delete”. For consistency with the earlier experiment/model tables, either drop the parenthetical or clarify for all rows whether the permission is evaluated on the experiment or on the scorer resource itself.
Also applies to: 99-103, 584-629
mlflow/store/model_registry/file_store.py (1)
135-143: More actionable deprecation warning; consider registry-specific wordingThe expanded warning with an explicit database example and GitHub issue link is a nice improvement for guiding users off the filesystem backend.
Since this class is the model registry
FileStore, you might want to tweak the text from “Filesystem tracking backend (e.g., './mlruns')” to something that also (or primarily) mentions the filesystem-backed model registry, or to a generic “filesystem backend”, so the warning feels accurate in both tracking and registry contexts.mlflow/server/auth/db/migrations/versions/0965eb92f5f0_add_scorer_permissions.py (1)
19-29: Consider adding a foreign key constraint for experiment_id.The
scorer_permissionstable referencesexperiment_idbut doesn't define a foreign key constraint to theexperimentstable. While this might be intentional (e.g., to allow soft references or avoid cascading issues), it differs from typical relational design patterns.Additionally, verify that
scorer_namelength of 256 is sufficient for all expected scorer identifiers.If a foreign key is appropriate, apply this diff:
def upgrade() -> None: op.create_table( "scorer_permissions", sa.Column("id", sa.Integer(), nullable=False, primary_key=True), sa.Column("experiment_id", sa.String(length=255), nullable=False), sa.Column("scorer_name", sa.String(length=256), nullable=False), sa.Column("user_id", sa.Integer(), nullable=False), sa.Column("permission", sa.String(length=255), nullable=True), + sa.ForeignKeyConstraint(["experiment_id"], ["experiments.experiment_id"], name="fk_scorer_perm_experiment_id"), sa.ForeignKeyConstraint(["user_id"], ["users.id"], name="fk_scorer_perm_user_id"), sa.UniqueConstraint("experiment_id", "scorer_name", "user_id", name="unique_scorer_user"), )mlflow/ml-package-versions.yml (1)
753-755: Temporary constraint with open issue - consider tracking.The
qdrant-client<1.16.0constraint is marked as temporary pending resolution of an upstream issue. Consider creating a tracking issue in your repository to ensure this constraint is revisited once the upstream issue is resolved.tests/store/tracking/test_sqlalchemy_store_schema.py (1)
13-15: Temporary import - track cleanup.This import is marked as temporary and should be removed once workspace models are properly integrated into the store layer. Consider creating a tracking issue to ensure this technical debt is addressed.
The comment indicates the proper solution:
# This can be removed once we have a workspace store imported.mlflow/server/auth/entities.py (1)
10-11:ScorerPermissionwiring is consistent; consider defending againstNone/ missing keysThe new
scorer_permissionsfield onUserand theScorerPermissionclass mirror the existing experiment and registered-model permission patterns, which keeps the API surface uniform.To make this a bit more robust:
- If there’s any code path that instantiates
Userwithout explicitly providing a list,self.scorer_permissionsmay beNone, and"[p.to_json() for p in self.scorer_permissions]"will fail. You could defensively treatNoneas[]into_json(and optionally in the property getter), e.g.,for p in self.scorer_permissions or [].- Likewise,
from_jsoncurrently assumes"scorer_permissions"is always present. If you expect to ever deserialize older payloads that don’t include this key,dictionary.get("scorer_permissions", [])would make the upgrade path smoother.If the store layer always sets a concrete list and all callers are version-locked, current code is fine; the above would just future-proof the entity API a bit.
Also applies to: 18-18, 56-63, 73-73, 90-92, 182-230
tests/server/auth/test_sqlalchemy_store.py (1)
10-15: Scorer-permission test coverage is solid and mirrors existing patternsThe
_sp_makerhelper and new tests (test_create_scorer_permission,test_get_scorer_permission, list/update/delete, anddelete_scorer_permissions_for_scorer) nicely parallel the existing experiment and registered-model permission tests, covering:
- happy paths (entity fields, user_id wiring),
- duplicate creation (
RESOURCE_ALREADY_EXISTS),- invalid permission (
INVALID_PARAMETER_VALUE), and- not-found scenarios (
RESOURCE_DOES_NOT_EXIST).Optionally, you could:
- Add an
isinstance(sp1, ScorerPermission)assertion intest_get_scorer_permission(as done for other permission types), and- Add an explicit “invalid permission on update” test for scorer permissions, mirroring
test_update_experiment_permission/test_update_registered_model_permission.Not required, but they’d keep parity across all three permission types.
Also applies to: 48-50, 468-623
tests/server/auth/test_auth.py (1)
635-805: Scorer permission tests cover main flows; a couple of small consistency nitsThe three scorer tests exercise registration, MANAGE default, READ-only access, and delete-denial paths end‑to‑end and look logically correct.
Two minor polish suggestions:
- In
test_register_and_delete_scorer, the admin request toGET_SCORER_PERMISSIONuses hard‑coded"admin", "password1234"under aUser(username1, ...)context. Consider either:
- Using
ADMIN_USERNAME/ADMIN_PASSWORDfor consistency with the rest of the file, or- Dropping the
Usercontext entirely around that admin call since you already pass explicitauth=....- The
scorer_json = '{"name": "test_scorer", "type": "pyfunc"}'literal is duplicated across tests; extracting a small helper/constant would slightly reduce duplication, but is purely optional.tests/db/test_workspace_migration.py (2)
47-55: Consider includingwebhooksin workspace column and index assertionsThe migration adds
workspaceandidx_webhooks_workspaceonwebhooks, but_WORKSPACE_TABLESand_assert_workspace_columns()skip that table, and the upgrade test only checks workspace indexes for experiments, registered_models, and evaluation_datasets. This leaves the webhooks workspace column and index untested.Suggest adding
"webhooks"to_WORKSPACE_TABLESand asserting its workspace index similarly to the other tables.Also applies to: 143-145, 829-842
57-95: Deduplicate workspace-aware table definitions to reduce confusion
_REGISTERED_MODEL_TAGS,_MODEL_VERSION_TAGS,_REGISTERED_MODEL_ALIASES, and_EVALUATION_DATASETSare defined twice with identical workspace-aware schemas. The earlier_LEGACY_*definitions already cover the pre-workspace shape.Recommend keeping a single workspace-aware definition for each of these and relying on the
_LEGACY_*tables for the legacy shape to make the test helpers easier to follow.Also applies to: 587-684
mlflow/server/auth/sqlalchemy_store.py (1)
12-17: Scorer permission CRUD is consistent with existing permission patternsThe scorer permission methods mirror the experiment and registered model permission implementations:
- Use
_validate_permission()before writes.- Resolve users via
_get_user()with properRESOURCE_DOES_NOT_EXIST/INVALID_STATEhandling.- Map
IntegrityErrortoRESOURCE_ALREADY_EXISTSfor creates.- Return
ScorerPermissionentities viato_mlflow_entity().This is idiomatic for the existing store and should integrate cleanly with the new REST and client APIs.
Also applies to: 18-23, 269-356
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
docs/static/images/deep-learning/dl-training-ui.gifis excluded by!**/*.gif
📒 Files selected for processing (69)
.github/workflows/master.yml(1 hunks)README.md(1 hunks)docs/api_reference/api_inventory.txt(4 hunks)docs/docs/classic-ml/deep-learning/pytorch/index.mdx(1 hunks)docs/docs/community/usage-tracking.mdx(1 hunks)docs/docs/self-hosting/security/basic-http-auth.mdx(3 hunks)docs/docusaurus.config.ts(3 hunks)docs/sidebars.ts(1 hunks)docs/sidebarsClassicML.ts(3 hunks)libs/skinny/README_SKINNY.md(1 hunks)libs/typescript/core/tests/clients/client.test.ts(4 hunks)mlflow/R/mlflow/NAMESPACE(0 hunks)mlflow/R/mlflow/R/databricks-utils.R(2 hunks)mlflow/R/mlflow/R/tracking-rest.R(4 hunks)mlflow/R/mlflow/tests/testthat/test-client.R(7 hunks)mlflow/R/mlflow/tests/testthat/test-databricks-utils.R(5 hunks)mlflow/R/mlflow/tests/testthat/test-model-h2o.R(2 hunks)mlflow/R/mlflow/tests/testthat/test-model-registry.R(16 hunks)mlflow/R/mlflow/tests/testthat/test-model-xgboost.R(3 hunks)mlflow/R/mlflow/tests/testthat/test-model.R(4 hunks)mlflow/R/mlflow/tests/testthat/test-params.R(1 hunks)mlflow/R/mlflow/tests/testthat/test-rest.R(1 hunks)mlflow/R/mlflow/tests/testthat/test-run.R(4 hunks)mlflow/cli/__init__.py(4 hunks)mlflow/cli/eval.py(1 hunks)mlflow/entities/__init__.py(2 hunks)mlflow/entities/workspace.py(1 hunks)mlflow/ml-package-versions.yml(31 hunks)mlflow/server/auth/__init__.py(11 hunks)mlflow/server/auth/client.py(2 hunks)mlflow/server/auth/db/migrations/versions/0965eb92f5f0_add_scorer_permissions.py(1 hunks)mlflow/server/auth/db/models.py(4 hunks)mlflow/server/auth/entities.py(5 hunks)mlflow/server/auth/routes.py(1 hunks)mlflow/server/auth/sqlalchemy_store.py(2 hunks)mlflow/server/otel_api.py(5 hunks)mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py(1 hunks)mlflow/store/model_registry/dbmodels/models.py(11 hunks)mlflow/store/model_registry/file_store.py(1 hunks)mlflow/store/model_registry/sqlalchemy_store.py(6 hunks)mlflow/store/tracking/__init__.py(1 hunks)mlflow/store/tracking/dbmodels/models.py(5 hunks)mlflow/store/tracking/file_store.py(1 hunks)mlflow/store/workspace/dbmodels/__init__.py(1 hunks)mlflow/store/workspace/dbmodels/models.py(1 hunks)mlflow/telemetry/events.py(2 hunks)mlflow/tracking/_tracking_service/utils.py(3 hunks)mlflow/tracking/request_header/default_request_header_provider.py(1 hunks)mlflow/utils/cli_args.py(3 hunks)mlflow/utils/server_cli_utils.py(2 hunks)mlflow/utils/workspace_utils.py(1 hunks)tests/cli/test_eval.py(5 hunks)tests/db/check_migration.py(3 hunks)tests/db/compose.yml(3 hunks)tests/db/schemas/mssql.sql(13 hunks)tests/db/schemas/mysql.sql(6 hunks)tests/db/schemas/postgresql.sql(6 hunks)tests/db/schemas/sqlite.sql(6 hunks)tests/db/test_schema.py(3 hunks)tests/db/test_workspace_migration.py(1 hunks)tests/resources/db/latest_schema.sql(6 hunks)tests/server/auth/db/test_cli.py(2 hunks)tests/server/auth/test_auth.py(2 hunks)tests/server/auth/test_sqlalchemy_store.py(3 hunks)tests/store/model_registry/test_file_store.py(1 hunks)tests/store/tracking/test_file_store.py(1 hunks)tests/store/tracking/test_sqlalchemy_store_schema.py(1 hunks)tests/tracing/test_otel_logging.py(2 hunks)tests/tracking/_tracking_service/test_utils.py(1 hunks)
💤 Files with no reviewable changes (1)
- mlflow/R/mlflow/NAMESPACE
✅ Files skipped from review due to trivial changes (3)
- README.md
- mlflow/R/mlflow/tests/testthat/test-params.R
- mlflow/store/tracking/file_store.py
🚧 Files skipped from review as they are similar to previous changes (7)
- mlflow/store/workspace/dbmodels/init.py
- .github/workflows/master.yml
- mlflow/utils/workspace_utils.py
- mlflow/store/workspace/dbmodels/models.py
- mlflow/store/tracking/dbmodels/models.py
- tests/db/schemas/mysql.sql
- docs/docusaurus.config.ts
🧰 Additional context used
🧬 Code graph analysis (21)
docs/sidebars.ts (1)
docs/docusaurusConfigUtils.ts (1)
apiReferencePrefix(6-16)
tests/db/check_migration.py (1)
mlflow/store/model_registry/sqlalchemy_store.py (4)
create_model_version(730-838)set_registered_model_alias(1247-1270)create_registered_model(172-214)create_webhook(1335-1367)
mlflow/server/otel_api.py (4)
mlflow/entities/span.py (5)
events(191-207)trace_id(116-118)trace_id(785-789)parent_id(146-150)parent_id(812-813)mlflow/telemetry/events.py (2)
TraceSource(326-330)TracesReceivedByServerEvent(333-334)mlflow/telemetry/track.py (1)
_record_event(86-101)mlflow/store/tracking/sqlalchemy_store.py (1)
log_spans(3289-3468)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
tests/db/test_schema.py (1)
_get_unique_constraints(90-107)
docs/sidebarsClassicML.ts (1)
docs/docusaurusConfigUtils.ts (1)
apiReferencePrefix(6-16)
mlflow/server/auth/routes.py (1)
mlflow/server/handlers.py (1)
_get_rest_path(3710-3711)
mlflow/store/model_registry/sqlalchemy_store.py (1)
mlflow/store/model_registry/dbmodels/models.py (3)
SqlRegisteredModelTag(155-190)SqlModelVersionTag(193-230)SqlRegisteredModelAlias(233-268)
tests/db/test_workspace_migration.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (2)
upgrade(144-484)downgrade(487-732)
tests/cli/test_eval.py (3)
mlflow/genai/scorers/base.py (1)
scorer(847-1031)mlflow/tracking/fluent.py (1)
set_experiment(130-230)mlflow/cli/eval.py (1)
evaluate_traces(61-128)
mlflow/utils/server_cli_utils.py (1)
mlflow/utils/uri.py (1)
is_local_uri(28-69)
tests/tracking/_tracking_service/test_utils.py (3)
mlflow/tracking/_tracking_service/utils.py (1)
_get_store(252-253)mlflow/store/tracking/file_store.py (1)
FileStore(187-2811)mlflow/store/tracking/sqlalchemy_store.py (1)
SqlAlchemyStore(188-4511)
tests/tracing/test_otel_logging.py (2)
mlflow/telemetry/client.py (2)
TelemetryClient(27-335)add_record(113-129)mlflow/telemetry/events.py (2)
TraceSource(326-330)TracesReceivedByServerEvent(333-334)
mlflow/entities/__init__.py (1)
mlflow/entities/workspace.py (1)
Workspace(9-13)
tests/server/auth/test_auth.py (3)
tests/server/auth/auth_test_utils.py (2)
create_user(13-26)User(29-41)mlflow/server/auth/entities.py (7)
User(1-93)permission(116-117)permission(120-121)permission(159-160)permission(163-164)permission(208-209)permission(212-213)tests/tracking/integration_test_utils.py (1)
_send_rest_tracking_post_request(112-120)
mlflow/tracking/_tracking_service/utils.py (4)
tests/store/tracking/test_file_store.py (1)
store(78-79)mlflow/store/model_registry/file_store.py (1)
FileStore(122-1115)mlflow/utils/file_utils.py (3)
exists(127-128)path_to_local_file_uri(481-485)path(281-282)mlflow/environment_variables.py (2)
get(41-51)get(75-98)
mlflow/cli/__init__.py (1)
mlflow/tracking/_tracking_service/utils.py (3)
_get_default_tracking_uri(55-60)is_tracking_uri_set(63-67)set_tracking_uri(70-121)
tests/server/auth/test_sqlalchemy_store.py (6)
mlflow/server/auth/entities.py (19)
ExperimentPermission(96-136)RegisteredModelPermission(139-179)ScorerPermission(182-230)User(1-93)experiment_id(108-109)experiment_id(196-197)scorer_name(200-201)username(25-26)permission(116-117)permission(120-121)permission(159-160)permission(163-164)permission(208-209)permission(212-213)id(21-22)name(151-152)user_id(112-113)user_id(155-156)user_id(204-205)mlflow/server/auth/__init__.py (4)
create_scorer_permission(1398-1404)get_scorer_permission(1408-1413)update_scorer_permission(1417-1423)delete_scorer_permission(1427-1432)mlflow/server/auth/sqlalchemy_store.py (6)
create_scorer_permission(269-290)get_scorer_permission(319-325)list_scorer_permissions(327-335)update_scorer_permission(337-344)delete_scorer_permission(346-349)delete_scorer_permissions_for_scorer(351-356)tests/helper_functions.py (1)
random_str(49-56)mlflow/exceptions.py (1)
MlflowException(67-113)mlflow/protos/databricks_pb2.pyi (1)
ErrorCode(17-97)
mlflow/server/auth/client.py (3)
mlflow/server/auth/entities.py (18)
ExperimentPermission(96-136)RegisteredModelPermission(139-179)ScorerPermission(182-230)User(1-93)experiment_id(108-109)experiment_id(196-197)scorer_name(200-201)username(25-26)permission(116-117)permission(120-121)permission(159-160)permission(163-164)permission(208-209)permission(212-213)from_json(77-93)from_json(131-136)from_json(174-179)from_json(224-230)mlflow/server/auth/__init__.py (4)
create_scorer_permission(1398-1404)get_scorer_permission(1408-1413)update_scorer_permission(1417-1423)delete_scorer_permission(1427-1432)mlflow/server/auth/sqlalchemy_store.py (4)
create_scorer_permission(269-290)get_scorer_permission(319-325)update_scorer_permission(337-344)delete_scorer_permission(346-349)
mlflow/server/auth/db/models.py (1)
mlflow/server/auth/entities.py (20)
ExperimentPermission(96-136)RegisteredModelPermission(139-179)ScorerPermission(182-230)User(1-93)scorer_permissions(57-58)scorer_permissions(61-62)id(21-22)experiment_id(108-109)experiment_id(196-197)scorer_name(200-201)user_id(112-113)user_id(155-156)user_id(204-205)permission(116-117)permission(120-121)permission(159-160)permission(163-164)permission(208-209)permission(212-213)name(151-152)
mlflow/server/auth/sqlalchemy_store.py (5)
mlflow/server/auth/db/models.py (6)
SqlScorerPermission(77-94)SqlUser(21-42)to_mlflow_entity(31-42)to_mlflow_entity(53-58)to_mlflow_entity(69-74)to_mlflow_entity(88-94)mlflow/server/auth/entities.py (18)
ExperimentPermission(96-136)RegisteredModelPermission(139-179)ScorerPermission(182-230)User(1-93)experiment_id(108-109)experiment_id(196-197)scorer_name(200-201)username(25-26)permission(116-117)permission(120-121)permission(159-160)permission(163-164)permission(208-209)permission(212-213)user_id(112-113)user_id(155-156)user_id(204-205)id(21-22)mlflow/server/auth/__init__.py (4)
create_scorer_permission(1398-1404)get_scorer_permission(1408-1413)update_scorer_permission(1417-1423)delete_scorer_permission(1427-1432)mlflow/server/auth/client.py (4)
create_scorer_permission(548-578)get_scorer_permission(580-605)update_scorer_permission(607-634)delete_scorer_permission(636-657)mlflow/exceptions.py (1)
MlflowException(67-113)
mlflow/server/auth/__init__.py (6)
mlflow/java/client/src/main/java/org/mlflow/api/proto/Service.java (17)
DeleteScorer(183911-185270)Response(21470-22086)Response(23709-24774)Response(26225-26881)Response(27544-27953)Response(28616-29025)Response(29724-30133)Response(31176-31832)Response(33594-34250)Response(35528-35937)Response(36600-37009)Response(37936-38345)Response(40450-40859)Response(42191-42600)Response(43705-44114)Response(45088-45497)Response(46805-47214)mlflow/server/auth/entities.py (15)
experiment_id(108-109)experiment_id(196-197)name(151-152)username(25-26)permission(116-117)permission(120-121)permission(159-160)permission(163-164)permission(208-209)permission(212-213)scorer_name(200-201)to_json(64-74)to_json(123-128)to_json(166-171)to_json(215-221)mlflow/server/auth/client.py (4)
get_scorer_permission(580-605)create_scorer_permission(548-578)update_scorer_permission(607-634)delete_scorer_permission(636-657)mlflow/server/auth/sqlalchemy_store.py (5)
get_scorer_permission(319-325)create_scorer_permission(269-290)delete_scorer_permissions_for_scorer(351-356)update_scorer_permission(337-344)delete_scorer_permission(346-349)mlflow/utils/proto_json_utils.py (1)
parse_dict(196-199)mlflow/server/handlers.py (1)
catch_mlflow_exception(723-734)
| def train_distributed(): | ||
| # Initialize distributed training | ||
| dist.init_process_group(backend="nccl") | ||
| rank = dist.get_rank() | ||
|
|
||
| # Wrap model with DDP | ||
| model_ddp = DDP(model.to(rank), device_ids=[rank]) | ||
|
|
||
| # Create distributed sampler | ||
| from torch.utils.data.distributed import DistributedSampler | ||
|
|
||
| sampler = DistributedSampler( | ||
| train_dataset, num_replicas=dist.get_world_size(), rank=rank | ||
| ) | ||
| train_loader = DataLoader(train_dataset, batch_size=32, sampler=sampler) | ||
|
|
||
| # Only log from rank 0 | ||
| if rank == 0: | ||
| mlflow.start_run() | ||
| mlflow.log_params({"world_size": dist.get_world_size(), "backend": "nccl"}) | ||
|
|
||
| # Training loop | ||
| epochs = 10 | ||
| for epoch in range(epochs): | ||
| sampler.set_epoch(epoch) # Shuffle data differently each epoch | ||
| train_loss = train_epoch(model_ddp, train_loader) | ||
|
|
||
| # Log metrics from rank 0 only | ||
| if rank == 0: | ||
| mlflow.log_metric("train_loss", train_loss, step=epoch) | ||
|
|
||
| # Save model from rank 0 | ||
| if rank == 0: | ||
| mlflow.pytorch.log_model(model, name="distributed_model") | ||
| mlflow.end_run() | ||
| ``` |
There was a problem hiding this comment.
Add cleanup and clarify distributed training execution.
The distributed training example has several issues:
- Missing cleanup: The process group is initialized but never destroyed, which can lead to resource leaks.
- Incomplete execution context: The example doesn't show how to launch the distributed training script (e.g., using
torch.distributed.launchortorchrun).
Apply this diff to add proper cleanup:
# Save model from rank 0
if rank == 0:
mlflow.pytorch.log_model(model, name="distributed_model")
mlflow.end_run()
+
+ # Cleanup distributed resources
+ dist.destroy_process_group()Additionally, consider adding a comment explaining how to run the distributed script:
# Run this script with torchrun:
# torchrun --nproc_per_node=2 train_distributed.py
def train_distributed():
# Initialize distributed training
dist.init_process_group(backend="nccl")
...🤖 Prompt for AI Agents
In docs/docs/classic-ml/deep-learning/pytorch/index.mdx around lines 439 to 474,
the distributed training snippet never tears down the process group and lacks
guidance on how to launch the script; wrap the init/ training/ logging in a
try/finally (or try/except/finally) so that on all code paths you call
dist.destroy_process_group(), ensure mlflow.end_run() is called from rank 0 in
the finally block to avoid orphaned runs, and add a top-line comment showing the
recommended launch command (e.g., "Run this script with torchrun: torchrun
--nproc_per_node=2 train_distributed.py") so users know how to execute the
distributed job.
| - [Daniel Lok](https://github.com/daniellok-db) | ||
| - [Gabriel Fu](https://github.com/gabrielfu) | ||
| - [Harutaka Kawamura](https://github.com/harupy) | ||
| - [Joel Robin P](https://github.com/joelrobin18) |
There was a problem hiding this comment.
Manual edit to autogenerated file—change the source instead.
Line 1 declares this file as autogenerated and forbids manual edits. The addition of Joel Robin P should be made in dev/pyproject.py (the source generator) rather than manually edited here. This ensures the change persists if the file is regenerated.
Regenerate the file after updating the source:
python dev/pyproject.pyThen remove the manual edit and commit the regenerated version.
🤖 Prompt for AI Agents
In libs/skinny/README_SKINNY.md around line 333, the file was manually edited to
add "Joel Robin P"; because the file is autogenerated you must instead add that
contributor to the generator source dev/pyproject.py, run the generator (python
dev/pyproject.py) to regenerate README_SKINNY.md, verify the change appears,
remove the manual edit from this file, and commit the regenerated file so the
change persists on future generations.
| from mlflow.telemetry.events import TraceSource, TracesReceivedByServerEvent | ||
| from mlflow.telemetry.track import _record_event | ||
| from mlflow.tracing.utils.otlp import MLFLOW_EXPERIMENT_ID_HEADER, OTLP_TRACES_PATH | ||
| from mlflow.tracking.request_header.default_request_header_provider import ( | ||
| _MLFLOW_PYTHON_CLIENT_USER_AGENT_PREFIX, | ||
| _USER_AGENT, | ||
| ) |
There was a problem hiding this comment.
🧩 Analysis chain
Emit a string for source instead of the raw Enum in telemetry params
The OTEL handler wiring looks correct: you derive trace_source from the User-Agent prefix and emit a single TracesReceivedByServerEvent summarizing the request. One thing to tighten:
- You currently pass the Enum instance as a param:
trace_source = (
TraceSource.MLFLOW_PYTHON_CLIENT
if user_agent and user_agent.startswith(_MLFLOW_PYTHON_CLIENT_USER_AGENT_PREFIX)
else TraceSource.UNKNOWN
)
_record_event(
TracesReceivedByServerEvent,
{
"source": trace_source,
"count": len(completed_trace_ids),
},
)Depending on how telemetry records are serialized, storing an Enum instead of a plain string can be brittle (e.g., JSON encoding or schema validation).
I recommend emitting the Enum’s value:
- _record_event(
- TracesReceivedByServerEvent,
- {
- "source": trace_source,
- "count": len(completed_trace_ids),
- },
- )
+ _record_event(
+ TracesReceivedByServerEvent,
+ {
+ "source": trace_source.value,
+ "count": len(completed_trace_ids),
+ },
+ )This also matches the documented example payload {"source": "MLFLOW_PYTHON_CLIENT", "count": 3}.
Please double-check how telemetry records are serialized (e.g., via Pydantic / JSON) to confirm you don’t rely on Enum objects being stored as-is; if you already normalize Enums centrally, this diff is still a harmless simplification.
Also applies to: 159-172
I'll verify the review comment by examining how telemetry records are serialized and whether the Enum needs to be converted to a string value.
---
🏁 Script executed:
```shell
#!/bin/bash
# Inspect the actual code at lines 23-29 and 159-172 in mlflow/server/otel_api.py
echo "=== Lines 23-29 ==="
sed -n '23,29p' mlflow/server/otel_api.py
echo ""
echo "=== Lines 159-172 ==="
sed -n '159,172p' mlflow/server/otel_api.py
Length of output: 878
🏁 Script executed:
#!/bin/bash
# Find and inspect _record_event function
echo "=== Finding _record_event definition ==="
rg -nP 'def _record_event' --type=py -A 10Length of output: 846
🏁 Script executed:
#!/bin/bash
# Find TracesReceivedByServerEvent definition to check if it's Pydantic
echo "=== Finding TracesReceivedByServerEvent ==="
rg -nP 'class TracesReceivedByServerEvent' --type=py -A 15Length of output: 191
🏁 Script executed:
#!/bin/bash
# Check how telemetry params are serialized
echo "=== Checking telemetry serialization ==="
rg -nP 'TracesReceivedByServerEvent' --type=py -B 3 -A 3 | head -50Length of output: 3671
Use .value when emitting Enum to telemetry params
The review comment is correct. Test assertions in tests/tracing/test_otel_logging.py (lines 658 and 747) explicitly validate that the telemetry record params contain string values, not Enum instances:
assert record.params["source"] == TraceSource.MLFLOW_PYTHON_CLIENT.value
assert record.params["source"] == TraceSource.UNKNOWN.valueThe current code at lines 159–172 in mlflow/server/otel_api.py passes the raw Enum ("source": trace_source), which would fail these tests. The _record_event function does not normalize Enums centrally; it passes the params dict as-is to the telemetry Record. Therefore, you must emit the Enum's string value:
_record_event(
TracesReceivedByServerEvent,
{
- "source": trace_source,
+ "source": trace_source.value,
"count": len(completed_trace_ids),
},
)This fix applies to both occurrences in the file.
🤖 Prompt for AI Agents
In mlflow/server/otel_api.py around lines 159 to 172, the telemetry params
currently pass the Enum object for the "source" field (trace_source) which
causes tests to fail; update both occurrences to emit the Enum's string value by
using trace_source.value when building the params dict (e.g., "source":
trace_source.value) so that _record_event receives plain strings instead of Enum
instances.
| # Also used as default location for artifacts, when not provided, in non local file based backends | ||
| # (eg MySQL) | ||
| DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH = "./mlruns" | ||
| DEFAULT_TRACKING_URI = "sqlite:///mlflow.db" |
There was a problem hiding this comment.
Consider the implications of using a relative path for the default SQLite database.
The value "sqlite:///mlflow.db" uses three slashes, indicating a relative path. This means the database file will be created in the current working directory when MLflow commands are executed.
Users who run MLflow commands from different directories may inadvertently create multiple separate database files, leading to confusion when their experiments appear to be "missing."
Verify that:
- This behavior is intentional and documented for users
- The backward compatibility logic (checking for existing mlruns data) properly handles directory changes
- Users are guided to set
MLFLOW_TRACKING_URIexplicitly for production use
🤖 Prompt for AI Agents
mlflow/store/tracking/__init__.py around line 13: the default
DEFAULT_TRACKING_URI is a relative SQLite path ("sqlite:///mlflow.db") which
causes DB files to be created in whatever CWD a user runs MLflow from; change
the default to compute an absolute path at runtime (e.g., under the user home or
a well-known per-user directory like ~/.mlflow/mlflow.db) instead of a
hard-coded relative URI, update the backward-compatibility lookup logic to also
check the previous relative locations (CWD/mlflow.db and mlruns directories)
when migrating/locating existing data, and add/update docs and an inline comment
advising users to set MLFLOW_TRACKING_URI explicitly for production and
describing the default fallback behavior.
|
|
||
| """ | ||
| return backend_store_uri == DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH | ||
| return backend_store_uri == DEFAULT_TRACKING_URI |
There was a problem hiding this comment.
🧩 Analysis chain
Verify this change aligns with the new default tracking URI semantics.
The comparison now checks against DEFAULT_TRACKING_URI (sqlite:///mlflow.db) instead of DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH (./mlruns). This means the validation logic at line 49 will now consider sqlite:///mlflow.db as the "default" backend store URI when checking artifacts-only mode.
Ensure that the CLI runtime default computation (in mlflow/cli/init.py) can return either value depending on existing data, and that this validation correctly handles both cases.
Run the following script to verify how the default is computed and used:
🏁 Script executed:
#!/bin/bash
# Description: Check how _get_default_tracking_uri is used and what values it can return
# Find usages of _get_default_tracking_uri
rg -nP --type=py -C3 '_get_default_tracking_uri\(\)'
# Find usages of _is_default_backend_store_uri
rg -nP --type=py -C3 '_is_default_backend_store_uri'
# Check artifacts_only_config_validation usage
rg -nP --type=py -C5 'artifacts_only_config_validation'Length of output: 5028
🏁 Script executed:
#!/bin/bash
# Get the full implementation of _get_default_tracking_uri
rg -nP --type=py -A10 'def _get_default_tracking_uri'
# Get constant definitions
rg -nP --type=py 'DEFAULT_TRACKING_URI\s*='
rg -nP --type=py 'DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH\s*='
# Check path_to_local_file_uri usage
rg -nP --type=py -B2 -A2 'path_to_local_file_uri'Length of output: 14303
🏁 Script executed:
#!/bin/bash
# Check what constants are imported in server_cli_utils.py
rg -nP --type=py -B10 '_is_default_backend_store_uri' mlflow/utils/server_cli_utils.py | head -30
# Check the imports at top of server_cli_utils.py
head -50 mlflow/utils/server_cli_utils.pyLength of output: 2464
Fix validation logic to recognize both possible default backend store URIs.
The change at line 45 introduces a critical bug: _is_default_backend_store_uri() now only checks for DEFAULT_TRACKING_URI ("sqlite:///mlflow.db"), but _get_default_tracking_uri() can return either "sqlite:///mlflow.db" or "./mlruns" (when legacy data exists).
When artifacts-only mode is used with existing mlruns data, the validation will incorrectly reject "./mlruns" as a non-default value, even though it is a legitimate default. The comparison at line 45 must account for both possible defaults:
return backend_store_uri in (DEFAULT_TRACKING_URI, DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH)🤖 Prompt for AI Agents
In mlflow/utils/server_cli_utils.py around line 45, the helper
_is_default_backend_store_uri currently only checks equality with
DEFAULT_TRACKING_URI and fails to recognize the alternate legacy default
"./mlruns"; update the validation so it returns True when backend_store_uri
matches either DEFAULT_TRACKING_URI or DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH.
Replace the single equality check with a membership test against the two known
default constants to accept both possible default backend URIs.
| def test_otel_trace_received_telemetry_from_mlflow_client(mlflow_server: str): | ||
| """ | ||
| Test TraceReceivedByServerEvent telemetry shows source=MLFLOW_PYTHON_CLIENT for standard client. | ||
|
|
||
| Uses @mlflow.trace with standard MLflow client configuration, which automatically sends | ||
| User-Agent and X-MLflow-Client-Version headers to identify traces from MLflow client. | ||
| """ | ||
| mlflow.set_tracking_uri(mlflow_server) | ||
| mlflow.set_experiment("otel-telemetry-mlflow-client-test") | ||
|
|
||
| with mock.patch("mlflow.telemetry.track.get_telemetry_client") as mock_get_client: | ||
| mock_client = mock.MagicMock(spec=TelemetryClient) | ||
| mock_get_client.return_value = mock_client | ||
|
|
||
| @mlflow.trace | ||
| def test_function(): | ||
| return "test result" | ||
|
|
||
| result = test_function() | ||
| assert result == "test result" | ||
|
|
||
| time.sleep(1) | ||
|
|
||
| if mock_client.add_record.called: | ||
| record = mock_client.add_record.call_args[0][0] | ||
| assert record.event_name == TracesReceivedByServerEvent.name | ||
| assert record.params["source"] == TraceSource.MLFLOW_PYTHON_CLIENT.value | ||
| assert record.params["count"] == 1 | ||
|
|
There was a problem hiding this comment.
Critical: Test may pass silently without verifying telemetry.
The conditional check at line 655 means the test will pass even if add_record is never called, defeating the test's purpose. Unlike the second test (line 742), which properly uses assert_called_once(), this test uses optional assertions inside an if block.
Apply this diff to fix the assertion logic:
- time.sleep(1)
-
- if mock_client.add_record.called:
- record = mock_client.add_record.call_args[0][0]
- assert record.event_name == TracesReceivedByServerEvent.name
- assert record.params["source"] == TraceSource.MLFLOW_PYTHON_CLIENT.value
- assert record.params["count"] == 1
+ time.sleep(1)
+
+ mock_client.add_record.assert_called_once()
+ record = mock_client.add_record.call_args[0][0]
+ assert record.event_name == TracesReceivedByServerEvent.name
+ assert record.params["source"] == TraceSource.MLFLOW_PYTHON_CLIENT.value
+ assert record.params["count"] == 1📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def test_otel_trace_received_telemetry_from_mlflow_client(mlflow_server: str): | |
| """ | |
| Test TraceReceivedByServerEvent telemetry shows source=MLFLOW_PYTHON_CLIENT for standard client. | |
| Uses @mlflow.trace with standard MLflow client configuration, which automatically sends | |
| User-Agent and X-MLflow-Client-Version headers to identify traces from MLflow client. | |
| """ | |
| mlflow.set_tracking_uri(mlflow_server) | |
| mlflow.set_experiment("otel-telemetry-mlflow-client-test") | |
| with mock.patch("mlflow.telemetry.track.get_telemetry_client") as mock_get_client: | |
| mock_client = mock.MagicMock(spec=TelemetryClient) | |
| mock_get_client.return_value = mock_client | |
| @mlflow.trace | |
| def test_function(): | |
| return "test result" | |
| result = test_function() | |
| assert result == "test result" | |
| time.sleep(1) | |
| if mock_client.add_record.called: | |
| record = mock_client.add_record.call_args[0][0] | |
| assert record.event_name == TracesReceivedByServerEvent.name | |
| assert record.params["source"] == TraceSource.MLFLOW_PYTHON_CLIENT.value | |
| assert record.params["count"] == 1 | |
| def test_otel_trace_received_telemetry_from_mlflow_client(mlflow_server: str): | |
| """ | |
| Test TraceReceivedByServerEvent telemetry shows source=MLFLOW_PYTHON_CLIENT for standard client. | |
| Uses @mlflow.trace with standard MLflow client configuration, which automatically sends | |
| User-Agent and X-MLflow-Client-Version headers to identify traces from MLflow client. | |
| """ | |
| mlflow.set_tracking_uri(mlflow_server) | |
| mlflow.set_experiment("otel-telemetry-mlflow-client-test") | |
| with mock.patch("mlflow.telemetry.track.get_telemetry_client") as mock_get_client: | |
| mock_client = mock.MagicMock(spec=TelemetryClient) | |
| mock_get_client.return_value = mock_client | |
| @mlflow.trace | |
| def test_function(): | |
| return "test result" | |
| result = test_function() | |
| assert result == "test result" | |
| time.sleep(1) | |
| mock_client.add_record.assert_called_once() | |
| record = mock_client.add_record.call_args[0][0] | |
| assert record.event_name == TracesReceivedByServerEvent.name | |
| assert record.params["source"] == TraceSource.MLFLOW_PYTHON_CLIENT.value | |
| assert record.params["count"] == 1 |
🤖 Prompt for AI Agents
In tests/tracing/test_otel_logging.py around lines 632 to 660, the test
currently wraps assertions about mock_client.add_record inside an if block so
the test can silently pass when add_record was never called; replace the
conditional with a strict assertion (e.g., assert mock_client.add_record.called
or mock_client.add_record.assert_called_once()) and then extract the recorded
call (record = mock_client.add_record.call_args[0][0]) and assert
record.event_name, record.params["source"] and record.params["count"] as before
so the test fails if no telemetry was sent.
This adds the required workspace columns and the workspace catalog table with the default workspace precreated. All workspace columns default to "default" for now and we may choose to remove the defaults once the tracking store and model registry store are made workspace aware to catch application logic issues not properly setting the workspace. Some model registry store changes were needed to account for the new composite foreign key. Signed-off-by: mprahl <mprahl@users.noreply.github.com>
mprahl
force-pushed
the
workspaces-db-model
branch
from
c0691c6 to
1f0b624
Compare
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
mlflow/genai/evaluation/entities.py (1)
171-190:tablesproperty looks good, but__repr__has a pandas truth‑value bugThe new
tablesproperty is straightforward and matches the new tests’ expectations ("eval_results"mapped to the sameresult_dfwhen present, empty dict otherwise) — no issues there.However,
__repr__currently does:result_df_str = ( f"{len(self.result_df)} rows x {len(self.result_df.columns)} cols" if self.result_df else "None" )Using
if self.result_dfon apd.DataFramewill raiseValueError: The truth value of a DataFrame is ambiguouswhenresult_dfis a non‑empty DataFrame. This is a latent correctness bug that will surface wheneverEvaluationResultis printed or logged.A safe fix is to test explicitly for
None:- result_df_str = ( - f"{len(self.result_df)} rows x {len(self.result_df.columns)} cols" - if self.result_df - else "None" - ) + if self.result_df is not None: + result_df_str = ( + f"{len(self.result_df)} rows x {len(self.result_df.columns)} cols" + ) + else: + result_df_str = "None"With that change, the new
tablesproperty and the existing__repr__are both safe and consistent.mlflow/store/tracking/databricks_rest_store.py (1)
165-189: Critical: HandleNonelocation in endpoint construction.Making
locationoptional with a default ofNonealigns with the abstract base class and other store implementations, but the endpoint construction on line 187 will fail iflocationisNone. The f-stringf"{_V4_TRACE_REST_API_PATH_PREFIX}/{location}/batchGet"will interpolateNoneinto the URL, resulting in an invalid endpoint like.../None/batchGet.Consider one of these approaches:
- Add validation to require
locationfor this Databricks implementation:def batch_get_traces(self, trace_ids: list[str], location: str | None = None) -> list[Trace]: """ Get a batch of complete traces with spans for given trace ids. Args: trace_ids: List of trace IDs to fetch. location: Location of the trace. For example, "catalog.schema" for UC schema. Returns: List of Trace objects. """ + if location is None: + raise MlflowException( + "location parameter is required for Databricks trace store", + error_code=INVALID_PARAMETER_VALUE, + ) trace_ids = [parse_trace_id_v4(trace_id)[1] for trace_id in trace_ids]
- Or handle the
Nonecase by extracting location from trace IDs or falling back to a different behavior (if applicable).mlflow/store/model_registry/sqlalchemy_store.py (1)
216-243: Add workspace filtering to all read/search helper methodsThe review comment is accurate. Verification confirms that all the identified helper methods lack workspace filtering despite tables having workspace-keyed primary keys:
- SqlRegisteredModel uses PrimaryKeyConstraint("workspace", "name") but
_get_registered_model()filters only on name- SqlModelVersion uses PrimaryKeyConstraint("workspace", "name", "version") but
_get_sql_model_version()conditions include only name, version, and stage—no workspace- Similar gaps exist in
_get_registered_model_tag(),_get_model_version_tag(), and_get_registered_model_alias()While workspace=DEFAULT_WORKSPACE_NAME is set when creating records, this is insufficient for multi-workspace safety. The query methods must explicitly filter by workspace to prevent data collisions and eliminate the "Expected only 1 … Found {len(rms)}" errors that would occur with multiple workspaces.
Recommend adding workspace filtering (either explicit DEFAULT_WORKSPACE_NAME or parametrized) to all seven helper methods at lines 216–243, 377–440, 849–885, 1165–1185, 1237–1245, 1292–1326, and the associated tag/alias queries.
♻️ Duplicate comments (10)
tests/db/check_migration.py (1)
161-164: Avoid hard‑coding"default"; reuse the shared workspace constant and consider empty‑table semantics.Two follow‑ups here:
- Shared constant instead of string literal (repeat of earlier feedback).
Using the literal"default"for the workspace name will drift ifDEFAULT_WORKSPACE_NAMEever changes. Importing and using the shared constant keeps this CI check aligned with the rest of the codebase:-from mlflow.store.tracking.dbmodels.models import ( +from mlflow.store.tracking.dbmodels.models import ( SqlModelVersion, ... ) +from mlflow.utils.workspace_utils import DEFAULT_WORKSPACE_NAME ... - for table in WORKSPACE_TABLES: - df = pd.read_sql(sa.text(f"SELECT DISTINCT workspace FROM {table}"), conn) - assert not df["workspace"].isna().any(), f"{table} contains NULL workspace values" - assert set(df["workspace"]) == {"default"}, f"{table} contains non-default workspaces" + for table in WORKSPACE_TABLES: + df = pd.read_sql(sa.text(f"SELECT DISTINCT workspace FROM {table}"), conn) + assert not df["workspace"].isna().any(), f"{table} contains NULL workspace values" + assert set(df["workspace"]) == {DEFAULT_WORKSPACE_NAME}, ( + f"{table} contains non-default workspaces" + )
- Optional: bail out early for empty tables.
You’ve mitigated the previous empty‑table issue by seeding more rows, but if any of these tables are ever legitimately empty in a migration scenario,set(df["workspace"]) == {DEFAULT_WORKSPACE_NAME}will still fail spuriously. Adding a simpleif df.empty: continuekeeps the check robust without weakening the signal.mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (2)
219-252: Still missing foreign keys from workspace columns toworkspaces.nameThe migration creates the
workspacescatalog and adds non‑nullableworkspacecolumns, but it never declares FKs from those columns toworkspaces.name(forexperiments,registered_models,evaluation_datasets, andwebhooks). This allows rows with arbitraryworkspacevalues that don’t exist in the catalog, undermining the integrity the catalog is meant to provide.Please add
ForeignKey/create_foreign_keyconstraints from each root table’sworkspacecolumn toworkspaces(name)in both the SQLite batch blocks and the non‑SQLite path, and ensure the downgrade drops those FKs before removing the column/table.Also applies to: 268-278, 351-356, 377-385, 478-485, 730-732
288-301:model_versionsand tag tables still lackON DELETE CASCADEto registered modelsThe upgraded schema continues the previous inconsistency where:
registered_model_aliasesFKs toregistered_modelsspecifyondelete="CASCADE", butmodel_versionsandregistered_model_tagsFKs toregistered_models, andmodel_version_tagsFKs tomodel_versions, only useonupdate="CASCADE".Deleting a registered model can therefore leave orphaned versions and tags. To keep referential behavior consistent with aliases and avoid manual cleanup, extend these FKs in both the SQLite batch operations and the non‑SQLite path to include
ondelete="CASCADE"alongside the existingonupdate="CASCADE".For example:
batch_op.create_foreign_key( "fk_model_versions_registered_models", "registered_models", ["workspace", "name"], ["workspace", "name"], - onupdate="CASCADE", + onupdate="CASCADE", + ondelete="CASCADE", )and similarly for the corresponding
registered_model_tags/model_version_tagsFKs and their non‑workspace counterparts indowngrade().Also applies to: 309-314, 321-331, 411-418, 430-437, 449-456, 574-580, 590-596, 606-611, 618-624, 696-703, 705-711, 713-720, 721-728
tests/resources/db/latest_schema.sql (2)
19-31: Workspaces catalog exists but workspace columns still lack foreign key constraints.
evaluation_datasets.workspace,experiments.workspace,registered_models.workspace, andwebhooks.workspaceare now present, andworkspaces(name)is defined, but none of these columns are constrained to the catalog. That still allows invalid workspace values and breaks the intent of the workspace table.Recommend adding FKs like:
ALTER TABLE evaluation_datasets ADD CONSTRAINT fk_evaluation_datasets_workspace FOREIGN KEY (workspace) REFERENCES workspaces (name); ALTER TABLE experiments ADD CONSTRAINT fk_experiments_workspace FOREIGN KEY (workspace) REFERENCES workspaces (name); ALTER TABLE registered_models ADD CONSTRAINT fk_registered_models_workspace FOREIGN KEY (workspace) REFERENCES workspaces (name); ALTER TABLE webhooks ADD CONSTRAINT fk_webhooks_workspace FOREIGN KEY (workspace) REFERENCES workspaces (name);Also applies to: 34-45, 81-88, 91-103, 106-110
184-201: Model‑related foreign keys still have inconsistent ON DELETE CASCADE settings.
registered_model_aliasescascades deletes correctly, but:
model_versions→registered_models(fk_model_versions_registered_models)registered_model_tags→registered_models(fk_registered_model_tags_registered_models)model_version_tags→model_versions(fk_model_version_tags_model_versions)all lack
ON DELETE CASCADE. Deleting a registered model will leave orphaned versions/tags unless they’re cleaned up separately.Recommend updating these constraints to include
ON DELETE CASCADE ON UPDATE CASCADEfor consistency with aliases and to match expected cleanup semantics.Also applies to: 214-221, 367-375
tests/db/schemas/mysql.sql (2)
19-31: MySQL schema: workspace columns are not constrained to the workspaces catalog.The MySQL schema mirrors the new
workspacecolumns andworkspacestable but still lacks FKs:
evaluation_datasets.workspaceexperiments.workspaceregistered_models.workspacewebhooks.workspaceNone reference
workspaces(name), so invalid workspace names remain possible. This file should reflect the finalized migration by adding correspondingFOREIGN KEY (workspace) REFERENCES workspaces (name)constraints (with explicitCONSTRAINTnames) on each of these tables.Also applies to: 34-45, 81-88, 91-103, 106-110
183-200: MySQL schema: foreign keys for model versions and tags still lack ON DELETE CASCADE.As in
latest_schema.sql, the MySQL schema shows:
model_versions→registered_models(fk_model_versions_registered_models) with onlyON UPDATE CASCADE.registered_model_tags→registered_models(fk_registered_model_tags_registered_models) with onlyON UPDATE CASCADE.model_version_tags→model_versions(fk_model_version_tags_model_versions) with onlyON UPDATE CASCADE.while
registered_model_aliasescascades on delete.For consistency and to avoid orphaned rows after deleting a registered model or version, these constraints should also include
ON DELETE CASCADE.Also applies to: 213-220, 367-375
tests/db/schemas/sqlite.sql (1)
19-31: SQLite workspace schema still missing FKs toworkspacesand consistentON DELETE CASCADEThe workspace-aware schema is wired up (columns +
workspacestable), but two previously flagged issues remain:
No FK from workspace columns to
workspaces(name)
evaluation_datasets.workspace(Line 29),experiments.workspace(Line 41),registered_models.workspace(Line 86), andwebhooks.workspace(Line 101) are not constrained toworkspaces(name). This allows invalid workspace values and defeats the point of the catalog table.Model registry FKs still lack
ON DELETE CASCADE
model_versions→registered_models(Line 200)registered_model_tags→registered_models(Line 220)model_version_tags→model_versions(Line 374)
These only specifyON UPDATE CASCADE, so deleting a registered model / version can leave orphan versions/tags in SQLite tests.Please extend the migration and regenerate this schema so:
- Those four root tables have FKs
... FOREIGN KEY(workspace) REFERENCES workspaces(name) ....- The three model registry FKs also include
ON DELETE CASCADE(matching other dialects and ORM intent).Also applies to: 34-45, 81-88, 90-103, 106-110, 184-201, 204-221, 367-375
tests/db/schemas/mssql.sql (1)
19-31: MSSQL workspace schema still lacks FKs toworkspacesand full cascade semanticsThe MSSQL schema mirrors the unresolved issues seen in SQLite:
No FK from workspace columns to
workspaces(name)
evaluation_datasets.workspace(Line 29),experiments.workspace(Line 41),registered_models.workspace(Line 85), andwebhooks.workspace(Line 100) are unconstrained with respect toworkspaces. With a dedicatedworkspacestable (Lines 105-109), these should each have a FK like:ALTER TABLE evaluation_datasets ADD CONSTRAINT fk_evaluation_datasets_workspace FOREIGN KEY(workspace) REFERENCES workspaces(name);and analogous constraints for
experiments,registered_models, andwebhooks.Inconsistent
ON DELETE CASCADEfor model registry FKs
fk_model_versions_registered_models(Line 197),fk_registered_model_tags_registered_models(Line 217),fk_model_version_tags_model_versions(Line 366)only specify
ON UPDATE CASCADE, whilefk_registered_model_aliases_registered_models(Line 207) includesON DELETE CASCADE. For consistency and to prevent orphan rows, these three FKs should also includeON DELETE CASCADE.Please update the migration that generates this schema and then regenerate
mssql.sqlso these constraints are reflected here.Also applies to: 34-44, 80-87, 90-102, 105-109, 181-198, 201-208, 211-218, 359-367
mlflow/store/model_registry/dbmodels/models.py (1)
57-57: Add missing FK toworkspacesandondelete="cascade"on registry FKsTwo sets of integrity gaps called out earlier still exist:
Root models lack FK to
workspaces
SqlRegisteredModel.__table_args__(Line 57) defines only the composite PK.SqlWebhook.__table_args__(Lines 317-322) has PK + indexes only.With a
SqlWorkspace/workspacescatalog present,workspaceon these root tables should be constrained, e.g.:class SqlRegisteredModel(Base): ...
- table_args = (PrimaryKeyConstraint("workspace", "name", name="registered_model_pk"),)
- table_args = (
ForeignKeyConstraint(["workspace"],["workspaces.name"],name="fk_registered_models_workspace",),PrimaryKeyConstraint("workspace", "name", name="registered_model_pk"),- )
and similarly for `SqlWebhook`: ```diff class SqlWebhook(Base): ... - __table_args__ = ( - PrimaryKeyConstraint("webhook_id", name="webhook_pk"), - Index("idx_webhooks_status", "status"), - Index("idx_webhooks_name", "name"), - Index("idx_webhooks_workspace", "workspace"), - ) + __table_args__ = ( + ForeignKeyConstraint( + ["workspace"], + ["workspaces.name"], + name="fk_webhooks_workspace", + ), + PrimaryKeyConstraint("webhook_id", name="webhook_pk"), + Index("idx_webhooks_status", "status"), + Index("idx_webhooks_name", "name"), + Index("idx_webhooks_workspace", "workspace"), + )This prevents invalid workspace values on the root entities.
Model registry FKs still omit
ondelete="cascade"
SqlModelVersion.__table_args__FK toregistered_models(Lines 126-131).SqlRegisteredModelTag.__table_args__FK toregistered_models(Lines 176-181).SqlModelVersionTag.__table_args__FK tomodel_versions(Lines 211-216).To align DB-level behavior with the ORM
cascade="all"relationships and with other dialect schemas, these should addondelete="cascade":ForeignKeyConstraint( ["workspace", "name"], ["registered_models.workspace", "registered_models.name"],
onupdate="cascade",
onupdate="cascade",ondelete="cascade",and similarly for the other two constraints.Tightening these constraints will make workspace and registry behavior consistent across DB backends and direct SQL/ORM operations.
Also applies to: 126-133, 176-183, 211-218, 252-261, 301-306, 317-322
🧹 Nitpick comments (13)
tests/genai/evaluate/test_evaluation.py (1)
133-154: Backward‑compat check is sound; consider making the expectation a bit more explicitThe new assertion correctly enforces that
tables["eval_results"]exists and has the same row count asresult_df, which is sufficient for the stated backward‑compat goal. If you want to make the contract clearer and guard against future regressions, you could optionally add:
- An explicit membership check:
'eval_results' in result.tables- (Only if desired) a stronger assertion on contents, e.g.
result.tables["eval_results"].equals(result.result_df)Not required for correctness, but would better document the intended API behavior.
tests/store/tracking/test_sqlalchemy_store.py (1)
11352-11500: get_trace tests accurately cover success, not-found, partial, and complete cases; consider relaxing message matchingThe new tests exercise:
- Happy path where
log_spans+get_tracereturns spans with correct IDs and timing.- Not-found behavior mapping directly to the
_get_trace“not found” exception.- Partial traces with
TraceMetadataKey.SIZE_STATSboth whenallow_partialisTrue(returns 1 span) andFalse(retries then raises the “not fully exported yet” error).All of this aligns with the current
get_trace/_get_trace/_get_spans_with_trace_infosemantics and gives solid coverage.One minor robustness tweak you might consider: the tests that assert on exception messages (
test_get_trace_not_foundand the partial-trace branch) currently match almost the full message string. If those messages ever change slightly, otherwise-correct behavior will fail the tests. You could instead:
- Match only the invariant part of the message with a looser regex (e.g.,
r"Trace with ID .* is not found\."), or- Assert on the
error_codein addition to or instead of the full text, if exposed.Not strictly necessary, but it would make the tests less brittle to future wording changes.
mlflow/server/handlers.py (1)
667-674: Stricter bool query parsing looks good; consider factoring into a helper if reused more.Converting
"true"/"false"query params into actual booleans with validation is a nice improvement for GET handlers (e.g.,GetTrace.allow_partial). If more boolean fields show up in query strings later, you might want to centralize this pattern into a small helper to keep_get_request_messagelean, but it's not required right now.mlflow/store/tracking/rest_store.py (1)
80-81:RestStore.get_tracecorrectly targets the V3 endpoint; consider whether a fallback is needed.The implementation is consistent with the other V3 trace APIs (
/traces/search,/traces/batchGet): it builds aGetTracerequest includingallow_partial, hitsf"{_V3_TRACE_REST_API_PATH_PREFIX}/get", and returnsTrace.from_proto(response_proto.trace). That looks correct.If you expect this client to talk to older tracking servers that don’t implement
GetTrace, you may want a compatibility path similar toget_trace_info(e.g., detectENDPOINT_NOT_FOUNDand fall back tobatch_get_tracesor the previous trace‑info + artifact flow). If all call sites are version‑gated to new servers, the current behavior is fine.Also applies to: 466-472
tests/store/tracking/test_rest_store.py (1)
49-71: Nice coverage ofRestStore.get_trace; consider also asserting endpoint pathThe new tests correctly validate that
GetTraceis used and thatallow_partialis wired through in the JSON body. You may also want to assert theendpointkwarg in_call_endpoint(e.g. equalsf"{_V3_TRACE_REST_API_PATH_PREFIX}/get") to lock in the V3 path and catch future regressions in routing.Also applies to: 1310-1363
mlflow/tracing/otel/translation/__init__.py (1)
35-54:sanitize_attributescorrectly handles double‑dumped JSON; tighten scope to string valuesThe normalization logic for double‑dumped strings/dicts looks right and being applied up front in
translate_span_when_storingshould prevent redundant JSON encoding. To reduce unnecessary exceptions and make the intent clearer, you could first guard onisinstance(value, str)before callingjson.loads, and (if desired) decide whether lists should also be normalized the same way as dicts rather than remaining double‑encoded.Also applies to: 207-233
tests/db/test_schema.py (1)
11-16: Dialect‑aware reattachment logic is sound; minor helper cleanupThe
_reattach_missing_unique_constraintsflow looks correct for MySQL/MSSQL and should stabilize schema dumps across dialects._get_unique_constraintsdoesn’t use itsdialectargument, though—dropping that parameter (and its callers) or wiring it into behavior would simplify the API a bit.Also applies to: 46-107
mlflow/store/tracking/dbmodels/models.py (2)
64-65: Workspace column and per-workspace uniqueness look consistentUsing
DEFAULT_WORKSPACE_NAMEfor both ORM default andserver_defaultplus moving global name uniqueness toUniqueConstraint("workspace", "name")is consistent with the new schema. No functional issues here as long as all callers still treat the tracking store as single-workspace for now. If/when Experiment entities gain an explicit workspace field, this model will need to propagate it into_mlflow_entity/ constructors.Also applies to: 92-103, 131-138
1303-1323: Evaluation dataset workspace default is safe but not yet surfaced in APIAdding a non-null
workspacecolumn with defaultDEFAULT_WORKSPACE_NAMEand indexing it is aligned with the multi-tenant schema. However,SqlEvaluationDataset.to_mlflow_entity/from_mlflow_entitydon’t surface or accept workspace, so new datasets are effectively pinned to the default workspace. That’s fine for a first step, but once non-default workspaces are supported you’ll need to thread workspace through these methods (and through any callers) rather than relying solely on the DB default.Also applies to: 1375-1380, 1421-1444
mlflow/store/model_registry/sqlalchemy_store.py (1)
64-65: Registry writes are correctly workspace-qualified but hard-coded to the default
create_registered_model,create_model_version, and the variousset_*_tag/set_registered_model_aliashelpers now populateworkspace=DEFAULT_WORKSPACE_NAME, and tag/alias rows are created with matching workspace. That’s consistent with the new composite keys and keeps the current store effectively single-tenant.To avoid repeating the literal everywhere and to ease future evolution to multi-workspace support, consider centralizing workspace resolution (e.g., a
_get_workspace()helper or constructor parameter) rather than inliningDEFAULT_WORKSPACE_NAMEat each call site.Also applies to: 172-207, 686-707, 730-821, 1186-1212, 1247-1270
mlflow/genai/judges/instructions_judge/__init__.py (2)
31-37: Multi-turn session support is well-structured; consider surfacingsessionin input fields and polishing messagesThe new
{{ conversation }}variable,is_multi_turnproperty,sessionargument to__call__, and_validate_sessionlogic together give a clear and robust story:
- All traces in a session must share
TraceMetadataKey.TRACE_SESSION, with very helpful error messages (including a docs link) when that’s not true.- Conversation is derived via
resolve_conversation_from_sessionand injected into the template/user message only when present.- The template validator enforces that
{{ conversation }}can only be combined with{{ expectations }}, which keeps the prompting surface simple.Two improvements worth considering:
Expose
sessionviaget_input_fields
Right nowget_input_fields()only returns fields forinputs,outputs,expectations, andtrace. For templates that use{{ conversation }}, adding aJudgeField(name="session", description=...)would make it clearer to tooling and UI layers that this judge expects a session argument.Align warnings with the public API parameter name
_warn_unused_parameterswarns about"conversation"being unused, but from a caller’s perspective the parameter they pass issession. Either adjust the warning text to referencesessionor have_warn_unused_parametersaccept the originalsessionflag and infer whether it produced a conversation value, to avoid confusing users.These are ergonomics / API-clarity tweaks; the core design and validation look sound.
Also applies to: 52-63, 154-158, 189-219, 248-279, 364-425, 433-475, 550-575
220-247: Internal parameter naming aroundsessionvsconversationis a bit confusing
_check_required_parameters’s fifth parameter is still namedsession: list[Trace] | None, but__call__now passes the derivedconversationlist into that slot. Given the current implementation only checks forNonevs non-None, this works, but it obscures the real intent of the parameter and may trip up future maintainers or type-checking.Similarly,
_warn_unused_parametersoperates on the derivedconversationvalue, which is fine, but combined with the warning text (“parameters … 'conversation'”) it can be hard to trace back to the originalsessionargument.I’d suggest, when convenient:
- Either rename the
_check_required_parametersargument fromsessiontoconversation(and update the docstring), or change the call site to pass the originalsessionand check for that instead.- Consider passing both
sessionandconversationinto_warn_unused_parameters, so you can warn in terms of the public parameter name while still deciding based on the derived conversation’s presence.These are internal cleanups; they don’t affect runtime behavior today but will make the multi-turn path easier to reason about.
Also applies to: 280-308, 392-398, 433-438, 477-508, 550-575
mlflow/store/tracking/sqlalchemy_store.py (1)
788-799: Model metric logging behavior looks correct; consider avoiding duplicate run lookups.The new
_log_model_metricsflow correctly:
- Filters to
metric.model_id is not None.- Reuses
_validate_metricand metric de-duplication.- Sanitizes values via
sanitize_metric_valueto align with regular metric logging.- Persists
experiment_idand dataset fields throughSqlLoggedModelMetric.Two minor nits:
log_metricnow calls_get_run/_check_run_is_activetwice (once here, once in_log_metrics), which is redundant I/O; if this path becomes hot, consider refactoring to reuse the loadedrunor acceptexperiment_iddirectly into_log_metrics._log_model_metrics’sexperiment_idparameter is annotated asstrbut is passed anintfromrun.experiment_idinlog_batch/log_metric; consider widening the type hint toint | strfor static type sanity.Also applies to: 891-941, 1598-1612
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (41)
docs/api_reference/api_inventory.txt(4 hunks)docs/docs/self-hosting/security/basic-http-auth.mdx(4 hunks)mlflow/entities/__init__.py(2 hunks)mlflow/entities/workspace.py(1 hunks)mlflow/genai/evaluation/entities.py(1 hunks)mlflow/genai/judges/instructions_judge/__init__.py(17 hunks)mlflow/genai/judges/make_judge.py(2 hunks)mlflow/genai/scorers/base.py(3 hunks)mlflow/genai/utils/trace_utils.py(1 hunks)mlflow/java/client/src/main/java/org/mlflow/api/proto/Service.java(10 hunks)mlflow/protos/service.proto(1 hunks)mlflow/server/handlers.py(7 hunks)mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py(1 hunks)mlflow/store/model_registry/dbmodels/models.py(11 hunks)mlflow/store/model_registry/sqlalchemy_store.py(6 hunks)mlflow/store/tracking/abstract_store.py(1 hunks)mlflow/store/tracking/databricks_rest_store.py(1 hunks)mlflow/store/tracking/dbmodels/models.py(5 hunks)mlflow/store/tracking/rest_store.py(2 hunks)mlflow/store/tracking/sqlalchemy_store.py(6 hunks)mlflow/store/workspace/dbmodels/__init__.py(1 hunks)mlflow/store/workspace/dbmodels/models.py(1 hunks)mlflow/tracing/otel/translation/__init__.py(4 hunks)mlflow/tracing/otel/translation/base.py(1 hunks)mlflow/utils/workspace_utils.py(1 hunks)tests/db/check_migration.py(3 hunks)tests/db/schemas/mssql.sql(13 hunks)tests/db/schemas/mysql.sql(6 hunks)tests/db/schemas/postgresql.sql(6 hunks)tests/db/schemas/sqlite.sql(6 hunks)tests/db/test_schema.py(3 hunks)tests/db/test_workspace_migration.py(1 hunks)tests/genai/evaluate/test_evaluation.py(1 hunks)tests/genai/judges/test_make_judge.py(3 hunks)tests/resources/db/latest_schema.sql(6 hunks)tests/server/test_handlers.py(5 hunks)tests/store/tracking/test_rest_store.py(2 hunks)tests/store/tracking/test_sqlalchemy_store.py(2 hunks)tests/store/tracking/test_sqlalchemy_store_schema.py(1 hunks)tests/tracing/otel/test_span_translation.py(4 hunks)tests/tracking/test_rest_tracking.py(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- tests/genai/judges/test_make_judge.py
🚧 Files skipped from review as they are similar to previous changes (5)
- mlflow/store/workspace/dbmodels/init.py
- mlflow/store/workspace/dbmodels/models.py
- tests/store/tracking/test_sqlalchemy_store_schema.py
- mlflow/entities/workspace.py
- tests/db/test_workspace_migration.py
🧰 Additional context used
🧬 Code graph analysis (20)
mlflow/entities/__init__.py (1)
mlflow/entities/workspace.py (1)
Workspace(9-13)
tests/db/test_schema.py (1)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
_get_unique_constraints(150-164)
tests/store/tracking/test_sqlalchemy_store.py (1)
mlflow/store/tracking/sqlalchemy_store.py (4)
log_metric(788-798)log_spans(3301-3480)get_trace(3554-3567)start_trace(2566-2625)
mlflow/store/db_migrations/versions/1b5f0d9ad7c1_add_workspace_columns_and_catalog.py (1)
tests/db/test_schema.py (1)
_get_unique_constraints(90-107)
tests/db/check_migration.py (1)
mlflow/store/model_registry/sqlalchemy_store.py (4)
create_model_version(730-838)set_registered_model_alias(1247-1270)create_registered_model(172-214)create_webhook(1335-1367)
mlflow/genai/scorers/base.py (2)
mlflow/genai/judges/instructions_judge/__init__.py (1)
is_multi_turn(155-157)mlflow/protos/service_pb2.pyi (1)
Trace(983-989)
tests/server/test_handlers.py (6)
mlflow/exceptions.py (2)
MlflowException(67-113)MlflowNotImplementedException(215-219)mlflow/store/tracking/sqlalchemy_store.py (4)
_get_trace(3569-3594)get_trace(3554-3567)batch_get_traces(3596-3635)get_trace_info(2627-2639)mlflow/server/handlers.py (2)
_get_trace(2961-2976)get_trace_artifact_handler(3186-3212)mlflow/store/tracking/abstract_store.py (3)
get_trace(332-346)batch_get_traces(348-362)get_trace_info(320-330)mlflow/store/tracking/rest_store.py (3)
get_trace(466-471)batch_get_traces(473-488)get_trace_info(438-464)mlflow/tracing/client.py (2)
get_trace(138-201)get_trace_info(122-136)
mlflow/genai/judges/instructions_judge/__init__.py (2)
mlflow/genai/utils/trace_utils.py (4)
resolve_conversation_from_session(152-184)resolve_expectations_from_trace(187-213)resolve_inputs_from_trace(104-125)resolve_outputs_from_trace(128-149)mlflow/exceptions.py (2)
MlflowException(67-113)invalid_parameter_value(104-113)
mlflow/store/model_registry/sqlalchemy_store.py (1)
mlflow/store/model_registry/dbmodels/models.py (3)
SqlRegisteredModelTag(155-190)SqlModelVersionTag(193-230)SqlRegisteredModelAlias(233-268)
mlflow/genai/utils/trace_utils.py (3)
mlflow/protos/service_pb2.pyi (1)
Trace(983-989)mlflow/entities/trace.py (1)
Trace(26-332)mlflow/entities/span.py (2)
inputs(158-160)outputs(163-165)
mlflow/server/handlers.py (9)
mlflow/protos/service_pb2.pyi (1)
GetTrace(798-809)mlflow/exceptions.py (4)
MlflowException(67-113)invalid_parameter_value(104-113)MlflowTracingException(172-181)MlflowNotImplementedException(215-219)mlflow/store/tracking/sqlalchemy_store.py (4)
_get_trace(3569-3594)get_trace(3554-3567)batch_get_traces(3596-3635)get_trace_info(2627-2639)mlflow/store/tracking/abstract_store.py (3)
get_trace(332-346)batch_get_traces(348-362)get_trace_info(320-330)mlflow/store/tracking/rest_store.py (3)
get_trace(466-471)batch_get_traces(473-488)get_trace_info(438-464)mlflow/tracing/client.py (3)
get_trace(138-201)store(74-75)get_trace_info(122-136)mlflow/entities/trace.py (2)
to_proto(317-325)to_dict(44-45)mlflow/store/tracking/databricks_rest_store.py (2)
batch_get_traces(165-189)get_trace_info(191-213)mlflow/store/tracking/file_store.py (1)
get_trace_info(1755-1765)
mlflow/store/tracking/rest_store.py (5)
mlflow/java/client/src/main/java/org/mlflow/api/proto/Service.java (2)
GetTrace(90511-91932)Trace(111624-112580)mlflow/protos/service_pb2.pyi (2)
GetTrace(798-809)Trace(983-989)mlflow/store/tracking/abstract_store.py (1)
get_trace(332-346)mlflow/entities/trace.py (2)
Trace(26-332)from_proto(328-332)mlflow/utils/proto_json_utils.py (1)
message_to_json(112-129)
mlflow/store/tracking/databricks_rest_store.py (4)
mlflow/store/tracking/abstract_store.py (1)
batch_get_traces(348-362)mlflow/store/tracking/sqlalchemy_store.py (1)
batch_get_traces(3596-3635)mlflow/store/tracking/rest_store.py (1)
batch_get_traces(473-488)mlflow/entities/trace.py (1)
Trace(26-332)
mlflow/store/tracking/abstract_store.py (7)
mlflow/store/tracking/sqlalchemy_store.py (2)
get_trace(3554-3567)batch_get_traces(3596-3635)mlflow/store/tracking/rest_store.py (2)
get_trace(466-471)batch_get_traces(473-488)mlflow/tracking/client.py (1)
get_trace(1116-1145)mlflow/tracing/fluent.py (1)
get_trace(619-660)mlflow/tracing/client.py (1)
get_trace(138-201)mlflow/exceptions.py (1)
MlflowNotImplementedException(215-219)mlflow/store/tracking/databricks_rest_store.py (1)
batch_get_traces(165-189)
mlflow/store/tracking/sqlalchemy_store.py (8)
mlflow/server/handlers.py (1)
_get_trace(2961-2976)mlflow/entities/trace_info.py (4)
experiment_id(173-181)experiment_id(184-185)request_id(168-170)TraceInfo(20-239)mlflow/store/tracking/dbmodels/models.py (19)
SqlLoggedModelMetric(1115-1207)SqlTraceInfo(686-753)to_mlflow_entity(143-158)to_mlflow_entity(254-284)to_mlflow_entity(317-324)to_mlflow_entity(358-365)to_mlflow_entity(411-423)to_mlflow_entity(467-479)to_mlflow_entity(510-517)to_mlflow_entity(580-594)to_mlflow_entity(671-678)to_mlflow_entity(734-753)to_mlflow_entity(904-957)to_mlflow_entity(1084-1099)to_mlflow_entity(1197-1207)to_mlflow_entity(1252-1253)to_mlflow_entity(1298-1299)to_mlflow_entity(1382-1419)SqlSpan(1705-1791)mlflow/store/tracking/abstract_store.py (1)
get_trace(332-346)mlflow/tracing/client.py (1)
get_trace(138-201)mlflow/exceptions.py (2)
MlflowException(67-113)MlflowTracingException(172-181)mlflow/tracing/constant.py (2)
TraceTagKey(24-32)SpansLocation(134-136)mlflow/tracing/otel/translation/__init__.py (1)
translate_loaded_span(143-167)
tests/genai/evaluate/test_evaluation.py (2)
mlflow/genai/evaluation/entities.py (1)
tables(189-190)mlflow/models/evaluation/base.py (1)
tables(681-699)
tests/tracing/otel/test_span_translation.py (3)
mlflow/tracing/otel/translation/__init__.py (1)
sanitize_attributes(207-233)mlflow/tracing/otel/translation/vercel_ai.py (1)
VercelAITranslator(8-88)mlflow/entities/span.py (2)
SpanType(44-63)attributes(181-188)
tests/tracking/test_rest_tracking.py (4)
tests/telemetry/test_tracked_events.py (1)
mlflow_client(61-62)mlflow/tracking/_tracking_service/utils.py (1)
set_tracking_uri(70-121)mlflow/tracking/client.py (2)
tracking_uri(228-229)start_span(1368-1519)mlflow/tracing/fluent.py (2)
start_span(414-511)trace(63-204)
tests/store/tracking/test_rest_store.py (2)
mlflow/protos/service_pb2.pyi (2)
GetTrace(798-809)Trace(983-989)mlflow/store/tracking/rest_store.py (1)
get_trace(466-471)
mlflow/tracing/otel/translation/__init__.py (2)
mlflow/tracing/otel/translation/vercel_ai.py (1)
VercelAITranslator(8-88)mlflow/entities/span.py (3)
attributes(181-188)get(871-879)get(902-903)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (40)
- GitHub Check: core_tests (1, 3)
- GitHub Check: core_tests (2, 3)
- GitHub Check: core_tests (2, 5)
- GitHub Check: core_tests (1, 4)
- GitHub Check: core_tests (2, 4)
- GitHub Check: core_tests (1, 5)
- GitHub Check: pyfunc (2)
- GitHub Check: pyfunc (4)
- GitHub Check: pyfunc (3)
- GitHub Check: pyfunc (1)
- GitHub Check: evaluate (1)
- GitHub Check: genai
- GitHub Check: evaluate (2)
- GitHub Check: models (1)
- GitHub Check: java
- GitHub Check: python (1)
- GitHub Check: models (2)
- GitHub Check: flavors
- GitHub Check: python (2)
- GitHub Check: python (3)
- GitHub Check: windows (4)
- GitHub Check: windows (2)
- GitHub Check: windows (1)
- GitHub Check: python-skinny
- GitHub Check: windows (3)
- GitHub Check: database
- GitHub Check: r
- GitHub Check: build
- GitHub Check: typescript-sdk (20)
- GitHub Check: typescript-sdk (24)
- GitHub Check: typescript-sdk (22)
- GitHub Check: core
- GitHub Check: build (skinny)
- GitHub Check: build (dev)
- GitHub Check: build (tracing)
- GitHub Check: lint-macos
- GitHub Check: js (windows-latest, --testPathIgnorePatterns)
- GitHub Check: js (ubuntu-latest, --testPathPattern)
- GitHub Check: js (ubuntu-latest, --testPathIgnorePatterns)
- GitHub Check: js (windows-latest, --testPathPattern)
🔇 Additional comments (26)
docs/docs/self-hosting/security/basic-http-auth.mdx (1)
7-7: Scorer permissions documentation is well-integrated.The additions to support scorers in the authentication documentation are consistent with the existing patterns for experiments and registered models. The new "Required Permissions for accessing scorers" section (lines 584–627) is appropriately positioned and clearly documents the five scorer operations with their respective endpoints, HTTP methods, and permission requirements.
Please verify that the scorer endpoint versions (3.0) and permission mappings are accurate with the actual implementation changes in this PR.
Also applies to: 100-100, 584-627
mlflow/entities/__init__.py (2)
65-65: LGTM! Import correctly added.The import statement follows the existing pattern and is placed in the correct alphabetical position.
129-129: LGTM! Export correctly added.The
Workspaceentity is properly added to the__all__list, making it part of the public API.mlflow/protos/service.proto (1)
2567-2568: Web search didn't provide specific information about the default change. Let me search the repository directly for any comments or commit hints.Let me search for any release notes or changelog that might document this change.Based on my analysis of the code and search results, I cannot definitively verify the claim in the review comment. Here's what I found:
Current State:
- The proto at lines 2567-2568 shows:
optional bool allow_partial = 2 [default = false]- The RestStore Python method at line 466 defaults to:
allow_partial: bool = False- Both currently default to
falseEvidence from the Search:
- The handler at
mlflow/server/handlers.py:3160explicitly setsallow_partial=Truewith a comment about rendering in-progress traces—suggesting this is an intentional override for UI display- Tests show mixed patterns: some explicitly set
allow_partial(e.g., lines 2009, 2055), while test at line 2094 omits it- No entry in the MLflow 3 breaking changes documentation mentions this change
- Web search found no changelog entry documenting this as a breaking change
The Issue:
The review comment asserts the change was "fromdefault = truetodefault = false," but I cannot confirm the original default wastrueusing available tools. Without git history access or version-specific documentation, I cannot verify this factual claim.Verification Status:
While the current state does showdefault = false, and some callers do explicitly set the parameter, I cannot conclusively confirm whether this is actually a breaking change without evidence of what the previous default was.mlflow/java/client/src/main/java/org/mlflow/api/proto/Service.java (5)
90488-90506: LGTM: Documentation consistently updated.The documentation comments accurately reflect the new default value (
false) for theallow_partialfield across all interface methods, class fields, and builder methods.Also applies to: 91352-91373, 91824-91880
91591-91597: LGTM: Builder reset methods correctly updated.The
clear()andclearAllowPartial()methods now properly reset the field tofalse, consistent with the new default value.Also applies to: 91875-91880
91626-91634: LGTM: Correct optional field semantics.The
buildPartial()method now conditionally assignsallowPartial_only when explicitly set (bit field check). This ensures proper handling of optional fields with explicit presence tracking, where unset fields correctly default tofalsevia Java's default initialization.
187932-188478: LGTM: Protobuf descriptor updated.The serialized protobuf descriptor correctly reflects the
allow_partialdefault change fromtruetofalsein the GetTrace message definition. As generated code, this descriptor is automatically computed from the proto definition.
90488-91880: Now let me search for information specific to the MLflow proto definition and verify the actual change:I need to verify the actual proto file definition and understand whether this truly represents a breaking change. Let me search for the source proto file:
Based on my verification, the original review comment appears to contain incorrect analysis. Here's why:
In proto3, for bools the default value is false, and for numeric types the default value is zero. This is the implicit default that applies to all boolean fields in proto3.
The
[default = false]annotation visible in the Java generated code is documentation only, not a functional behavioral change. Proto3 only allows the implicit default value (the natural zero value for numbers/enums, false for booleans, or empty string for strings); no other default values are permitted.The claim that changing
allow_partialfromtruetofalserepresents a breaking change is misleading without verifying:
- What the actual proto file specifies
- Whether this field previously had a non-default value specified
- Whether proto2 or proto3 syntax is used
To properly verify this concern, the actual
.protofile definition would need to be examined to determine if this represents a real semantic change or if it's just documentation syntax in the generated Java code.
tests/store/tracking/test_sqlalchemy_store.py (1)
7675-7699: Logged-model metric test correctly validates experiment_id propagationThis test is well-constructed: it exercises
log_metricwith amodel_id, then inspectsSqlLoggedModelMetricto ensure theexperiment_idpersisted matches the run’s experiment (castingexp_idtointis correct given the DB schema). The flow matches the currentlog_metric/_log_model_metricsimplementation, so this gives good regression coverage for the new behavior.mlflow/utils/workspace_utils.py (1)
1-5: DEFAULT_WORKSPACE_NAME helper is clear and well-scopedCentralizing the default workspace name here and exporting it via
__all__looks good and will help keep DB models/migrations consistent.mlflow/genai/scorers/base.py (2)
120-127:is_multi_turnproperty is a good, conservative defaultExposing
is_multi_turnas a read-only property with a default ofFalseis a clean way to let subclasses opt into multi‑turn behavior without affecting existing scorers.
443-505: Verify whether multi-turn evaluation flows throughScorer.run()or bypass itVerification confirms the gap exists:
Scorer.run()at line 383 constructs a merged dict with onlyinputs,outputs,expectations, andtrace—nosessionparameter. Both call sites (optimize/util.py and evaluation/utils.py) do not passsessionthroughrun().However,
__call__at lines 443–505 does acceptsession. The filtering logic inrun()means session is never forwarded to__call__even if a scorer expects it.To resolve this, either:
- Update
run()to accept and forwardsessionto__call__, OR- Confirm that all multi-turn scorers call
__call__directly rather than throughrun()Please verify which path aligns with your multi-turn evaluation design.
tests/tracking/test_rest_tracking.py (1)
2848-2871:test_get_trace_handlerprovides solid end-to-end coverageNice integration test: it exercises the new
/ajax-api/3.0/mlflow/traces/getendpoint for bothallow_partial=True/False, verifies status, trace ID, span count/name, and attributes. This should catch regressions in the REST wiring and JSON shape.mlflow/genai/judges/make_judge.py (1)
113-117: Documentation and validation constraints verifiedThe underlying
InstructionsJudgecode properly enforces the documented constraints: lines 566–572 inmlflow/genai/judges/instructions_judge/__init__.pyvalidate that{{ conversation }}can only coexist with{{ expectations }}, rejecting any combination with{{ inputs }},{{ outputs }}, or{{ trace }}. This aligns with the updated docstring and multi-turn example.mlflow/server/handlers.py (2)
2959-2976: _get_trace handler matches store API and proto contract.The new
_get_traceendpoint correctly validatestrace_idandallow_partial, delegates toAbstractTrackingStore.get_trace, and returnsGetTrace.Response(trace=trace.to_proto()), which aligns with the abstract store and proto definitions. I don’t see functional issues here.
4143-4143: Handler registration forGetTraceis consistent with other tracing APIs.Wiring
GetTraceto_get_traceinHANDLERSkeeps the service surface consistent with the new V3 tracing APIs. No issues here.tests/db/check_migration.py (1)
52-61: Workspace test coverage via seeding looks comprehensive.Defining
WORKSPACE_TABLESand then explicitly exercising all of them inlog_everything()(extra experiment + registered model, alias, webhook, and direct insert intoevaluation_datasets) is a good way to ensure the migration actually backfills real rows across all workspace‑backed tables rather than only the “happy path” ones. The raw SQL insert againstevaluation_datasetsis scoped to a single row and wrapped inengine.begin(), so it won’t interfere with the higher‑level ORM tests.Also applies to: 69-127
tests/tracing/otel/test_span_translation.py (1)
10-11: Vercel span kind coverage andsanitize_attributestests align well with the translation logic.Including
VercelAITranslatorintest_json_serialized_valuesand assertingai.generateText→LLM,ai.toolCall→TOOLvia the sharedSPAN_KIND_ATTRIBUTE_KEYexercises the new translator wiring nicely. The newtest_sanitize_attributescases mirror the implementation’s behavior for double‑dumped strings/dicts and primitive JSON types, so regressions in that normalization logic should be caught early. I don’t see any issues here.Also applies to: 20-21, 69-76, 335-367
mlflow/store/tracking/abstract_store.py (1)
332-347: Trace retrieval API surface looks consistent and future‑proof
get_traceand the updatedbatch_get_tracessignature align with the concrete store implementations (SQLAlchemy, REST, Databricks) and useMlflowNotImplementedExceptionto surface 501s cleanly; the docstrings accurately describeallow_partialandlocation. No changes needed.Also applies to: 348-363
tests/server/test_handlers.py (3)
18-24: Imports and handler wiring for traces look correctBringing in
MlflowNotImplementedException,RESOURCE_DOES_NOT_EXIST,GetTrace,_get_trace, andget_trace_artifact_handleraligns the tests with the new trace APIs and error semantics. Nothing to change here.Also applies to: 48-64, 72-103, 128-129
2006-2111: _get_trace handler tests cover success and not-found paths appropriatelyThe new
test_get_trace_handlerandtest_get_trace_handler_with_allow_partial_falsecorrectly exercise propagation of theallow_partialflag intoget_traceand validate response serialization. The not-found test asserts a 404 withRESOURCE_DOES_NOT_EXIST, matching the expectedMlflowException→ HTTP mapping. These look solid.
2113-2311: Comprehensive coverage of get_trace_artifact_handler fallback behaviorThe suite of
get_trace_artifact_handler*tests nicely exercises:
- Happy path when
get_trace(..., allow_partial=True)is implemented.- 400 when
request_idis missing.- 404 when
get_traceraisesRESOURCE_DOES_NOT_EXIST.- Fallback to
batch_get_traceswhenget_traceraisesMlflowNotImplementedException.- 404 when
batch_get_tracesreturns an empty list.- Final fallback to
get_trace_info+_get_trace_artifact_repo().download_trace_data()when both trace APIs raiseMlflowNotImplementedException.Two minor notes:
- The tests assume
_fetch_trace_data_from_storecallsget_trace(trace_id, allow_partial=True)first and thenbatch_get_traces([trace_id], None); any change in that helper’s contract will require updating these expectations.- The assertions on error messages (
"Trace with id={trace_id} not found") effectively bake the exact message into the public contract; keep that stable or relax the match if you expect wording changes.Otherwise, the scenarios and mocks line up well with the intended handler behavior.
tests/db/schemas/postgresql.sql (1)
19-31: PostgreSQL schema matches workspace-aware ORM changesThe added
workspace VARCHAR(63) DEFAULT 'default' NOT NULLcolumns and updated keys/constraints for:
evaluation_datasets,experiments,registered_models,model_versions,registered_model_aliases,registered_model_tags,model_version_tags,- plus
webhooksand the newworkspacestableare all consistent with the workspace-qualified ORM definitions and composite keys in the model registry and tracking models.
One thing to keep in mind is the tight coupling between the literal
'default'here andDEFAULT_WORKSPACE_NAME; if that constant ever changes, these schema snapshots and migrations will need to be regenerated together. Otherwise the DDL looks coherent.Also applies to: 34-45, 81-88, 91-103, 106-110, 184-201, 204-211, 214-221, 365-372
docs/api_reference/api_inventory.txt (1)
531-531: Workspace and scorer permission APIs correctly surfaced in public inventoryThe new
Workspaceentities and scorer permission client/entity symbols are consistently listed and aligned with existing naming patterns. No issues from the API-inventory perspective.Also applies to: 631-631, 1182-1208
mlflow/store/model_registry/dbmodels/models.py (1)
36-37: Workspace-aware registry and webhook models look structurally consistentThe introduction of
workspacecolumns withDEFAULT_WORKSPACE_NAMEand composite primary keys acrossSqlRegisteredModel,SqlModelVersion,SqlRegisteredModelTag,SqlModelVersionTag, andSqlRegisteredModelAlias, plus the workspace index onSqlWebhook, all align with the workspace-scoped schemas in the test DDL. The ORM-side relationships (model_versions,registered_model_tags,model_version_tags,registered_model_aliases) remain coherent with the new composite keys.Also applies to: 42-47, 49-57, 88-96, 126-133, 158-166, 176-183, 196-204, 211-223, 236-243, 252-261, 301-306, 317-322
| <tr> | ||
| <td>`admin`</td> | ||
| <td>`password`</td> | ||
| <td>`password1234`</td> |
There was a problem hiding this comment.
Update example code to use the correct admin password.
The default admin password table on line 895 shows password1234, but the example code on line 905 still references the old password password. This inconsistency could cause authentication failures for users following the documentation.
Apply this diff to fix the example:
# authenticate as built-in admin user
export MLFLOW_TRACKING_USERNAME=admin
- export MLFLOW_TRACKING_PASSWORD=password
+ export MLFLOW_TRACKING_PASSWORD=password1234Also applies to: 905-905
🤖 Prompt for AI Agents
docs/docs/self-hosting/security/basic-http-auth.mdx lines 895 and 905: the
default admin password in the table on line 895 is `password1234` but the
example code on line 905 still uses `password`; update the example on line 905
to use `password1234` so both match, and scan nearby lines for any other
references to `password` and replace them with `password1234` to keep the
documentation consistent.
| def resolve_conversation_from_session( | ||
| session: list[Trace], | ||
| ) -> list[dict[str, str]]: | ||
| """ | ||
| Extract conversation history from traces in session. | ||
|
|
||
| Args: | ||
| session: List of traces from the same session. | ||
|
|
||
| Returns: | ||
| List of conversation messages in the format [{"role": "user"|"assistant", "content": str}]. | ||
| Each trace contributes two messages: user (from input) and assistant (from output). | ||
| """ | ||
| # Sort traces by creation time (timestamp_ms) | ||
| sorted_traces = sorted(session, key=lambda t: t.info.timestamp_ms) | ||
|
|
||
| conversation = [] | ||
| for trace in sorted_traces: | ||
| # Extract and parse input (user message) | ||
| inputs = extract_inputs_from_trace(trace) | ||
| if inputs: | ||
| user_content = parse_inputs_to_str(inputs) | ||
| if user_content and user_content.strip(): | ||
| conversation.append({"role": "user", "content": user_content}) | ||
|
|
||
| # Extract and parse output (assistant message) | ||
| outputs = extract_outputs_from_trace(trace) | ||
| if outputs: | ||
| assistant_content = parse_outputs_to_str(outputs) | ||
| if assistant_content and assistant_content.strip(): | ||
| conversation.append({"role": "assistant", "content": assistant_content}) | ||
|
|
||
| return conversation | ||
|
|
There was a problem hiding this comment.
Avoid dropping valid but falsy inputs/outputs when building the conversation
resolve_conversation_from_session currently does if inputs: / if outputs:, so values like 0, False, {}, or [] will be skipped even though they may serialize to meaningful content. You already filter out empty/whitespace after string conversion, so it’s safer to only guard against None at this stage.
Consider:
- for trace in sorted_traces:
- # Extract and parse input (user message)
- inputs = extract_inputs_from_trace(trace)
- if inputs:
- user_content = parse_inputs_to_str(inputs)
- if user_content and user_content.strip():
- conversation.append({"role": "user", "content": user_content})
+ for trace in sorted_traces:
+ # Extract and parse input (user message)
+ inputs = extract_inputs_from_trace(trace)
+ if inputs is not None:
+ user_content = parse_inputs_to_str(inputs)
+ if user_content and user_content.strip():
+ conversation.append({"role": "user", "content": user_content})
@@
- # Extract and parse output (assistant message)
- outputs = extract_outputs_from_trace(trace)
- if outputs:
- assistant_content = parse_outputs_to_str(outputs)
- if assistant_content and assistant_content.strip():
- conversation.append({"role": "assistant", "content": assistant_content})
+ # Extract and parse output (assistant message)
+ outputs = extract_outputs_from_trace(trace)
+ if outputs is not None:
+ assistant_content = parse_outputs_to_str(outputs)
+ if assistant_content and assistant_content.strip():
+ conversation.append({"role": "assistant", "content": assistant_content})This preserves legitimate falsy values while still excluding empty/whitespace-only messages.
🤖 Prompt for AI Agents
In mlflow/genai/utils/trace_utils.py around lines 152 to 185, the current guards
use `if inputs:` and `if outputs:` which drop valid but falsy values (e.g., 0,
False, {}, []). Change those checks to explicitly test for None (e.g., `if
inputs is not None:` and `if outputs is not None:`) so falsy-but-valid values
are passed to the parsing functions; keep the existing conversion and
`strip()`-based empty-string filtering after parsing.
| def _fetch_trace_data_from_store( | ||
| store: AbstractTrackingStore, request_id: str | ||
| ) -> dict[str, Any] | None: | ||
| try: | ||
| # allow partial so the frontend can render in-progress traces | ||
| trace = store.get_trace(request_id, allow_partial=True) | ||
| return trace.data.to_dict() | ||
| except MlflowTracingException: | ||
| return None | ||
| except MlflowNotImplementedException: | ||
| # fallback to batch_get_traces if get_trace is not implemented | ||
| pass | ||
|
|
||
| try: | ||
| traces = store.batch_get_traces([request_id], None) | ||
| match traces: | ||
| case [trace]: | ||
| return trace.data.to_dict() | ||
| case _: | ||
| raise MlflowException( | ||
| f"Trace with id={request_id} not found.", | ||
| error_code=RESOURCE_DOES_NOT_EXIST, | ||
| ) | ||
| # For stores that don't support batch get traces, or if trace data is not in the store, | ||
| # return None to signal fallback to artifact repository | ||
| except (MlflowTracingException, MlflowNotImplementedException): | ||
| return None |
There was a problem hiding this comment.
🧩 Analysis chain
Fallback behavior from get_trace to batch_get_traces / artifacts may be too narrow for older servers.
The overall flow (prefer store.get_trace(allow_partial=True), fall back to batch_get_traces, then finally to the artifact repo) is solid, but there are a couple of edge cases worth tightening:
-
Older tracking servers without GetTrace:
- For RestStore‑backed tracking stores talking to an older server that doesn’t support
GetTrace,store.get_trace()will likely raise a genericMlflowExceptionwitherror_code=ENDPOINT_NOT_FOUND, notMlflowNotImplementedException. _fetch_trace_data_from_storecurrently only treatsMlflowNotImplementedExceptionas “feature not available”, soENDPOINT_NOT_FOUNDwill bubble out and prevent fallback tobatch_get_traces/ artifacts, even though those older paths might still work.- Consider treating
MlflowExceptionwithENDPOINT_NOT_FOUNDlike “not implemented” here so the new handler degrades gracefully on older backends.
- For RestStore‑backed tracking stores talking to an older server that doesn’t support
-
Distinguishing “trace truly missing” from “data not in store”:
- When
batch_get_tracesreturns anything other than a single trace, you immediately raiseMlflowException("Trace with id=... not found", RESOURCE_DOES_NOT_EXIST). - If there are scenarios where
batch_get_tracesreturns an empty list whileget_trace_info+ artifact JSON still work (e.g., older traces whose spans live only in artifacts), this will now short‑circuit the artifact fallback you added. - If that scenario is supported, you may want
batch_get_tracesto signal “no DB data but trace exists” viaMlflowTracingException, or you can explicitly check fortraces == []and returnNoneto letget_trace_artifact_handlerfall back.
- When
A minimal adjustment for (1) could look like:
- except MlflowTracingException:
- return None
- except MlflowNotImplementedException:
- # fallback to batch_get_traces if get_trace is not implemented
- pass
+ except MlflowTracingException:
+ return None
+ except MlflowNotImplementedException:
+ # fallback to batch_get_traces if get_trace is not implemented
+ pass
+ except MlflowException as e:
+ # Older tracking servers may not implement GetTrace and return ENDPOINT_NOT_FOUND.
+ if e.error_code == databricks_pb2.ErrorCode.Name(databricks_pb2.ENDPOINT_NOT_FOUND):
+ return None
+ raiseThis keeps real server‑side failures visible while preserving the intended fallback for unsupported endpoints.
Also applies to: 3186-3212
I'll verify this review comment by examining the codebase to understand the error handling behavior and whether the suggested fix is necessary and correct.
---
🏁 Script executed:
```shell
#!/bin/bash
# Locate and inspect the handlers.py file to see the current implementation
fd -e "handlers.py" -p "mlflow/server" --exec cat {} \;
Length of output: 0
🏁 Script executed:
#!/bin/bash
# Search for RestStore implementation to understand how it handles older servers
fd -e "py" -p "mlflow" | xargs rg -l "class RestStore" | head -5 | xargs -I {} sh -c 'echo "=== {} ===" && cat "{}"'Length of output: 88081
🏁 Script executed:
#!/bin/bash
# Check for ENDPOINT_NOT_FOUND error code usage in the codebase
rg -n "ENDPOINT_NOT_FOUND" --type py | head -20Length of output: 2360
🏁 Script executed:
#!/bin/bash
# Search for databricks_pb2 imports and ErrorCode references
rg -n "databricks_pb2.*ErrorCode" --type py | head -20Length of output: 2539
🏁 Script executed:
#!/bin/bash
# Look for existing patterns of handling MlflowException with specific error codes
rg -n "MlflowException.*error_code" --type py -A 3 | head -40Length of output: 1957
---
🏁 Script executed:
```shell
#!/bin/bash
# Find handlers.py using ripgrep
rg -l "handlers\.py" --type py | head -5
Length of output: 0
🏁 Script executed:
#!/bin/bash
# Search for _fetch_trace_data_from_store function
rg -n "_fetch_trace_data_from_store" --type py -A 30Length of output: 4085
🏁 Script executed:
#!/bin/bash
# Check get_trace method in RestStore for exception handling
rg -n "def get_trace" mlflow/store/tracking/rest_store.py -A 10Length of output: 1011
🏁 Script executed:
#!/bin/bash
# Check batch_get_traces return behavior for empty traces
rg -n "def batch_get_traces" mlflow/store/tracking/rest_store.py -A 10Length of output: 432
🏁 Script executed:
#!/bin/bash
# Check if handlers.py exists and list its location
find . -name "handlers.py" -path "*/mlflow/server/*" 2>/dev/null | head -5Length of output: 27
---
🏁 Script executed:
```shell
#!/bin/bash
# Check _call_endpoint implementation in rest_store.py
rg -n "def _call_endpoint" mlflow/store/tracking/rest_store.py -A 15
Length of output: 604
🏁 Script executed:
#!/bin/bash
# Check how get_trace_info handles ENDPOINT_NOT_FOUND (lines 438-465)
sed -n '438,465p' mlflow/store/tracking/rest_store.pyLength of output: 1333
🏁 Script executed:
#!/bin/bash
# Check exception imports in handlers.py
rg -n "^from mlflow.*import|^import" mlflow/server/handlers.py | grep -E "(Exception|Error)" | head -20Length of output: 82
🏁 Script executed:
#!/bin/bash
# Search for MlflowNotImplementedException to understand when it's raised
rg -n "class MlflowNotImplementedException" --type py -A 5Length of output: 338
🏁 Script executed:
#!/bin/bash
# Check if ENDPOINT_NOT_FOUND is imported in handlers.py
rg -n "ENDPOINT_NOT_FOUND" mlflow/server/handlers.pyLength of output: 0
🏁 Script executed:
#!/bin/bash
# Look at call_endpoint function to see how it handles missing endpoints
rg -n "def call_endpoint" mlflow/utils/rest_utils.py -A 30Length of output: 2880
---
🏁 Script executed:
```shell
#!/bin/bash
# Check verify_rest_response function
rg -n "def verify_rest_response" mlflow/utils/rest_utils.py -A 40
Length of output: 1929
🏁 Script executed:
#!/bin/bash
# Search for test cases using batch_get_traces
rg -n "batch_get_traces" --type py -B 2 -A 5 | head -60Length of output: 4225
🏁 Script executed:
#!/bin/bash
# Check AbstractTrackingStore interface for batch_get_traces signature
rg -n "def batch_get_traces" mlflow/store/tracking/abstract_store.py -A 10Length of output: 445
🏁 Script executed:
#!/bin/bash
# Look for handlers that call _fetch_trace_data_from_store
rg -n "get_trace_artifact_handler\|_fetch_trace_data_from_store" mlflow/server/handlers.py -B 5 -A 10 | head -50Length of output: 0
🏁 Script executed:
#!/bin/bash
# Check RestException class and how ENDPOINT_NOT_FOUND is set
rg -n "class RestException" mlflow/utils/rest_utils.py -A 20Length of output: 0
---
🏁 Script executed:
```shell
#!/bin/bash
# Search for handler function that uses _fetch_trace_data_from_store
rg -n "_fetch_trace_data_from_store" mlflow/server/handlers.py -B 20 | tail -40
Length of output: 1506
🏁 Script executed:
#!/bin/bash
# Get the full context of the handler that calls _fetch_trace_data_from_store
sed -n '3183,3215p' mlflow/server/handlers.pyLength of output: 1039
🏁 Script executed:
#!/bin/bash
# Search for RestException definition and error_code
rg -n "class RestException\|raise RestException" mlflow/utils/rest_utils.py -A 3 | head -50Length of output: 0
🏁 Script executed:
#!/bin/bash
# Check imports in mlflow/utils/rest_utils.py for RestException
rg -n "RestException|from.*RestException" mlflow/utils/rest_utils.py | head -10Length of output: 123
🏁 Script executed:
#!/bin/bash
# Check what MlflowException gets raised from call_endpoint
rg -n "MlflowException.*ENDPOINT_NOT_FOUND\|RestException" mlflow/utils/rest_utils.py -B 3 -A 3 | head -60Length of output: 0
🏁 Script executed:
#!/bin/bash
# Trace exception conversion from RestException to MlflowException
rg -n "except RestException" mlflow/ --type py -A 3 | head -40Length of output: 1955
---
🏁 Script executed:
```shell
#!/bin/bash
# Find where RestException is imported from
rg -n "from.*RestException\|import.*RestException" mlflow/ --type py | head -10
Length of output: 0
🏁 Script executed:
#!/bin/bash
# Check the exception module for RestException
rg -n "class RestException" mlflow/ --type py -A 5Length of output: 316
🏁 Script executed:
#!/bin/bash
# Look at how call_endpoint handles exceptions (full implementation)
sed -n '568,610p' mlflow/utils/rest_utils.pyLength of output: 1378
🏁 Script executed:
#!/bin/bash
# Check if RestException inherits from MlflowException
rg -n "RestException" mlflow/exceptions.py -B 2 -A 10Length of output: 913
🏁 Script executed:
#!/bin/bash
# Search for where exceptions are caught and re-raised in tracking
rg -n "except.*Exception.*get_trace\|except.*Exception.*batch_get" mlflow/store/tracking/rest_store.py -B 2 -A 5Length of output: 0
🏁 Script executed:
#!/bin/bash
# Check test file for _fetch_trace_data_from_store behavior
rg -n "_fetch_trace_data_from_store\|get_trace_artifact_handler" tests/ --type py -B 3 -A 10Length of output: 0
Fallback behavior from get_trace to batch_get_traces / artifacts needs to handle ENDPOINT_NOT_FOUND and empty responses.
Analysis of the code confirms two edge cases require handling:
-
Older servers without GetTrace: When
store.get_trace()is called against an older MLflow server lacking the GetTrace endpoint, it raisesRestException(subclass ofMlflowException) witherror_code=ENDPOINT_NOT_FOUND. The function currently only catchesMlflowNotImplementedException, so this exception propagates instead of falling back tobatch_get_traces. The pattern inRestStore.get_trace_info()(line 457 ofrest_store.py) handles this correctly by explicitly checking forENDPOINT_NOT_FOUND. -
Empty batch_get_traces responses: When
batch_get_traces()returns an empty list (possible for traces with spans only in artifacts), the currentmatchstatement raisesMlflowExceptioninstead of returningNoneto trigger the artifact repository fallback at line 3199.
Add a handler for ENDPOINT_NOT_FOUND in the first exception block and check for empty traces in the batch_get_traces branch:
- except MlflowTracingException:
- return None
- except MlflowNotImplementedException:
- # fallback to batch_get_traces if get_trace is not implemented
- pass
+ except MlflowTracingException:
+ return None
+ except MlflowNotImplementedException:
+ # fallback to batch_get_traces if get_trace is not implemented
+ pass
+ except MlflowException as e:
+ # Older tracking servers may not implement GetTrace and return ENDPOINT_NOT_FOUND.
+ if e.error_code == databricks_pb2.ErrorCode.Name(databricks_pb2.ENDPOINT_NOT_FOUND):
+ pass # fallback to batch_get_traces
+ else:
+ raise
try:
traces = store.batch_get_traces([request_id], None)
match traces:
case [trace]:
return trace.data.to_dict()
- case _:
- raise MlflowException(
- f"Trace with id={request_id} not found.",
- error_code=RESOURCE_DOES_NOT_EXIST,
- )
+ case []:
+ # Empty result from batch_get_traces; fallback to artifact repository
+ return None
+ case _:
+ raise MlflowException(
+ f"Trace with id={request_id} not found.",
+ error_code=RESOURCE_DOES_NOT_EXIST,
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def _fetch_trace_data_from_store( | |
| store: AbstractTrackingStore, request_id: str | |
| ) -> dict[str, Any] | None: | |
| try: | |
| # allow partial so the frontend can render in-progress traces | |
| trace = store.get_trace(request_id, allow_partial=True) | |
| return trace.data.to_dict() | |
| except MlflowTracingException: | |
| return None | |
| except MlflowNotImplementedException: | |
| # fallback to batch_get_traces if get_trace is not implemented | |
| pass | |
| try: | |
| traces = store.batch_get_traces([request_id], None) | |
| match traces: | |
| case [trace]: | |
| return trace.data.to_dict() | |
| case _: | |
| raise MlflowException( | |
| f"Trace with id={request_id} not found.", | |
| error_code=RESOURCE_DOES_NOT_EXIST, | |
| ) | |
| # For stores that don't support batch get traces, or if trace data is not in the store, | |
| # return None to signal fallback to artifact repository | |
| except (MlflowTracingException, MlflowNotImplementedException): | |
| return None | |
| def _fetch_trace_data_from_store( | |
| store: AbstractTrackingStore, request_id: str | |
| ) -> dict[str, Any] | None: | |
| try: | |
| # allow partial so the frontend can render in-progress traces | |
| trace = store.get_trace(request_id, allow_partial=True) | |
| return trace.data.to_dict() | |
| except MlflowTracingException: | |
| return None | |
| except MlflowNotImplementedException: | |
| # fallback to batch_get_traces if get_trace is not implemented | |
| pass | |
| except MlflowException as e: | |
| # Older tracking servers may not implement GetTrace and return ENDPOINT_NOT_FOUND. | |
| if e.error_code == databricks_pb2.ErrorCode.Name(databricks_pb2.ENDPOINT_NOT_FOUND): | |
| pass # fallback to batch_get_traces | |
| else: | |
| raise | |
| try: | |
| traces = store.batch_get_traces([request_id], None) | |
| match traces: | |
| case [trace]: | |
| return trace.data.to_dict() | |
| case []: | |
| # Empty result from batch_get_traces; fallback to artifact repository | |
| return None | |
| case _: | |
| raise MlflowException( | |
| f"Trace with id={request_id} not found.", | |
| error_code=RESOURCE_DOES_NOT_EXIST, | |
| ) | |
| # For stores that don't support batch get traces, or if trace data is not in the store, | |
| # return None to signal fallback to artifact repository | |
| except (MlflowTracingException, MlflowNotImplementedException): | |
| return None |
🤖 Prompt for AI Agents
In mlflow/server/handlers.py around lines 3155 to 3181, the fallback from
get_trace to batch_get_traces must handle servers that return an MlflowException
with error_code ENDPOINT_NOT_FOUND and treat empty batch_get_traces responses as
missing trace data; update the first except to also catch MlflowException and,
when caught, check if exc.error_code == ENDPOINT_NOT_FOUND then fall through to
batch_get_traces (re-raise otherwise), and in the batch_get_traces branch, if
traces is an empty list return None instead of raising an MlflowException so the
caller will fall back to artifact repository.
| def get_trace(self, trace_id: str, *, allow_partial: bool = False) -> Trace: | ||
| if not allow_partial: | ||
| for retry_count in range(3): | ||
| # only retry if the spans are not fully exported | ||
| if trace := self._get_trace(trace_id, allow_partial): | ||
| return trace | ||
| elif retry_count < 2: | ||
| time.sleep(2**retry_count) | ||
| raise MlflowException( | ||
| message=f"Trace with ID {trace_id} is not fully exported yet, " | ||
| "please try again later.", | ||
| error_code=RESOURCE_DOES_NOT_EXIST, | ||
| ) | ||
| return self._get_trace(trace_id, allow_partial) |
There was a problem hiding this comment.
🧩 Analysis chain
Trace retrieval mostly matches allow_partial semantics; edge cases around zero spans and non‑tracking locations may need clarification.
Behavior is generally sound (retry loop for allow_partial=False, completeness check via SIZE_STATS, and centralized span translation), but there are two subtle points to verify:
- When
allow_partial=False,_get_tracenever returns a trace ifspansis empty, even whenSIZE_STATSindicates 0 spans or no stats are present; such traces are always treated as “not fully exported yet”. If 0‑span traces are considered valid in your model, this will surface as spurious “try again later” errors instead of returning an empty trace. _get_spans_with_trace_inforaisesMlflowTracingExceptionwhenSPANS_LOCATION != TRACKING_STORE.get_trace/batch_get_tracesdon’t catch this, so callers must either rely on that exception or handle artifact‑repo fallbacks themselves; the comment about “handler can capture the exception and load data from artifact repo instead” suggests there may be missing or external handling to confirm.
If 0‑span traces should be returnable and handler‑side fallback for non‑tracking locations is expected, a small follow‑up adjustment might be warranted.
Also applies to: 3569-3594, 3616-3634, 3637-3670
I'll verify the review comment by examining the actual implementation of the tracing methods and their handling of edge cases.
🏁 Script executed:
#!/bin/bash
# Examine the get_trace, _get_trace, and related methods around the specified lines
sed -n '3554,3670p' mlflow/store/tracking/sqlalchemy_store.pyLength of output: 5032
🏁 Script executed:
#!/bin/bash
# Look at the _get_trace implementation to see how it handles empty spans
grep -n "_get_trace" mlflow/store/tracking/sqlalchemy_store.py | head -20Length of output: 1047
🏁 Script executed:
#!/bin/bash
# Search for _get_spans_with_trace_info and its exception handling
rg -n "_get_spans_with_trace_info" -A 10 -B 2 mlflow/store/tracking/sqlalchemy_store.pyLength of output: 1809
🏁 Script executed:
#!/bin/bash
# Look for SIZE_STATS and its usage pattern
rg -n "SIZE_STATS" -B 3 -A 3 mlflow/store/tracking/sqlalchemy_store.pyLength of output: 381
🏁 Script executed:
#!/bin/bash
# Search for SPANS_LOCATION and TRACKING_STORE logic
rg -n "SPANS_LOCATION|TRACKING_STORE" -B 2 -A 2 mlflow/store/tracking/sqlalchemy_store.pyLength of output: 946
🏁 Script executed:
#!/bin/bash
# Check for exception handling around batch_get_traces
grep -n "batch_get_traces" mlflow/store/tracking/sqlalchemy_store.py | head -5Length of output: 193
🏁 Script executed:
#!/bin/bash
# Look for test files that might clarify expected behavior
fd -e py "test.*trace" mlflow/store/tracking/ | head -10Length of output: 0
🏁 Script executed:
#!/bin/bash
# Search for MlflowTracingException handling patterns
rg -n "MlflowTracingException" -B 2 -A 2 mlflow/store/tracking/sqlalchemy_store.pyLength of output: 548
🏁 Script executed:
#!/bin/bash
# Look for any comments about handlers or artifact repo fallbacks
rg -n "handler|artifact.*repo|fallback" -i -B 2 -A 2 mlflow/store/tracking/sqlalchemy_store.py | head -40Length of output: 969
🏁 Script executed:
#!/bin/bash
# Examine batch_get_traces around lines 3569-3594
sed -n '3569,3594p' mlflow/store/tracking/sqlalchemy_store.pyLength of output: 1176
🏁 Script executed:
#!/bin/bash
# Check lines 3616-3634 and 3637-3670
sed -n '3616,3670p' mlflow/store/tracking/sqlalchemy_store.pyLength of output: 2493
🏁 Script executed:
#!/bin/bash
# Look for how empty spans are handled in _get_trace
rg -n "def _get_trace" -A 40 mlflow/store/tracking/sqlalchemy_store.pyLength of output: 7536
🏁 Script executed:
#!/bin/bash
# Search for the trace translation and span collection logic
rg -n "def _spans_to_trace|_collect.*span" -A 15 mlflow/store/tracking/sqlalchemy_store.pyLength of output: 0
🏁 Script executed:
#!/bin/bash
# Look for SIZE_STATS in the context of completeness checking
rg -n "completeness|SIZE_STATS.*0|spans.*empty" -i -B 3 -A 3 mlflow/store/tracking/sqlalchemy_store.pyLength of output: 0
🏁 Script executed:
#!/bin/bash
# Check for any documentation about valid traces with 0 spans
rg -n "Trace|trace.*empty|zero.*span" -i -B 2 -A 2 mlflow/store/tracking/sqlalchemy_store.py | head -50Length of output: 1490
🏁 Script executed:
#!/bin/bash
# Search for Trace dataclass or model definition
fd -e py | xargs grep -l "class Trace" | head -5Length of output: 171
🏁 Script executed:
#!/bin/bash
# Look for where the exception about "not fully exported" is used
rg -n "not fully exported" -B 5 -A 5 mlflow/store/tracking/sqlalchemy_store.pyLength of output: 1753
🏁 Script executed:
#!/bin/bash
# Check the actual implementation area around line 3554
sed -n '3540,3580p' mlflow/store/tracking/sqlalchemy_store.pyLength of output: 1938
🏁 Script executed:
#!/bin/bash
# Look for the complete _get_trace definition
ast-grep --pattern 'def _get_trace($_, $_) { $$$ }'Length of output: 0
🏁 Script executed:
#!/bin/bash
# Search for where batch_get_traces calls or handles these methods
sed -n '3569,3610p' mlflow/store/tracking/sqlalchemy_store.pyLength of output: 1649
🏁 Script executed:
#!/bin/bash
# Check if there are any try-except blocks around _get_spans_with_trace_info calls
rg -n "try:" -A 10 mlflow/store/tracking/sqlalchemy_store.py | grep -A 10 "_get_spans_with_trace_info"Length of output: 0
🏁 Script executed:
#!/bin/bash
# Look for test files that test get_trace with edge cases
fd -e py | xargs grep -l "get_trace.*allow_partial" | head -5Length of output: 180
🏁 Script executed:
#!/bin/bash
# Check for any comments or documentation about the retry logic
sed -n '3550,3570p' mlflow/store/tracking/sqlalchemy_store.pyLength of output: 928
🏁 Script executed:
#!/bin/bash
# Look at the broader context of how traces are constructed
grep -n "def.*_get_trace\|def.*_spans_to_trace\|def get_trace\|def batch_get_traces" mlflow/store/tracking/sqlalchemy_store.pyLength of output: 559
🏁 Script executed:
#!/bin/bash
# Check what happens when _get_trace returns None or empty trace
rg -n "if.*_get_trace|return.*_get_trace" -A 2 -B 2 mlflow/store/tracking/sqlalchemy_store.pyLength of output: 500
🏁 Script executed:
#!/bin/bash
# Search for SIZE_STATS definition and how it's populated
rg -n "SIZE_STATS\s*=" -B 5 -A 10 mlflow/store/tracking/sqlalchemy_store.py | head -50Length of output: 0
🏁 Script executed:
#!/bin/bash
# Look for the complete definition of _get_spans_with_trace_info
rg -n "def _get_spans_with_trace_info" -A 30 mlflow/store/tracking/sqlalchemy_store.pyLength of output: 1663
🏁 Script executed:
#!/bin/bash
# Check if there are any handlers or callbacks for exception cases
rg -n "handler|Handler|callback" -i -B 2 -A 2 mlflow/store/tracking/sqlalchemy_store.py | head -50Length of output: 412
🏁 Script executed:
#!/bin/bash
# Look for imports that might hint at exception handling patterns
head -50 mlflow/store/tracking/sqlalchemy_store.pyLength of output: 1161
🏁 Script executed:
#!/bin/bash
# Check MlflowTracingException definition
fd -e py | xargs grep -l "class MlflowTracingException" | head -3Length of output: 20
🏁 Script executed:
#!/bin/bash
# Search for documentation or comments about spans location
rg -n "SPANS_LOCATION.*TRACKING_STORE|non.*tracking.*location" -i -B 3 -A 3 mlflow/store/tracking/sqlalchemy_store.pyLength of output: 537
🏁 Script executed:
#!/bin/bash
# Look for the critical section around lines 3554-3570 with context
sed -n '3545,3600p' mlflow/store/tracking/sqlalchemy_store.pyLength of output: 2501
🏁 Script executed:
#!/bin/bash
# Search for all mentions of "artifact" in context of tracing
rg -n "artifact" -i -B 2 -A 2 mlflow/store/tracking/sqlalchemy_store.py | head -60Length of output: 2764
🏁 Script executed:
#!/bin/bash
# Check the RESOURCE_DOES_NOT_EXIST error code context
rg -n "RESOURCE_DOES_NOT_EXIST" -B 3 -A 3 mlflow/store/tracking/sqlalchemy_store.pyLength of output: 8176
🏁 Script executed:
#!/bin/bash
# Look at what _get_trace actually does - find its full implementation
rg -n "^\s*def _get_trace\(" -A 50 mlflow/store/tracking/sqlalchemy_store.py | head -70Length of output: 2292
🏁 Script executed:
#!/bin/bash
# Check line 3554 to confirm the method signature
sed -n '3554,3555p' mlflow/store/tracking/sqlalchemy_store.pyLength of output: 110
Both trace retrieval concerns verified; implementation matches described behavior.
Verification confirms the two raised issues:
-
Empty spans treated as incomplete: When
allow_partial=Falseand a trace has zero spans, line 3588's conditionif allow_partial or spans:evaluates to false (since[]is falsy), causing_get_traceto returnNone. This triggers the retry loop inget_trace, ultimately raising "not fully exported yet" error even for validly empty traces. TheSIZE_STATScheck only prevents this if metadata explicitly declares 0 expected spans; ifSIZE_STATSis absent, empty traces are always treated as incomplete. -
Unhandled
MlflowTracingExceptionfor non-tracking locations:_get_spans_with_trace_infointentionally raisesMlflowTracingExceptionwhenSPANS_LOCATION != TRACKING_STORE(line 3645), with a comment stating handlers should catch and load from artifact repo. However, neitherget_tracenorbatch_get_tracesincludes exception handling around these calls—the exception propagates directly to callers. Inbatch_get_traces, this means any trace stored outside the tracking store will fail the entire batch operation.
If 0-span traces should be returnable and artifact-repo fallback is implemented externally (e.g., in mlflow/server/handlers.py), these behaviors may be intentional. If not, _get_trace needs to distinguish between empty spans and missing data, and callers may need local fallback handling for non-tracking locations.
🤖 Prompt for AI Agents
In mlflow/store/tracking/sqlalchemy_store.py around lines 3554-3567, change the
logic so get_trace/_get_trace treat an empty list of spans as a valid,
fully-exported trace (i.e., distinguish between None meaning “not
found/incomplete” and [] meaning “found but zero spans”) by checking spans is
not None (or consulting SIZE_STATS==0) instead of truthiness; and add a local
try/except around calls that can raise MlflowTracingException (both in get_trace
and batch_get_traces) to catch MlflowTracingException, log/annotate the trace as
stored outside the tracking store, and allow callers (or the handler layer) to
perform the artifact-repo fallback instead of letting the exception propagate
and fail the operation.
| value = attributes.get(key) | ||
| if isinstance(value, str): | ||
| try: | ||
| result = json.loads(value) | ||
| if isinstance(result, str): | ||
| # the span attributes may be dumped several times in different places | ||
| # (e.g. Span.from_otel_proto, span.to_dict) | ||
| # so we try to load it twice here to get the dumped-once value | ||
| try: | ||
| if json.loads(result): | ||
| return result | ||
| return None | ||
| except json.JSONDecodeError: | ||
| pass | ||
| return value if result else None | ||
| return value if json.loads(value) else None | ||
| except json.JSONDecodeError: | ||
| pass # Use the string value as-is | ||
| return value |
There was a problem hiding this comment.
Consider potential data loss with falsy JSON values.
The expression return value if json.loads(value) else None filters out any JSON value that's falsy in Python, including "false", "0", "[]", "{}", and "null". While the docstring mentions filtering empty dictionaries and null (which seems intentional), filtering valid values like boolean false, numeric 0, or empty arrays [] could lead to data loss if these are legitimate attribute values in OTEL spans.
Additionally, parsing JSON solely to check truthiness (then discarding the result and returning the original string) is inefficient and unclear.
Consider either:
- Documenting the behavior explicitly if filtering all falsy JSON values is intentional:
def _get_and_check_attribute_value(self, attributes: dict[str, Any], key: str) -> Any:
"""
Get attribute value from OTEL attributes by checking whether the value is valid or not.
- This avoids fetching the value if it's empty dictionary or null.
+ This avoids fetching the value if it's a falsy JSON value (null, false, 0, {}, []).
Args:
attributes: Dictionary of span attributes
key: Attribute key
Returns:
Attribute value or None if not found
"""- Filtering only truly empty values if false/0/[] should be preserved:
value = attributes.get(key)
if isinstance(value, str):
try:
parsed = json.loads(value)
# Only filter null and empty containers
if parsed is None or (isinstance(parsed, (dict, list)) and not parsed):
return None
return value
except json.JSONDecodeError:
pass # Use the string value as-is
return value🤖 Prompt for AI Agents
In mlflow/tracing/otel/translation/base.py around lines 154 to 160, the current
code discards any JSON-parsable string that evaluates to falsy (e.g., "false",
"0", "[]", "{}"), which can lose legitimate attribute values; update the logic
to json.loads once into a variable and only return None when the parsed value is
explicitly None or an empty container (empty dict or list), otherwise return the
original string; keep the JSONDecodeError handling to fall back to the raw
string; alternatively, if the original broad falsy filtering was intentional,
add a docstring comment explaining that behavior.
1 participant
Summary by CodeRabbit
New Features
Removals
Documentation
Bug Fixes
Chores