Databricks Machine Learning Associate common mistakes
Databricks Machine Learning Associate · Understand the traps and how to avoid them.
Calls mlflow.start_run() and mlflow.log_metric() without ever calling mlflow.set_experiment(), then cannot locate the runs afterward or re-runs create duplicates in a different place.
When no experiment is set, MLflow writes to the notebook's implicit experiment path, which changes if the notebook is moved, cloned to another user, or run as a job under a different identity. Candidates assume there is a single default experiment for the workspace.
What to do instead
Call mlflow.set_experiment("/Users/[email protected]/my-experiment") (creating the experiment if needed) before any logging, or pass experiment_id explicitly. In jobs, set the experiment path explicitly so runs land in a stable, shared location.
Calls mlflow.register_model("runs:/<run_id>/model", "my_model") against a Unity Catalog metastore, or builds URIs like models:/my_model/Production, and then cannot resolve the model.
The workspace model registry (two-level name, stages) and the Unity Catalog registry (three-level catalog.schema.model, aliases, no stages) look similar but behave differently. The default registry URI is the workspace registry, not UC.
What to do instead
Set mlflow.set_registry_uri("databricks-uc") (or use the databricks-uc registry URI in the client call), register with a fully qualified catalog.schema.model name, and reference versions by alias (models:/catalog.schema.model@champion) rather than by stage.
Calls mlflow.sklearn.log_model(model, "model") and moves straight to registering and serving, without passing a signature or input_example.
Signature inference is often the first thing skipped when the notebook works locally; the training cell runs fine without it, so nothing signals a problem until the endpoint rejects or misinterprets payloads, or column order drifts between training and batch scoring.
What to do instead
Compute mlflow.models.infer_signature(X_train, model.predict(X_train)) and pass signature= and input_example= to log_model. Verify the logged schema in the run's artifacts before registering the version.
Enables autologging and assumes it covers everything (custom business metrics, plots, evaluation datasets), or enables autolog globally and is surprised when a later mlflow.log_metric raises 'no active run'.
Autolog records a fixed set of framework-specific params, metrics, and the model - not arbitrary metrics. It also manages its own run lifecycle when no run is active, so implicit run creation breaks code that assumes an existing run.
What to do instead
Treat autolog as a baseline: wrap training in an explicit mlflow.start_run() and log custom metrics, artifacts, and plots manually. Know which autolog flags (log_models, log_input_examples, log_model_signatures) exist so you can control what is recorded.
Runs Hyperopt fmin() with the default Trials object on a multi-node cluster and is puzzled that tuning is no faster than a single node, or that early stopping parameters are rejected.
The default Trials executes trials sequentially on the driver. Distributed tuning requires SparkTrials, and some features (parallelism control, early stopping via a SparkTrials argument) only exist there. Candidates conflate 'the cluster has many workers' with 'Hyperopt will use them'.
What to do instead
Wrap the objective in fmin(..., algo=tpe.suggest, max_evals=N, trials=SparkTrials(parallelism=K)) when running on a cluster, and size max_evals so the tuning run actually uses the parallelism you pay for.
Reaches for the legacy FeatureStoreClient (databricks.feature_store) to build a training set, or joins feature tables directly with a plain DataFrame join instead of creating a training set with FeatureLookup.
Older tutorials and legacy APIs still appear in search results, and a manual join 'works' on a small test so the candidate never sees the point-in-time / key-alignment problem it creates for production scoring.
What to do instead
Use FeatureEngineeringClient with create_training_set(training_df, feature_lookups=[FeatureLookup(...)], label=...), log the model with fe.log_model so feature metadata travels with it, and reuse the same lookups for batch scoring.
Installs a library with %pip install in a notebook cell and immediately imports it in the next cell, or installs it on the driver only and then expects a Spark UDF/Pandas UDF to find it on executors.
%pip is notebook-scoped and the Python process must be restarted to pick up the new package; executor-side availability depends on cluster libraries, not the notebook environment. Both fail in ways that look like unrelated ImportErrors.
What to do instead
After %pip install, call dbutils.library.restartPython() (or restart the cluster for cluster-scoped installs), and install anything used inside UDFs as a cluster library so it is present on every node.
Trains, registers, or serves a Unity Catalog model or feature table, then hands it to another user or a serving endpoint without checking USE CATALOG / USE SCHEMA / EXECUTE privileges.
Candidates test only with their own account, which typically owns the objects or has metastore admin rights, so permission gaps only surface later as opaque endpoint-creation or feature-lookup failures.
What to do instead
For any shared model or feature table, explicitly grant USE CATALOG on the catalog, USE SCHEMA on the schema, and EXECUTE (models) or SELECT (feature tables) to the consuming users or service principal, and verify by testing as that principal.
Registers a new model version and assumes an existing Model Serving endpoint will start serving it, or assumes a scale-to-zero endpoint has no latency trade-off.
Registration and deployment are separate steps, and scale-to-zero is chosen for cost without accounting for cold-start on the first request after idle. The candidate's mental model merges 'the model exists' with 'the endpoint is updated'.
What to do instead
Update the served entity/endpoint configuration to the new version or alias as an explicit deployment step (and automate it), and enable/verify an inference table so requests and responses are captured for monitoring. Size the workload against scale-to-zero cold start.
Scores a large table by collecting it to the driver with .toPandas() and looping model.predict(), or applies mlflow.pyfunc.spark_udf without the model's logged signature/environment in place.
The pandas path is the pattern from local data-science work and succeeds on a sample, so scaling and schema problems only appear on the full table - or the UDF silently returns wrong-typed predictions because no signature was recorded.
What to do instead
Load the model URI, create predict_udf = mlflow.pyfunc.spark_udf(spark, model_uri, result_type=...), and apply it with df.withColumn("prediction", predict_udf(*feature_cols)), keeping computation distributed and the schema enforced by the logged signature.