https://community.databricks.com/t5/technical-blog/train-and-deploy-yolo-vision-model-on-databricks-ai-runtime-air/ba-p/151558 <H2><STRONG><A id="intro" target="_blank"></A>Introduction</STRONG></H2> <P><A href="https://docs.ultralytics.com/" target="_blank" rel="noopener"><SPAN>Ultralytics YOLO</SPAN></A><SPAN><A href="#ultralytics-licensing" target="_self"><SUP>*</SUP></A> (You Only Look Once) is one of the most widely used </SPAN><STRONG>computer vision frameworks</STRONG><SPAN>. It is fast, accurate, and well supported, with a range of model sizes (from nano to extra-large) so you can trade off speed and accuracy for edge or server deployment. Training and inference are straightforward with a Python API and practical documentation, and the ecosystem features readily available pretrained weights, support for standard datasets (e.g. </SPAN><A href="https://cocodataset.org/" target="_blank" rel="noopener"><SPAN>COCO</SPAN></A><SPAN>), and ongoing active model development, as exemplified by recent advancements in </SPAN><A href="https://docs.ultralytics.com/models/yolo11/" target="_blank" rel="noopener"><SPAN>YOLO11</SPAN></A><SPAN>.</SPAN></P> <P><SPAN>For teams adopting computer vision (CV) tasks on Databricks, </SPAN><A href="https://docs.ultralytics.com/" target="_blank" rel="noopener"><SPAN>Ultralytics YOLO</SPAN></A><SPAN> is a practical choice for both prototyping and production pipelines. The framework supports multiple CV tasks — </SPAN><A href="https://docs.ultralytics.com/tasks/detect/" target="_blank" rel="noopener"><STRONG>object detection</STRONG></A><SPAN>, </SPAN><A href="https://docs.ultralytics.com/tasks/classify/" target="_blank" rel="noopener"><STRONG>classification</STRONG></A><SPAN>, </SPAN><A href="https://docs.ultralytics.com/tasks/segment/" target="_blank" rel="noopener"><STRONG>segmentation</STRONG></A><SPAN>, </SPAN><A href="https://docs.ultralytics.com/tasks/pose/" target="_blank" rel="noopener"><STRONG>pose estimation</STRONG></A><SPAN>, and </SPAN><A href="https://docs.ultralytics.com/tasks/obb/" target="_blank" rel="noopener"><STRONG>oriented bounding boxes (OBB)</STRONG></A><SPAN> — each with models in several sizes (nano to extra-large, often denoted as </SPAN><SPAN>n</SPAN><SPAN>, </SPAN><SPAN>s</SPAN><SPAN>, </SPAN><SPAN>m</SPAN><SPAN>, </SPAN><SPAN>l</SPAN><SPAN>, </SPAN><SPAN>x</SPAN><SPAN>).</SPAN></P> <P class="lia-align-center"><span class="lia-inline-image-display-wrapper lia-image-align-inline" image-alt="mmt_0-1781055272540.png" style="width: 999px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/27680i4FFA6D3D7A180B94/image-size/large?v=v2&amp;px=999" role="button" title="mmt_0-1781055272540.png" alt="mmt_0-1781055272540.png" /></span></P> <P class="lia-align-center"><I><SPAN>Figure 1: Common computer vision tasks and their associated annotation type.</SPAN></I></P> <P>&nbsp;</P> <P><SPAN>This post demonstrates a single-node workflow for training an </SPAN><STRONG>object detection</STRONG><SPAN> model on </SPAN><A href="https://www.databricks.com/blog/introducing-ai-runtime-scalable-serverless-nvidia-gpus-databricks-training-and-finetuning" target="_blank" rel="noopener"><SPAN>Databricks AI Runtime</SPAN></A><SPAN> — scalable, serverless NVIDIA GPU compute. We use the nano YOLO model, </SPAN><A href="https://docs.ultralytics.com/models/yolo11/" target="_blank" rel="noopener"><SPAN>YOLO11n</SPAN></A><SPAN>, for real-time performance that outputs bounding boxes, class labels, and confidence scores. The process covers training YOLO11n on the </SPAN><A href="https://docs.ultralytics.com/datasets/detect/coco128/" target="_blank" rel="noopener"><SPAN>COCO128</SPAN></A><SPAN> dataset (</SPAN><I><SPAN>demo-only</SPAN></I><SPAN>; refer to </SPAN><STRONG>Data preparation</STRONG><SPAN> for production guidance) and deploying it to </SPAN><A href="https://docs.databricks.com/en/machine-learning/model-serving/index.html" target="_blank" rel="noopener"><SPAN>Model Serving</SPAN></A><SPAN>. Deployment includes a </SPAN><A href="https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#creating-custom-pyfunc-models" target="_blank" rel="noopener"><SPAN>custom MLflow PyFunc wrapper</SPAN></A><SPAN> to handle base64 image input to the YOLO model and structured bounding-box output from model prediction.</SPAN></P> <P>&nbsp;</P> <HR /> <H2><STRONG>Why AI Runtime and single-node?</STRONG></H2> <P><SPAN>Critically, running YOLO on </SPAN><A href="https://docs.databricks.com/aws/en/machine-learning/ai-runtime" target="_blank" rel="noopener"><SPAN>Databricks AI Runtime</SPAN></A><SPAN> lets you train and iterate </SPAN><STRONG>without provisioning or managing clusters</STRONG><SPAN>: you get GPU compute on demand, pay for what you use, and when you are done the compute is terminated. This makes it ideal for experimentation, proof-of-concept, and small-to-medium training jobs — and </SPAN><A href="https://docs.databricks.com/en/mlflow/index.html" target="_blank" rel="noopener"><SPAN>MLflow</SPAN></A><SPAN> and </SPAN><A href="https://docs.databricks.com/en/connect/unity-catalog/volumes.html" target="_blank" rel="noopener"><SPAN>Unity Catalog</SPAN></A><SPAN> keep experiments and artifacts organized.</SPAN></P> <P><STRONG>Single-node</STRONG><SPAN> (one GPU instance) keeps the workflow simple and sufficient for many object-detection use cases. YOLO11n is a small model; training on datasets in the low thousands to tens of thousands of images often fits comfortably on one GPU (e.g. A10). A single node avoids distributed-training setup, multi-worker debugging, and extra cost — so you can focus on data, labels, and the MLflow-to-serving path.</SPAN></P> <P><SPAN>When your dataset or model grows and training time becomes a bottleneck, you can move to multi-GPU or multi-node patterns; the same registration and deployment steps in this post still apply.</SPAN></P> <P>&nbsp;</P> <HR /> <H2><STRONG>What you'll need to get started</STRONG></H2> <UL> <LI style="font-weight: 400;" aria-level="1"><STRONG>Workspace:</STRONG><SPAN> AI Runtime with access to an </SPAN><STRONG>AI</STRONG><SPAN> base environment (Serverless GPU).</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><STRONG>Permissions:</STRONG><SPAN> Create schemas and volumes in Unity Catalog; run MLflow experiments; register models; create Model Serving endpoints.</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><STRONG>Notebook:</STRONG><SPAN> From the </SPAN><A href="https://github.com/databricks-industry-solutions/cv-playground" target="_blank" rel="noopener"><SPAN>cv-playground</SPAN></A><SPAN> repository — e.g. </SPAN><A href="https://github.com/databricks-industry-solutions/cv-playground/blob/main/projects/ultralytics_databricks_examples/air-yolo11n-detect-coco128-singleGPU.ipynb" target="_blank" rel="noopener"><SPAN>air-yolo11n-detect-coco128-singleGPU.ipynb</SPAN></A><SPAN>.</SPAN></LI> </UL> <P><SPAN>Add the notebook via </SPAN><A href="https://docs.databricks.com/en/notebooks/notebook-export-import.html" target="_blank" rel="noopener"><SPAN>Import a notebook</SPAN></A><SPAN> or </SPAN><A href="https://docs.databricks.com/en/repos/repos-setup.html" target="_blank" rel="noopener"><SPAN>Clone a Git repo (Repos)</SPAN></A><SPAN>, then attach to Serverless GPU.</SPAN></P> <P><SPAN>Navigate to the </SPAN><STRONG>Compute</STRONG><SPAN> dropdown, </SPAN><STRONG>Connect</STRONG><SPAN> to and configure the notebook AI Runtime. Within the </SPAN><STRONG>Environment Panel</STRONG><SPAN> found on the </SPAN><SPAN>right-hand edge</SPAN><SPAN> of the notebook, </SPAN><SPAN>select an </SPAN><STRONG>Accelerator</STRONG><SPAN> and the </SPAN><STRONG>AI Base environment</STRONG><SPAN> (this walkthrough uses </SPAN><STRONG>A10</STRONG><SPAN> on </SPAN><STRONG>AI v5</STRONG><SPAN>, as shown below)</SPAN><SPAN>. Finally, click </SPAN><STRONG>Apply</STRONG><SPAN> and then </SPAN><STRONG>Confirm</STRONG><SPAN>, as shown in the figure below.</SPAN></P> <P><span class="lia-inline-image-display-wrapper lia-image-align-center" image-alt="env_config_setup_v5.png" style="width: 999px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/27687i4E18D75D08FA156E/image-size/large?v=v2&amp;px=999" role="button" title="env_config_setup_v5.png" alt="env_config_setup_v5.png" /></span></P> <P class="lia-align-center"><I><SPAN>Figure 2: Connecting to AI Runtime Serverless GPU cluster and configuring the Notebook Environment:&nbsp;</SPAN></I><I><SPAN>Connect → Serverless GPU → Environment → </SPAN></I><I><SPAN>choose your </SPAN></I><STRONG><I>Accelerator</I></STRONG><I><SPAN> + the </SPAN></I><STRONG><I>AI</I></STRONG><I><SPAN> base environment</SPAN></I><I><SPAN> → Apply &amp; Confirm.</SPAN></I></P> <P>&nbsp;</P> <P><SPAN>Given that packages and dependencies are installed in the first notebook cell, there is no need to install within the cluster environment panel.</SPAN></P> <P>&nbsp;</P> <HR /> <H2><STRONG>Workflow overview</STRONG></H2> <P><SPAN>The notebook walks through six steps in order; each builds on the previous one.</SPAN></P> <P><span class="lia-inline-image-display-wrapper lia-image-align-center" image-alt="mmt_2-1781055272547.png" style="width: 999px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/27681i4B60E65E631DC9F3/image-size/large?v=v2&amp;px=999" role="button" title="mmt_2-1781055272547.png" alt="mmt_2-1781055272547.png" /></span></P> <P class="lia-align-center"><I><SPAN>Figure 3: An end-to-end workflow.</SPAN></I></P> <P>&nbsp;</P> <H3><STRONG>Step 1: Setup — Environment and Unity Catalog</STRONG></H3> <P><SPAN>After attaching to Serverless GPU (see above: </SPAN><I><SPAN>Connect → Serverless GPU → choose your Accelerator + AI base environment → Apply and Confirm</SPAN></I><SPAN>), the first steps are to install the required Python packages and configure your Unity Catalog project structure.</SPAN></P> <P><SPAN>On AI v5, the base environment already bundles</SPAN> <SPAN>mlflow&gt;=3</SPAN><SPAN> (skinny build), </SPAN><SPAN>nvidia-ml-py</SPAN><SPAN>, </SPAN><SPAN>threadpoolctl</SPAN><SPAN>, and </SPAN><SPAN>torch</SPAN><SPAN> — so the only package you install is </SPAN><A href="https://docs.ultralytics.com/" target="_blank" rel="noopener"><SPAN>Ultralytics</SPAN></A><SPAN>. </SPAN><SPAN>%pip</SPAN><SPAN> restarts Python automatically after the install. Then set a writable YOLO config directory to avoid permission issues. A </SPAN><STRONG>Package Verification</STRONG><SPAN> cell (in the notebook) confirms the expected packages are present and flags the uncommon case where Model Serving needs the full </SPAN><SPAN>mlflow</SPAN><SPAN> re-added on top of the skinny build.</SPAN></P> <H4>&nbsp;</H4> <LI-CODE lang="python"># ============================================================ # PACKAGE INSTALLATION — AI v5 (Serverless GPU, single GPU) # ============================================================ # AI v5 pre-bundles mlflow&gt;=3 (skinny), nvidia-ml-py, threadpoolctl, torch. # Only ultralytics needs installing. %pip install ultralytics==8.3.204 -q # Note: %pip automatically restarts the Python environment after install. # Set a writable YOLO config dir (avoids permission errors) import os, uuid config_dir = f'/tmp/yolo_config_{uuid.uuid4().hex[:8]}' os.environ['YOLO_CONFIG_DIR'] = config_dir os.makedirs(config_dir, exist_ok=True) </LI-CODE> <H4>&nbsp;</H4> <P><SPAN>Next, create or use a catalog, schema, and </SPAN><A href="https://docs.databricks.com/en/connect/unity-catalog/volumes.html" target="_blank" rel="noopener"><SPAN>Unity Catalog Volume</SPAN></A><SPAN> for data, raw models, and model checkpoints from training runs. Use widgets for catalog, schema, volume, and model name so the same notebook can be reused across workspaces.</SPAN></P> <H4>&nbsp;</H4> <LI-CODE lang="python"># Widgets for catalog, schema, volume, model name catalog_name = dbutils.widgets.get("catalog_name") # e.g. "main" schema_name = dbutils.widgets.get("schema_name") # e.g. "default" volume_name = dbutils.widgets.get("volume_name") # e.g. "yolo_sgc" spark.sql(f"CREATE SCHEMA IF NOT EXISTS `{catalog_name}`.`{schema_name}`") spark.sql(f"CREATE VOLUME IF NOT EXISTS `{catalog_name}`.`{schema_name}`.`{volume_name}`") project_location = f'/Volumes/{catalog_name}/{schema_name}/{volume_name}/' os.makedirs(f'{project_location}runs/', exist_ok=True) os.makedirs(f'{project_location}data/', exist_ok=True) os.makedirs(f'{project_location}raw_model/', exist_ok=True) </LI-CODE> <H4>&nbsp;</H4> <H3><STRONG>Step 2: Data preparation</STRONG></H3> <P><SPAN>The dataset is configured via a YAML file (path, splits, class names). We download the COCO128 config and data to the volume, then split it into train (62.5%), validation (18.75%), and test (18.75%) with a fixed seed, updating the YAML with the new paths. For custom data, you typically adjust the YAML for your paths and classes. We use the Ultralytics </SPAN><SPAN>coco128.yaml</SPAN><SPAN>, downloaded to the UC Volume, but you can substitute your own config (e.g., </SPAN><SPAN>data.yaml</SPAN><SPAN>).</SPAN></P> <H4>&nbsp;</H4> <LI-CODE lang="python"># Download COCO128 dataset configuration to UC Volume import yaml os.makedirs(f'{project_location}data/coco128', exist_ok=True) config_url = "https://github.com/ultralytics/ultralytics/raw/main/ultralytics/cfg/datasets/coco128.yaml" config_path = f"{project_location}data/coco128.yaml" download_file(config_url, config_path, "COCO128 config") # Then load config, set data['path'] to volume path, download/extract dataset if needed, save updated YAML Split the data and update the YAML with train/val/test image paths: train_size, val_size, test_size = split_dataset( source_images_dir=f"{project_location}data/coco128/images/train2017", source_labels_dir=f"{project_location}data/coco128/labels/train2017", base_images_dir=f"{project_location}data/coco128/images", base_labels_dir=f"{project_location}data/coco128/labels", train_ratio=0.625, # 62.5% val_ratio=0.1875, # 18.75% random_seed=42, ) # In the notebook: update data.yaml so 'train', 'val', 'test' point to the new split dirs # (e.g. .../images/train, .../images/val, .../images/test) </LI-CODE> <H4>&nbsp;</H4> <DIV style="border-left: 4px solid #d0d0d0; padding-left: 18px; margin: 16px 0;"> <P style="margin: 0; font-size: 18px; line-height: 1.5;"><EM><FONT size="3" color="#808080"><STRONG>Important note:</STRONG> COCO128 is used here only for demonstration. With ~128 images it is too small for production and will overfit. For real use cases, use larger datasets (e.g. 100K+ images or 1K+ domain-specific images). The same workflow applies — update data paths and config as needed.</FONT></EM></P> </DIV> <H3>&nbsp;</H3> <H3><STRONG>Step 3: MLflow — Custom PyFunc Wrapper and Configuration</STRONG></H3> <P><SPAN>To deploy the trained YOLO model to </SPAN><A href="https://docs.databricks.com/en/machine-learning/model-serving/index.html" target="_blank" rel="noopener"><SPAN>Model Serving</SPAN></A><SPAN>, we need a single, serializable API: the endpoint will receive requests (e.g., base64-encoded images) and return structured responses (e.g., bounding boxes). YOLO's native API expects file paths or NumPy arrays and returns a rich in-memory object, which is not what the serving layer expects.</SPAN></P> <P><SPAN>The notebook therefore defines an </SPAN><A href="https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#creating-custom-pyfunc-models" target="_blank" rel="noopener"><SPAN>MLflow Custom PyFunc wrapper</SPAN></A> <SPAN>YOLOWrapper(mlflow.pyfunc.PythonModel)</SPAN><SPAN> that accepts a DataFrame with an </SPAN><SPAN>image_base64</SPAN><SPAN> column and returns a DataFrame of detections (class, confidence, bbox columns).</SPAN></P> <P><SPAN>The wrapper class has three methods:</SPAN></P> <UL> <LI style="font-weight: 400;" aria-level="1"><STRONG>load_context</STRONG><SPAN>&nbsp;loads the </SPAN><SPAN>.pt</SPAN><SPAN> artifact into </SPAN><SPAN>self.model</SPAN><SPAN> when the model is loaded (e.g. at serving startup).</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><STRONG>predict</STRONG><SPAN>&nbsp;accepts a DataFrame with </SPAN><SPAN>image_base64</SPAN><SPAN>, decodes each image, runs YOLO, and returns a DataFrame via </SPAN><SPAN>_format_predictions</SPAN><SPAN>.</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><STRONG>_format_predictions</STRONG><SPAN>&nbsp;converts YOLO's </SPAN><SPAN>Results</SPAN><SPAN> (</SPAN><SPAN>.boxes</SPAN><SPAN>, </SPAN><SPAN>.names</SPAN><SPAN>) into a single DataFrame with class name, class id, confidence, and bbox columns (</SPAN><SPAN>xyxy</SPAN><SPAN> and </SPAN><SPAN>xywh</SPAN><SPAN>).</SPAN></LI> </UL> <P><SPAN>We define the wrapper now so it's ready to use immediately after training completes in </SPAN><STRONG>Step 4</STRONG><SPAN>.</SPAN></P> <H4>&nbsp;</H4> <LI-CODE lang="python">class YOLOWrapper(mlflow.pyfunc.PythonModel): """Custom MLflow wrapper for YOLO models using base64-encoded images.""" def load_context(self, context): """Load YOLO model from artifacts (called once when model is loaded).""" from ultralytics import YOLO model_path = context.artifacts["yolo_model"] self.model = YOLO(model_path, task='detect') def _format_predictions(self, predictions): """Convert YOLO Results to a single DataFrame with class, confidence, bbox columns.""" import pandas as pd all_results = [] for prediction in predictions: if prediction.boxes is not None: boxes = prediction.boxes for i in range(len(boxes)): box_xyxy = boxes.xyxy[i].cpu().numpy() box_xywh = boxes.xywh[i].cpu().numpy() all_results.append({ "class_name": prediction.names[int(boxes.cls[i])], "class_num": int(boxes.cls[i]), "confidence": float(boxes.conf[i]), "bbox_x1": float(box_xyxy[0]), "bbox_y1": float(box_xyxy[1]), "bbox_x2": float(box_xyxy[2]), "bbox_y2": float(box_xyxy[3]), "bbox_center_x": float(box_xywh[0]), "bbox_center_y": float(box_xywh[1]), "bbox_width": float(box_xywh[2]), "bbox_height": float(box_xywh[3]), }) return pd.DataFrame(all_results) def predict(self, context, model_input): """Accept DataFrame with image_base64; decode, run YOLO, return DataFrame of detections.""" import pandas as pd import base64 from PIL import Image import io import numpy as np if 'image_base64' not in model_input.columns: raise ValueError("DataFrame must contain 'image_base64' column") all_predictions = [] for image_base64 in model_input['image_base64'].tolist(): image_bytes = base64.b64decode(image_base64) image_array = np.array(Image.open(io.BytesIO(image_bytes))) predictions = self.model.predict(image_array, verbose=False) all_predictions.extend(predictions) return self._format_predictions(all_predictions)</LI-CODE> <H4>&nbsp;</H4> <H4><STRONG>Step 3.1: MLflow configuration</STRONG></H4> <P><SPAN>We </SPAN><A href="https://mlflow.org/docs/latest/api_reference/_modules/mlflow/models/signature.html#infer_signature" target="_blank" rel="noopener"><SPAN>infer the model signature</SPAN></A><SPAN> from a sample prediction using this custom wrapper (DataFrame with </SPAN><SPAN>image_base64</SPAN><SPAN> input and detection columns output). We also set the MLflow </SPAN><A href="https://mlflow.org/docs/latest/api_reference/_modules/mlflow/entities/experiment_tag.html#ExperimentTag" target="_blank" rel="noopener"><SPAN>experiment</SPAN></A><SPAN> (e.g., under </SPAN><SPAN>/Workspace/Shared/</SPAN><SPAN>) and enable system metrics logging, along with YOLO's MLflow integration and MLflow autologging.</SPAN></P> <H4>&nbsp;</H4> <LI-CODE lang="python"># Infer signature from a sample image (input: base64, output: bbox columns) signature, input_example = infer_model_signature(model_path, sample_images[0]) # Enable system metrics and set experiment experiment_name, experiment_id = setup_mlflow_experiment( use_workspaceUsers_path=False, expt_name_suffix="Experiments_YOLO_CoCo", ) </LI-CODE> <H4>&nbsp;</H4> <P><SPAN>The model is then registered to Unity Catalog (</SPAN><SPAN>mlflow.pyfunc.log_model</SPAN><SPAN>) using this wrapper and the best </SPAN><SPAN>checkpoint.pt</SPAN><SPAN> artifact (called after training in </SPAN><STRONG>Step 4:</STRONG></P> <H4>&nbsp;</H4> <LI-CODE lang="python">mlflow.pyfunc.log_model( name="model", python_model=YOLOWrapper(), artifacts={"yolo_model": model_path}, signature=signature, input_example=input_example, registered_model_name=registered_model_name, pip_requirements=["ultralytics==...", "cloudpickle==...", "torch", "torchvision", "pillow", "numpy"], ) </LI-CODE> <H4>&nbsp;</H4> <H3><STRONG>Step 4: Model training</STRONG></H3> <P><SPAN>Train YOLO11n with your chosen hyperparameters (</SPAN><SPAN>epochs</SPAN><SPAN>, </SPAN><SPAN>batch size</SPAN><SPAN>, </SPAN><SPAN>learning rate</SPAN><SPAN>, </SPAN><SPAN>patience</SPAN><SPAN>, </SPAN><SPAN>dropout</SPAN><SPAN>, </SPAN><SPAN>weight decay</SPAN><SPAN>); these are specified in the </SPAN><A href="https://docs.ultralytics.com/usage/cfg/" target="_blank" rel="noopener"><SPAN>config</SPAN></A><SPAN> variables within </SPAN><SPAN>model.train()</SPAN><SPAN> as shown in the code snippet. Training runs in a unique temp directory, and the results and validation metrics are copied into the volume under a named run folder (</SPAN><SPAN>{task}_{model}_{dataset}_{timestamp}_run_{run_id}</SPAN><SPAN>). The best </SPAN><SPAN>checkpoint</SPAN><SPAN> is saved and is then registered to Unity Catalog with the custom PyFunc wrapper (base64 in, structured detections out) defined in the previous step.</SPAN></P> <H4>&nbsp;</H4> <LI-CODE lang="python">model = YOLO(model_path) results = model.train( task="detect", batch=4, device=0, # Single GPU for Serverless AI Runtime data=data_yaml_path, epochs=100, lr0=0.001, project=project_location, name=f"run_{timestamp}", patience=5, # Adjust as needed dropout=0.2, weight_decay=0.0005, save=True, ) run_id = mlflow.last_active_run().info.run_id # Register to Unity Catalog with custom PyFunc wrapper (base64 in, bbox out) registered_model_name = register_yolo_model( run_id=run_id, model_path=best_model_path, catalog_name=catalog_name, schema_name=schema_name, model_name=model_name, signature=signature, input_example=input_example, data_yaml_path=data_yaml_path, ) </LI-CODE> <H4>&nbsp;</H4> <H3><STRONG>Step 5: Model evaluation</STRONG></H3> <P><SPAN>Evaluate the registered model on </SPAN><SPAN>validation</SPAN><SPAN> and </SPAN><SPAN>test</SPAN><SPAN> sets (sample </SPAN><SPAN>predictions</SPAN><SPAN> and </SPAN><SPAN>metrics</SPAN><SPAN>), then run a </SPAN><STRONG>local serving test</STRONG><SPAN> by loading the </SPAN><SPAN>model</SPAN><SPAN> via </SPAN><SPAN>mlflow.pyfunc.load_model()</SPAN><SPAN> and calling it with </SPAN><SPAN>base64</SPAN><SPAN>-encoded images to confirm the same interface the endpoint will use.</SPAN></P> <H4>&nbsp;</H4> <LI-CODE lang="python"># Local serving test: same I/O as the deployed endpoint model_uri = f"models:/{registered_model_name}/{latest_version}" serving_model = mlflow.pyfunc.load_model(model_uri) with open(test_image_path, 'rb') as f: image_base64 = base64.b64encode(f.read()).decode('utf-8') input_df = pd.DataFrame({"image_base64": [image_base64]}) predictions = serving_model.predict(input_df) # DataFrame with class_name, confidence, bbox_* </LI-CODE> <H4>&nbsp;</H4> <H3><STRONG>Step 6: Model deployment</STRONG></H3> <P><SPAN>After a manual checkpoint (e.g. a </SPAN><I><SPAN>"Proceed with Deployment"</SPAN></I><SPAN> widget), you can create or update a </SPAN><A href="https://docs.databricks.com/aws/en/machine-learning/model-serving/custom-models" target="_blank" rel="noopener"><SPAN>Custom Model Serving</SPAN></A><SPAN> endpoint. The deployment configuration includes:</SPAN></P> <UL> <LI style="font-weight: 400;" aria-level="1"><A href="https://docs.databricks.com/en/dev-tools/sdk-python.html" target="_blank" rel="noopener"><STRONG>WorkspaceClient SDK</STRONG></A><SPAN> — enables programmatic endpoint management, ensuring deployments are repeatable, version-controlled, and integrated into automated workflows.</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><STRONG><A href="https://docs.databricks.com/aws/en/machine-learning/model-serving/create-manage-serving-endpoints#workload-types" target="_blank" rel="noopener">Small [endpoint workload compute size]</A></STRONG><SPAN>&nbsp;and enabling scale-to-zero — minimizes compute costs during development and evaluation by provisioning resources on demand and releasing them when the endpoint is idle.</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><A href="https://docs.databricks.com/en/ai-gateway/inference-tables/" target="_blank" rel="noopener">Unity AI Gateway inference tables</A><SPAN> — automatically logs all request and response payloads to a Unity Catalog Delta table, providing a built-in audit trail for monitoring, debugging, and downstream evaluation without additional instrumentation.</SPAN></LI> </UL> <H4>&nbsp;</H4> <H4><STRONG>Step 6.1: Create the endpoint with Unity AI Gateway enabled</STRONG></H4> <H4>&nbsp;</H4> <LI-CODE lang="python">from databricks.sdk import WorkspaceClient from databricks.sdk.service.serving import ( ServedEntityInput, EndpointCoreConfigInput, AiGatewayConfig, AiGatewayInferenceTableConfig, ) w = WorkspaceClient() w.serving_endpoints.create( name=endpoint_name, config=EndpointCoreConfigInput( served_entities=[ ServedEntityInput( entity_name=registered_model_name, entity_version=str(model_version), workload_size="Small", scale_to_zero_enabled=True, ) ] ), ai_gateway=AiGatewayConfig( inference_table_config=AiGatewayInferenceTableConfig( catalog_name=catalog_name, schema_name=schema_name, table_name_prefix=endpoint_name, enabled=True, ) ), ) </LI-CODE> <H4>&nbsp;</H4> <P><SPAN>Test the deployed endpoint by calling it with a </SPAN><SPAN>base64</SPAN><SPAN>-encoded image and verify the structured </SPAN><SPAN>bounding-box</SPAN><SPAN> response.</SPAN></P> <H4>&nbsp;</H4> <H4><STRONG>Step 6.2: Call the endpoint with base64 input (same as local PyFunc test)</STRONG></H4> <H4>&nbsp;</H4> <LI-CODE lang="python">import base64 with open(test_image_path, 'rb') as f: image_base64 = base64.b64encode(f.read()).decode('utf-8') response = w.serving_endpoints.query( name=endpoint_name, dataframe_records=[{"image_base64": image_base64}], ) # Response contains DataFrame with class_name, confidence, bbox_x1, bbox_y1, ... </LI-CODE> <H4>&nbsp;</H4> <P><SPAN>Successful testing confirms the custom PyFunc wrapper's ability to handle base64-encoded image input and return a structured bounding-box output from the Model Serving endpoint, which are critical technical considerations detailed next.</SPAN></P> <P>&nbsp;</P> <HR /> <H2><STRONG>Key technical details</STRONG></H2> <P><SPAN>A few details are worth calling out for implementation and operations:</SPAN></P> <UL> <LI style="font-weight: 400;" aria-level="1"><STRONG>Base64 input:</STRONG><SPAN> The custom </SPAN><A href="https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#creating-custom-pyfunc-models" target="_blank" rel="noopener"><SPAN>MLflow PyFunc</SPAN></A><SPAN> wrapper accepts a DataFrame with an </SPAN><SPAN>image_base64</SPAN><SPAN> column, where images are encoded as </SPAN><SPAN>base64</SPAN><SPAN> strings. </SPAN><I><SPAN>In our example, JPEG images are used; other formats (e.g. PNG, BMP) may work but have not been validated here.</SPAN></I><SPAN> Base64 encoding keeps the API simple and works across network boundaries for Model Serving.</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><STRONG>Bounding box output:</STRONG><SPAN> The wrapper returns a DataFrame with columns such as </SPAN><SPAN>class_name</SPAN><SPAN>, </SPAN><SPAN>class_num</SPAN><SPAN>, </SPAN><SPAN>confidence</SPAN><SPAN>, </SPAN><SPAN>bbox_x1</SPAN><SPAN>, </SPAN><SPAN>bbox_y1</SPAN><SPAN>, </SPAN><SPAN>bbox_x2</SPAN><SPAN>, </SPAN><SPAN>bbox_y2</SPAN><SPAN>, and center/width/height. This structure is inferred once and used for the registered model's </SPAN><A href="https://mlflow.org/docs/latest/ml/model/signatures/" target="_blank" rel="noopener"><SPAN>signature</SPAN></A><SPAN>.</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><STRONG>Why we use a custom wrapper:</STRONG><SPAN> YOLO's native API expects image paths or NumPy arrays and returns a rich in-memory object; </SPAN><A href="https://docs.databricks.com/en/machine-learning/model-serving/index.html" target="_blank" rel="noopener"><SPAN>Model Serving</SPAN></A><SPAN> expects a single, serializable contract (DataFrame input and output) for HTTP requests and responses. The wrapper (1) accepts base64-encoded images in a DataFrame column (JSON-friendly), (2) loads the </SPAN><SPAN>.pt</SPAN><SPAN> artifact and runs the YOLO object detection task inside </SPAN><SPAN>predict</SPAN><SPAN>, and (3) returns a structured DataFrame of detections that the endpoint can serialize. Without it, the raw YOLO model could not be deployed as a standard PyFunc. The full implementation is in </SPAN><STRONG>Step 3</STRONG><SPAN> above.</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><STRONG>Unity Catalog Volume structured layout:</STRONG><SPAN> Data is stored in </SPAN><SPAN>/Volumes/{catalog}/{schema}/{volume}/data/</SPAN><SPAN>. Pretrained weights are located in </SPAN><SPAN>raw_model/</SPAN><SPAN>. Each training run has its own dedicated folder under </SPAN><SPAN>runs/{task}_{model}_{dataset}_{timestamp}_run_{run_id}/</SPAN><SPAN>, which includes subfolders for </SPAN><SPAN>train/</SPAN><SPAN>, </SPAN><SPAN>validation_metrics/</SPAN><SPAN>, </SPAN><SPAN>validation_samples/</SPAN><SPAN>, and </SPAN><SPAN>test_samples/</SPAN><SPAN>.</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><STRONG>Deployment safety:</STRONG><SPAN> A parameter widget (e.g. </SPAN><I><SPAN>"Proceed with Deployment"</SPAN></I><SPAN>) gates the deployment cells so </SPAN><I><SPAN>"Run All"</SPAN></I><SPAN> doesn't deploy by accident. Endpoint creation/update can take on the order of 10–20 minutes; the notebook can exit early and direct you to re-run or check the UI.</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><SPAN> <A href="https://docs.databricks.com/aws/en/ai-gateway/overview-serving-endpoints" target="_blank" rel="noopener">Unity AI Gateway</A><STRONG>:</STRONG> New endpoints are created with Unity AI Gateway inference table config (catalog, schema, table name prefix). The payload table is created and populated after the first requests; there can be a short delay before rows appear. You can query the table (e.g. </SPAN><SPAN>SELECT * FROM catalog.schema.endpoint_payload ORDER BY timestamp_ms DESC</SPAN><SPAN>) to inspect logged requests.</SPAN></LI> </UL> <P>&nbsp;</P> <HR /> <H2><STRONG>Conclusion</STRONG></H2> <P><SPAN>The ability to train or fine-tune and deploy YOLO (You Only Look Once) models on Databricks Data Intelligence Platform provides enterprises with a high-performance, cost-optimized, and easily adoptable Computer Vision (CV) solution. Our walkthrough shows a complete path from raw images to a live YOLO endpoint on Databricks — no cluster provisioning, full MLflow tracking, Unity Catalog governance, and production-ready serving with built-in request logging. Swap COCO128 for your own dataset and the same workflow applies. As your data or model complexity grows, the same registration and deployment steps extend to multi-GPU and multi-node training patterns.</SPAN></P> <H4>&nbsp;</H4> <P class="lia-align-center"><span class="lia-inline-image-display-wrapper lia-image-align-inline" image-alt="mmt_3-1781055272554.png" style="width: 999px;"><img src="https://community.databricks.com/t5/image/serverpage/image-id/27683iFE126149736F5CEA/image-size/large?v=v2&amp;px=999" role="button" title="mmt_3-1781055272554.png" alt="mmt_3-1781055272554.png" /></span></P> <P class="lia-align-center"><I><SPAN>Figure 4: Example validation of YOLO object detection inference on a sample of COCO128 images.</SPAN></I></P> <P>&nbsp;</P> <DIV style="border-left: 4px solid #d0d0d0; padding-left: 18px; margin: 16px 0;"> <P style="margin: 0; font-size: 18px; line-height: 1.5;"><FONT color="#808080"><EM><FONT size="3"><STRONG>Single-GPU accelerator options:</STRONG> This walkthrough uses a single A10 GPU, a cost-efficient choice for YOLO11n on small-to-medium datasets. The same notebook code also runs on a single H100 (Beta) for faster training — just select H100 as the Accelerator in the same AI v5 Environment panel; the rest of the workflow is unchanged. Check the <A href="https://docs.databricks.com/release-notes/runtime/" target="_blank" rel="noopener">AI Runtime release notes</A> for current accelerator availability and Beta status in your region.</FONT></EM></FONT></P> </DIV> <P>&nbsp;</P> <HR /> <H2><STRONG>Next steps</STRONG></H2> <P><SPAN>Here's how you can try this out:</SPAN></P> <UL> <LI style="font-weight: 400;" aria-level="1"><STRONG>Single-node YOLO example:</STRONG><SPAN> Clone the </SPAN><A href="https://github.com/databricks-industry-solutions/cv-playground" target="_blank" rel="noopener"><SPAN>cv-playground</SPAN></A><SPAN> repository and run </SPAN><A href="https://github.com/databricks-industry-solutions/cv-playground/blob/main/projects/ultralytics_databricks_examples/air-yolo11n-detect-coco128-singleGPU.ipynb" target="_blank" rel="noopener"><SPAN>the notebook</SPAN></A><SPAN> in your Databricks workspace. </SPAN><SPAN>Attach it to an AI Runtime (Serverless GPU) with an AI base environment.</SPAN><SPAN> After running the cells, swap the COCO128 dataset for your own data and paths.</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><STRONG>Larger-scale instance segmentation:</STRONG><SPAN> For a bigger example using YOLO, check out the </SPAN><A href="https://github.com/databricks-industry-solutions/cv-playground/tree/main/projects/NuInsSeg" target="_blank" rel="noopener"><SPAN>NuInsSeg</SPAN></A><SPAN> project within the same repository.</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><STRONG>Monitor and improve:</STRONG><SPAN> Use the </SPAN><A href="https://docs.databricks.com/en/ai-gateway/inference-tables/" target="_blank" rel="noopener">Unity AI Gateway inference table</A><SPAN> to monitor model traffic, debug inputs and outputs, and feed insights back into your analytics or model improvement pipeline.</SPAN></LI> </UL> <H4>&nbsp;</H4> <P class="lia-align-center"><FONT size="3"><STRONG><I>Stay tuned for the follow-up post on multi-GPU and multi-node YOLO model training on AI Runtime!</I></STRONG></FONT></P> <P>&nbsp;</P> <HR /> <P><EM><STRONG>Acknowledgements:</STRONG> Thanks to <STRONG>Lin Yuan</STRONG> (Engineering, AI Runtime) for technical review and feedback.</EM></P> <P>&nbsp;</P> <DIV style="border-left: 4px solid #d0d0d0; padding-left: 16px; margin: 16px 0;"> <P style="margin: 0; font-size: 16px; line-height: 1.5;"><FONT color="#808080"><A id="ultralytics-licensing" target="_blank"></A><FONT size="3">*</FONT><EM><FONT size="3"><STRONG>Note on licensing:</STRONG> Ultralytics YOLO is dual-licensed: <A href="https://www.ultralytics.com/license" target="_blank" rel="noopener">AGPL-3.0</A> (default) or <A href="https://www.ultralytics.com/license" target="_blank" rel="noopener">Enterprise</A> for commercial use. Users should review <A href="https://www.ultralytics.com/license" target="_blank" rel="noopener">https://www.ultralytics.com/license</A> to determine which applies to their use case.</FONT></EM></FONT></P> </DIV> <P>&nbsp;</P> <P><A href="#intro" target="_self"><FONT size="2"><EM>^ Return to the top</EM></FONT></A></P> Fri, 12 Jun 2026 16:28:25 GMT mmt 2026-06-12T16:28:25Z https://community.databricks.com/t5/technical-blog/train-and-deploy-yolo-vision-model-on-databricks-ai-runtime-air/ba-p/151558 <P><I><SPAN>Train YOLO11n on <A href="https://www.databricks.com/blog/introducing-ai-runtime-scalable-serverless-nvidia-gpus-databricks-training-and-finetuning">Databricks AI Runtime</A>, register in Unity Catalog, and deploy to Model Serving with </SPAN></I><I><SPAN>base64</SPAN></I><I><SPAN> encoded image input and Unity AI Gateway inference logging.</SPAN></I></P> <UL> <LI style="font-weight: 400;" aria-level="1"><STRONG>End-to-end workflow for Ultralytics </STRONG><STRONG>YOLO11n</STRONG><STRONG> Object Detection:</STRONG><SPAN> Training on single-node Databricks AI Runtime, utilizing MLflow tracking, and using Unity Catalog for artifacts and model registry.</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><STRONG>Goal:</STRONG><SPAN> Provide a reproducible, production-ready path from GPU-trained </SPAN><SPAN>YOLO</SPAN><SPAN> to a Model Serving endpoint.</SPAN></LI> <LI style="font-weight: 400;" aria-level="1"><STRONG>Notebook Contents:</STRONG><SPAN> Covers environment setup, </SPAN><SPAN>COCO128</SPAN><SPAN> data preparation, MLflow autologging, custom </SPAN><SPAN>PyFunc</SPAN><SPAN> model registration (</SPAN><SPAN>base64</SPAN><SPAN> in, </SPAN><SPAN>bbox</SPAN><SPAN> out), evaluation, local serving validation, and deployment to a Model Serving endpoint with optional </SPAN><SPAN>request</SPAN><SPAN>/</SPAN><SPAN>response</SPAN><SPAN> logging via Unity AI Gateway inference tables.</SPAN></LI> </UL> <P>&nbsp;</P> Fri, 12 Jun 2026 16:28:25 GMT https://community.databricks.com/t5/technical-blog/train-and-deploy-yolo-vision-model-on-databricks-ai-runtime-air/ba-p/151558 mmt 2026-06-12T16:28:25Z