---
title: YOLO11n object detection on Databricks AI Runtime | Databricks on AWS
source: https://docs.databricks.com/aws/en/machine-learning/ai-runtime/examples/tutorials/sgc-yolo11n-detect-coco128
ingestedAt: "2026-06-18T08:09:15.270Z"
---

This notebook trains a [YOLO11n](https://docs.ultralytics.com/models/yolo11/) object detection model end to end on [Databricks AI Runtime](https://docs.databricks.com/aws/en/machine-learning/ai-runtime/). You train the model on the [COCO128](https://www.kaggle.com/datasets/ultralytics/coco128) dataset, track the run with [MLflow](https://docs.databricks.com/aws/en/mlflow/), register it to Unity Catalog, and optionally deploy it to a [Model Serving](https://docs.databricks.com/aws/en/machine-learning/model-serving/) endpoint.

The notebook prepares the dataset in [Unity Catalog volumes](https://docs.databricks.com/aws/en/volumes/) and wraps the trained model in a custom [MLflow PyFunc model](https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#creating-custom-pyfunc-models) that takes base64-encoded images and returns bounding boxes, so the same model works over a serving endpoint. It validates the model locally before deploying with the [Databricks SDK](https://docs.databricks.com/aws/en/dev-tools/sdk-python) and logging requests to [AI Gateway inference tables](https://docs.databricks.com/aws/en/ai-gateway/inference-tables).

> **⚠️ COCO128 is for demo only** (128 images). The model will overfit. For production, use larger datasets (1K+ images). This workflow scales directly by updating the data paths. See [NuInsSeg](https://github.com/databricks-industry-solutions/cv-playground/tree/main/projects/NuInsSeg) for a real-world example.

**Reference:** [Train and deploy a YOLO vision model on Databricks AI Runtime (serverless GPU)](https://community.databricks.com/t5/technical-blog/train-and-deploy-yolo-vision-model-on-databricks-ai-runtime-air/ba-p/151558) (Databricks Technical Blog)

## Connect to serverless GPU[​](#connect-to-serverless-gpu "Direct link to Connect to serverless GPU")

1.  Click **Serverless GPU** (top-right connect button).
2.  Open the ![Environment icon.](data:image/svg+xml;base64,PHN2ZyB3aWR0aD0iMTYiIGhlaWdodD0iMTYiIHZpZXdCb3g9IjAgMCAxNiAxNiIgZmlsbD0ibm9uZSIgeG1sbnM9Imh0dHA6Ly93d3cudzMub3JnLzIwMDAvc3ZnIj4KPHBhdGggZmlsbC1ydWxlPSJldmVub2RkIiBjbGlwLXJ1bGU9ImV2ZW5vZGQiIGQ9Ik00IDE1TDQgMTIuNjQ2NUMyLjg0NTc1IDEyLjMyIDIgMTEuMjU4OCAyIDEwQzIgOC43NDEyMiAyLjg0NTc1IDcuNjc5OTggNCA3LjM1MzUyTDQgMUg1LjVMNS41IDcuMzUzNTJDNi42NTQyNSA3LjY3OTk4IDcuNSA4Ljc0MTIyIDcuNSAxMEM3LjUgMTEuMjU4OCA2LjY1NDI1IDEyLjMyIDUuNSAxMi42NDY1TDUuNSAxNUg0Wk00Ljc1IDExLjI1QzQuMDU5NjQgMTEuMjUgMy41IDEwLjY5MDQgMy41IDEwQzMuNSA5LjMwOTY0IDQuMDU5NjQgOC43NSA0Ljc1IDguNzVDNS40NDAzNiA4Ljc1IDYgOS4zMDk2NCA2IDEwQzYgMTAuNjkwNCA1LjQ0MDM2IDExLjI1IDQuNzUgMTEuMjVaIiBmaWxsPSIjNkY2RjZGIi8+CjxwYXRoIGZpbGwtcnVsZT0iZXZlbm9kZCIgY2xpcC1ydWxlPSJldmVub2RkIiBkPSJNMTAuNSAxTDEwLjUgMy4zNTM1MkM5LjM0NTc1IDMuNjc5OTggOC41IDQuNzQxMjIgOC41IDZDOC41IDcuMjU4NzggOS4zNDU3NSA4LjMyMDAyIDEwLjUgOC42NDY0OFYxNUgxMlY4LjY0NjQ4QzEzLjE1NDMgOC4zMjAwMiAxNCA3LjI1ODc4IDE0IDZDMTQgNC43NDEyMiAxMy4xNTQzIDMuNjc5OTggMTIgMy4zNTM1MkwxMiAxSDEwLjVaTTExLjI1IDQuNzVDMTAuNTU5NiA0Ljc1IDEwIDUuMzA5NjQgMTAgNkMxMCA2LjY5MDM2IDEwLjU1OTYgNy4yNSAxMS4yNSA3LjI1QzExLjk0MDQgNy4yNSAxMi41IDYuNjkwMzYgMTIuNSA2QzEyLjUgNS4zMDk2NCAxMS45NDA0IDQuNzUgMTEuMjUgNC43NVoiIGZpbGw9IiM2RjZGNkYiLz4KPC9zdmc+Cg==) **Environment** panel in the right sidebar, then:
    *   Set **Accelerator**: **A10** (cost-efficient) or **H100** (faster).
    *   Set **Base environment**: **`AI v5`** (check the [release notes](https://docs.databricks.com/aws/en/release-notes/serverless/environment-version/five-gpu) for updates).
    *   Click **Apply**.

AI v5 pre-bundles `mlflow>=3` (skinny), `nvidia-ml-py`, `threadpoolctl`, and `torch`, so only `ultralytics` needs installing. If Model Serving fails on skinny MLflow, re-add `%pip install mlflow`; the Package Verification cell flags this.

## Environment setup[​](#environment-setup "Direct link to Environment setup")

Install required packages and configure Python environment for YOLO training on serverless GPU.

Python

    # ============================================================# PACKAGE INSTALLATION — AI v5 (Serverless GPU, single GPU)# ============================================================# AI v5 pre-bundles mlflow>=3, nvidia-ml-py, threadpoolctl (see top cell).# Only ultralytics needs install.%pip install ultralytics==8.3.204 -q# Note: %pip automatically restarts the Python environment after install.# YOLO_CONFIG_DIR is set in the 'Import Libraries' cell (cell 7).

Python

    # ============================================================# PACKAGE VERIFICATION# ============================================================import sysimport importlib.metadataprint("Checking required packages...\n")missing_packages = []installed_packages = {}# Check each required package using importlib.metadata# Check mlflow (full install) or mlflow-skinny (AI v5 pre-bundled) — either satisfies mlflow>=3.0_mlflow_found = Falsefor _pkg in ['mlflow', 'mlflow-skinny']:    try:        _ver = importlib.metadata.version(_pkg)        installed_packages['mlflow'] = _ver        print(f"✓ mlflow ({_pkg}): {_ver}")        _mlflow_found = True        break    except importlib.metadata.PackageNotFoundError:        passif not _mlflow_found:    missing_packages.append('mlflow>=3.0')    print("✗ mlflow: NOT INSTALLED")packages_to_check = [    ('ultralytics', 'ultralytics==8.3.204'),    ('opencv-python', 'opencv-python (provides cv2)'),    ('nvidia-ml-py', 'nvidia-ml-py==13.580.82'),    ('threadpoolctl', 'threadpoolctl==3.1.0')]for package_name, package_spec in packages_to_check:    try:        version = importlib.metadata.version(package_name)        installed_packages[package_name] = version        print(f"✓ {package_name}: {version}")    except importlib.metadata.PackageNotFoundError:        missing_packages.append(package_spec)        print(f"✗ {package_name}: NOT INSTALLED")print("\n" + "="*60)if missing_packages:    print("[ACTION REQUIRED] Missing packages detected!")    print("\nInstall missing packages and restart kernel.")    for pkg in missing_packages:        print(f"   - {pkg}")else:    print("[OK] All required packages are installed!")    print(f"   Python version: {sys.version.split()[0]}")print("="*60)# --- AI v5 environment probe (added for v5 validation) ---import torchprint("\n--- AI v5 environment probe ---")print("torch:", torch.__version__, "| CUDA available:", torch.cuda.is_available())try:    print("GPU:", torch.cuda.get_device_name(0))except Exception as e:    print("GPU: (could not query)", e)_dists = {d.metadata['Name'].lower() for d in importlib.metadata.distributions()}print("mlflow-skinny present:", "mlflow-skinny" in _dists, "| full mlflow present:", "mlflow" in _dists)try:    import mlflow.deployments, mlflow.models  # serving/registration path (the skinny risk)    print("[OK] mlflow.deployments + mlflow.models importable -> Model Serving step should work")except Exception as e:    print("[WARN] mlflow serving API import FAILED on skinny -> re-add `%pip install mlflow`:", e)

Python

    # Note: You may see a FutureWarning about pynvml being deprecated.# This is expected - PyTorch internally imports pynvml, but we have# the correct nvidia-ml-py package installed. The warning is harmless.import osimport pandas as pdimport torchimport mlflowimport uuid# Set Ultralytics config directory to a unique writable locationconfig_dir = f'/tmp/yolo_config_{uuid.uuid4().hex[:8]}'os.environ['YOLO_CONFIG_DIR'] = config_diros.makedirs(config_dir, exist_ok=True)from ultralytics import YOLO, settingsfrom mlflow.types.schema import Schema, ColSpecfrom mlflow.models.signature import ModelSignatureprint("[OK] All libraries imported successfully")

## Helper functions[​](#helper-functions "Direct link to Helper functions")

Utility functions for the complete YOLO training and deployment workflow:

**Data Management:**

*   `download_file()` - Download models and configs to UC Volume
*   `download_and_extract_dataset()` - Download and extract COCO128
*   `split_dataset()` - Create reproducible train/val/test splits

**MLflow Integration:**

*   `infer_model_signature()` - Automatically infer model signature from predictions
*   `setup_mlflow_experiment()` - Configure MLflow with system metrics
*   `register_yolo_model()` - Register model to Unity Catalog with custom wrapper

**Model Evaluation:**

*   `evaluate_model_on_split()` - Evaluate and visualize predictions on data splits

**Custom Wrapper:**

*   `YOLOWrapper` - [MLflow PyFunc wrapper](https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#creating-custom-pyfunc-models) for YOLO models
    *   **Input:** [Base64-encoded](https://en.wikipedia.org/wiki/Base64) images (universal format, works across network boundaries)
    *   **Output:** DataFrame with class, confidence, bounding boxes (11 columns)
    *   **Purpose:** Enables deployment to [Model Serving endpoints](https://docs.databricks.com/aws/en/machine-learning/model-serving/create-manage-serving-endpoints)
    *   **Production-ready:** Tested locally before deployment

Python

    # ============================================================# HELPER FUNCTIONS# ============================================================import osimport shutilimport requestsimport zipfileimport ioimport randomimport yamlimport globimport jsonfrom datetime import datetimeimport pandas as pdimport matplotlib.pyplot as pltfrom ultralytics import YOLOimport mlflowfrom mlflow import MlflowClientimport importlib.metadatadef download_file(url, destination, description="file"):    """Download a file from URL to destination path."""    if os.path.exists(destination):        print(f"[INFO] {description} already exists at: {destination}")        print(f"   Skipping download")        return True    print(f"Downloading {description}...")    print(f"   From: {url}")    print(f"   To: {destination}")    try:        response = requests.get(url, stream=True)        if response.status_code == 200:            os.makedirs(os.path.dirname(destination), exist_ok=True)            with open(destination, 'wb') as f:                for chunk in response.iter_content(chunk_size=8192):                    f.write(chunk)            print(f"[OK] Downloaded successfully")            if destination.endswith('.pt'):                print(f"   Size: {os.path.getsize(destination) / (1024*1024):.2f} MB")            return True        else:            print(f"[ERROR] Download failed with status code: {response.status_code}")            return False    except Exception as e:        print(f"[ERROR] Download failed: {e}")        return Falsedef download_and_extract_dataset(download_url, extraction_path):    """Download and extract a zip dataset."""    print("Downloading dataset...")    response = requests.get(download_url)    print("Extracting dataset...")    z = zipfile.ZipFile(io.BytesIO(response.content))    z.extractall(extraction_path)    print(f"[OK] Dataset downloaded and extracted to: {extraction_path}")    return Truedef split_dataset(source_images_dir, source_labels_dir, base_images_dir, base_labels_dir,                  train_ratio=0.625, val_ratio=0.1875, random_seed=42):    """Split dataset into train/val/test sets with reproducible random seed."""    print("=" * 60)    print("DATASET SPLITTING")    print("=" * 60)    random.seed(random_seed)    print(f"\nRandom seed: {random_seed}")    test_ratio = 1.0 - train_ratio - val_ratio    print(f"Split ratios: Train={train_ratio:.1%}, Val={val_ratio:.1%}, Test={test_ratio:.1%}\n")    # Get all images    all_images = sorted([f for f in os.listdir(source_images_dir) if f.endswith('.jpg')])    print(f"Total images: {len(all_images)}")    # Shuffle and split    random.shuffle(all_images)    train_size = int(len(all_images) * train_ratio)    val_size = int(len(all_images) * val_ratio)    train_images = all_images[:train_size]    val_images = all_images[train_size:train_size + val_size]    test_images = all_images[train_size + val_size:]    print(f"Split sizes: Train={len(train_images)}, Val={len(val_images)}, Test={len(test_images)}\n")    # Create directories    for split_name in ['train', 'val', 'test']:        os.makedirs(f"{base_images_dir}/{split_name}", exist_ok=True)        os.makedirs(f"{base_labels_dir}/{split_name}", exist_ok=True)    # Copy files    print("Copying files to splits...")    for split_name, image_list in [('train', train_images), ('val', val_images), ('test', test_images)]:        print(f"  Processing {split_name} split ({len(image_list)} images)...")        for img_name in image_list:            # Copy image            src_img = os.path.join(source_images_dir, img_name)            dst_img = os.path.join(base_images_dir, split_name, img_name)            shutil.copy2(src_img, dst_img)            # Copy label if exists            label_name = img_name.replace('.jpg', '.txt')            src_label = os.path.join(source_labels_dir, label_name)            dst_label = os.path.join(base_labels_dir, split_name, label_name)            if os.path.exists(src_label):                shutil.copy2(src_label, dst_label)        print(f"    [OK] {split_name}: {len(image_list)} images copied")    print(f"\n[OK] Dataset split complete!")    print("=" * 60)    return len(train_images), len(val_images), len(test_images)def infer_model_signature(model_path, sample_image_path):    """Infer MLflow model signature using actual model predictions."""    import base64    print("[INFO] Inferring model signature...\n")    # Load YOLO model    model = YOLO(model_path)    # Read and encode image as base64    with open(sample_image_path, 'rb') as f:        image_bytes = f.read()    image_base64 = base64.b64encode(image_bytes).decode('utf-8')    # Create input example    input_example = pd.DataFrame({"image_base64": [image_base64]})    # Create YOLOWrapper instance and get predictions to infer output schema    wrapper = YOLOWrapper()    # Simulate load_context    class MockContext:        def __init__(self, model_path):            self.artifacts = {"yolo_model": model_path}    wrapper.load_context(MockContext(model_path))    # Get output example by running prediction    output_example = wrapper.predict(None, input_example)    # Use MLflow's infer_signature to automatically create signature    signature = mlflow.models.infer_signature(input_example, output_example)    print(f"[OK] Model signature inferred successfully!")    print(f"   Input: DataFrame with 'image_base64' column (base64 string)")    print(f"   Output: DataFrame with {len(output_example.columns)} columns")    print(f"   Columns: {', '.join(output_example.columns.tolist())}")    # Optional: Show how to use manual schema (commented out)    # from mlflow.types.schema import Schema, ColSpec    # from mlflow.models.signature import ModelSignature    # input_schema = Schema([ColSpec("string", "image_base64")])    # output_schema = Schema([    #     ColSpec("string", "class_name"),    #     ColSpec("long", "class_num"),    #     ColSpec("double", "confidence"),    #     ColSpec("double", "bbox_x1"),    #     ColSpec("double", "bbox_y1"),    #     ColSpec("double", "bbox_x2"),    #     ColSpec("double", "bbox_y2"),    #     ColSpec("double", "bbox_center_x"),    #     ColSpec("double", "bbox_center_y"),    #     ColSpec("double", "bbox_width"),    #     ColSpec("double", "bbox_height")    # ])    # signature = ModelSignature(inputs=input_schema, outputs=output_schema)    return signature, input_exampledef setup_mlflow_experiment(use_workspaceUsers_path=False, expt_name_suffix="Experiments_YOLO_CoCo"):    """Setup MLflow experiment with system metrics enabled.    Args:        use_workspaceUsers_path: If True, derives experiment directory from the                                current notebook location under /Workspace/Users/.                                If False (default), uses '/Workspace/Shared'.        expt_name_suffix: Name of the experiment folder.                          Defaults to 'Experiments_YOLO_CoCo'.    Returns:        tuple: (experiment_name, experiment_id)    """    if use_workspaceUsers_path:        # Extract username from notebook path and build /Workspace/Users/{username}        # This avoids Git folder paths where MLflow experiment creation is not permitted        notebook_path = dbutils.notebook.entry_point.getDbutils().notebook().getContext().notebookPath().get()        path_parts = notebook_path.strip('/').split('/')        if 'Users' in path_parts:            user_idx = path_parts.index('Users')            username = path_parts[user_idx + 1]            experiment_base_path = f"/Workspace/Users/{username}"        else:            # Fallback if path doesn't contain Users            experiment_base_path = "/Workspace/Shared"    else:        # Default: Use /Workspace/Shared        experiment_base_path = "/Workspace/Shared"    experiment_name = f"{experiment_base_path}/{expt_name_suffix}"    # Ensure the experiment directory exists    os.makedirs(experiment_base_path, exist_ok=True)    os.environ['MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING'] = "true"    os.environ['MLFLOW_EXPERIMENT_NAME'] = experiment_name    mlflow.set_experiment(experiment_name)    experiment_id = mlflow.get_experiment_by_name(experiment_name).experiment_id    if 'MLFLOW_RUN_ID' in os.environ:        del os.environ['MLFLOW_RUN_ID']    print(f"[OK] MLflow experiment initialized: {experiment_name}")    print(f"   Experiment ID: {experiment_id}")    print(f"   System metrics: ENABLED")    return experiment_name, experiment_idclass YOLOWrapper(mlflow.pyfunc.PythonModel):    """Custom MLflow wrapper for YOLO models using base64-encoded images."""    def load_context(self, context):        """Load YOLO model from artifacts."""        from ultralytics import YOLO        model_path = context.artifacts["yolo_model"]        self.model = YOLO(model_path, task='detect')    def _format_predictions(self, predictions):        """Format YOLO prediction results with bounding boxes.        Args:            predictions: YOLO prediction results from model.predict()        Returns:            pd.DataFrame with class, confidence, and bounding box coordinates        """        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)):                    # Get bounding box coordinates in both formats                    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):        """Run YOLO prediction on base64-encoded images.        Args:            context: MLflow context            model_input: DataFrame with 'image_base64' column (base64-encoded images)        Returns:            pd.DataFrame with detection results including bounding boxes        """        import pandas as pd        import base64        from PIL import Image        import io        import numpy as np        if not isinstance(model_input, pd.DataFrame):            raise ValueError("Input must be a DataFrame with 'image_base64' column")        if 'image_base64' not in model_input.columns:            raise ValueError("DataFrame must contain 'image_base64' column with base64-encoded images")        # Process base64-encoded images        all_predictions = []        for image_base64 in model_input['image_base64'].tolist():            # Decode base64 to image            image_bytes = base64.b64decode(image_base64)            image = Image.open(io.BytesIO(image_bytes))            image_array = np.array(image)            # Run prediction            predictions = self.model.predict(image_array, verbose=False)            all_predictions.extend(predictions)        return self._format_predictions(all_predictions)def register_yolo_model(run_id, model_path, catalog_name, schema_name, model_name,                       signature=None, input_example=None, data_yaml_path=None):    """Register YOLO model to Unity Catalog with custom wrapper."""    registered_model_name = f"{catalog_name}.{schema_name}.{model_name}"    ultralytics_version = importlib.metadata.version('ultralytics')    cloudpickle_version = importlib.metadata.version('cloudpickle')    print(f"\n[INFO] Registering model to Unity Catalog...")    print(f"   Model name: {registered_model_name}")    print(f"   Using custom YOLO wrapper (base64 input, bbox output)")    print(f"   Pinning CloudPickle version: {cloudpickle_version}")    with mlflow.start_run(run_id=run_id):        if data_yaml_path:            mlflow.log_artifact(data_yaml_path, "input_data")        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=[                f"ultralytics=={ultralytics_version}",                f"cloudpickle=={cloudpickle_version}",                "torch",                "torchvision",                "pillow",                "numpy"            ]        )    print(f"   [OK] Model registered: {registered_model_name}")    return registered_model_namedef evaluate_model_on_split(model, image_dir, split_name, output_dir, run_id,                           registered_model_name, organized_run_name, num_samples=3):    """Evaluate model on a dataset split and save results."""    print("=" * 60)    print(f"{split_name.upper()} SET EVALUATION")    print("=" * 60)    os.makedirs(output_dir, exist_ok=True)    images = glob.glob(f"{image_dir}/*.jpg")    if not images:        print(f"[WARNING] No {split_name} images found")        return    print(f"\n{split_name.capitalize()} set: {len(images)} images\n")    # Visualize sample predictions    sample_images = images[:num_samples]    fig, axes = plt.subplots(1, len(sample_images), figsize=(15, 5))    if len(sample_images) == 1:        axes = [axes]    results = []    for i, img_path in enumerate(sample_images):        print(f"Sample {i+1}/{len(sample_images)}: {img_path.split('/')[-1]}")        predictions = model.predict(img_path, verbose=False)        if len(predictions) > 0:            result = predictions[0]            annotated_img = result.plot()            axes[i].imshow(annotated_img)            axes[i].axis('off')            if result.boxes is not None:                num_detections = len(result.boxes)                axes[i].set_title(f"{img_path.split('/')[-1]}\n{num_detections} objects", fontsize=10)                print(f"   [OK] Detections: {num_detections} objects")                img_results = {                    "image": img_path.split('/')[-1],                    "num_detections": num_detections,                    "detections": []                }                for j in range(min(num_detections, 3)):                    class_name = result.names[int(result.boxes.cls[j])]                    confidence = float(result.boxes.conf[j])                    print(f"      - {class_name}: {confidence:.3f}")                    img_results["detections"].append({                        "class_name": class_name,                        "confidence": confidence                    })                results.append(img_results)        print()    plt.tight_layout()    plt.suptitle(f"{split_name.capitalize()} Set Predictions - Run {run_id[:8]}", fontsize=14, y=1.02)    plot_path = os.path.join(output_dir, f"{split_name}_predictions.png")    plt.savefig(plot_path, dpi=150, bbox_inches='tight')    print(f"[OK] Plot saved to: {plot_path}")    plt.show()    # Save results JSON    json_path = os.path.join(output_dir, f"{split_name}_results.json")    with open(json_path, 'w') as f:        json.dump({            "run_id": run_id,            "registered_model": registered_model_name,            "timestamp": organized_run_name.split('_run_')[0],            "num_images": len(images),            "sample_results": results        }, f, indent=2)    print(f"[OK] Results saved to: {json_path}")    # Log to MLflow    with mlflow.start_run(run_id=run_id):        mlflow.log_artifact(plot_path, split_name)        mlflow.log_artifact(json_path, split_name)    print(f"\n[OK] {split_name.upper()} SET EVALUATION COMPLETE")    print("=" * 60)print("[OK] Helper functions loaded successfully")

## Unity Catalog configuration[​](#unity-catalog-configuration "Direct link to Unity Catalog configuration")

Configure catalog, schema, volume, and project paths.

Python

    dbutils.widgets.removeAll()

Python

    # Define widgets for catalog, schema, volume, model name, and deployment approvaldbutils.widgets.text("catalog_name", "main", "Catalog Name")dbutils.widgets.text("schema_name", "default", "Schema Name")dbutils.widgets.text("volume_name", "yolo_sgc", "Volume Name")dbutils.widgets.text("model_name", "yolo11n_coco128_sgc", "Model Name")dbutils.widgets.dropdown("proceed_with_deployment", "false", ["false", "true"], "Proceed with Deployment")# External artifact sources, pinned to fixed versions. For production or# network-restricted workspaces, stage these in a Unity Catalog volume or an# internal mirror and point these widgets there instead of the public internet.dbutils.widgets.text("model_source_url", "https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11n.pt", "Pretrained Model URL")dbutils.widgets.text("dataset_config_url", "https://raw.githubusercontent.com/ultralytics/ultralytics/v8.3.204/ultralytics/cfg/datasets/coco128.yaml", "Dataset Config URL")# Get widget valuescatalog_name = dbutils.widgets.get("catalog_name")schema_name = dbutils.widgets.get("schema_name")volume_name = dbutils.widgets.get("volume_name")model_name = dbutils.widgets.get("model_name")proceed_with_deployment_str = dbutils.widgets.get("proceed_with_deployment")model_source_url = dbutils.widgets.get("model_source_url")dataset_config_url = dbutils.widgets.get("dataset_config_url")print(f"[Configuration]")print(f"   Catalog: {catalog_name}")print(f"   Schema: {schema_name}")print(f"   Volume: {volume_name}")print(f"   Model: {model_name}")print(f"   Proceed with Deployment: {proceed_with_deployment_str}")print(f"\nUsing catalog: {catalog_name} (already exists)")# Create schema if it doesn't existspark.sql(f"CREATE SCHEMA IF NOT EXISTS `{catalog_name}`.`{schema_name}`")print(f"[OK] Schema: {catalog_name}.{schema_name}")# Create volume for persistent storagespark.sql(f"CREATE VOLUME IF NOT EXISTS `{catalog_name}`.`{schema_name}`.`{volume_name}`")print(f"[OK] Volume: {catalog_name}.{schema_name}.{volume_name}")

Python

    # Get Unity Catalog parameterscatalog_name = dbutils.widgets.get("catalog_name")schema_name = dbutils.widgets.get("schema_name")volume_name = dbutils.widgets.get("volume_name")# Construct volume path from parametersproject_location = f'/Volumes/{catalog_name}/{schema_name}/{volume_name}/'print(f"Using Unity Catalog Volume: {catalog_name}.{schema_name}.{volume_name}")print(f"Volume path: {project_location}")# Create subdirectories in the volumeos.makedirs(f'{project_location}runs/', exist_ok=True)       # Training runs (organized by task/model/dataset)os.makedirs(f'{project_location}data/', exist_ok=True)       # Dataset storageos.makedirs(f'{project_location}raw_model/', exist_ok=True)  # Pretrained models# Ephemeral /tmp/ location for faster I/O during trainingtmp_project_location = "/tmp/training_results/"os.makedirs(tmp_project_location, exist_ok=True)print(f"\n[OK] Project directories created:")print(f"   Runs: {project_location}runs/")print(f"   Data: {project_location}data/")print(f"   Raw models: {project_location}raw_model/")print(f"   Temp (training): {tmp_project_location}  # Ephemeral, fast I/O")

## Project folder structure[​](#project-folder-structure "Direct link to Project folder structure")

Unity Catalog volume organization:

    /Volumes/{catalog}/{schema}/{volume}/├── data/│   ├── coco128.yaml                    # Dataset configuration│   └── coco128/│       ├── images/│       │   ├── train2017/              # Original 128 images (from zip)│       │   ├── train/  val/  test/     # Custom splits (80/24/24)│       └── labels/│           ├── train2017/              # Original labels│           └── train/  val/  test/     # Split labels│├── raw_model/│   └── yolo11n.pt                      # Pretrained YOLO11n weights│└── runs/    └── {task}_{model}_{dataset}_{timestamp}_run_{mlflow_run_id}/        ├── train/                      # MLflow training outputs        │   ├── weights/ (best.pt, last.pt)        │   └── results.csv, confusion_matrix.png        ├── validation_metrics/         # YOLO validation outputs        ├── validation_samples/         # Custom evaluation samples        └── test_samples/               # Test evaluation samples

**Run Naming:** `detection_yolo11n_coco128_20260120_143052_run_{mlflow_run_id}`

*   Includes task, model, dataset, timestamp, and MLflow run ID for easy identification

## Download pretrained YOLO model[​](#download-pretrained-yolo-model "Direct link to Download pretrained YOLO model")

Download YOLO11n pretrained weights to Unity Catalog volume.

> **🔒 Artifact sources:** This notebook downloads the pretrained model and dataset config from public Ultralytics sources, pinned to fixed versions. For production or network-restricted workspaces, stage these artifacts in a Unity Catalog volume or an internal mirror and repoint the **Pretrained Model URL** and **Dataset Config URL** widgets to that location.

Python

    # Download pretrained YOLO11n model to Unity Catalog Volumemodel_path = f"{project_location}raw_model/yolo11n.pt"model_url = dbutils.widgets.get("model_source_url")download_file(model_url, model_path, "YOLO11n model")print(f"\n[OK] Pretrained model ready at: {model_path}")

## Dataset preparation[​](#dataset-preparation "Direct link to Dataset preparation")

Download and configure COCO128 dataset.

Python

    # Download COCO128 dataset configuration to UC Volumeimport yaml# Create data directory in UC Volumeos.makedirs(f'{project_location}data/coco128', exist_ok=True)# Download config directly to UC Volumeconfig_url = dbutils.widgets.get("dataset_config_url")config_path = f"{project_location}data/coco128.yaml"  # UC Volume pathdownload_file(config_url, config_path, "COCO128 config")# Load and update configurationwith open(config_path, 'r') as f:    data = yaml.safe_load(f)print(f"\n[Dataset Configuration]")print(f"   Dataset: {data.get('path', 'coco128')}")print(f"   Classes: {data.get('nc', 'unknown')}")print(f"   Download URL: {data.get('download', 'N/A')}")# Update paths for Unity Catalog Volumedata['path'] = f"{project_location}data/coco128"# Check if dataset already existsdataset_images_dir = f"{project_location}data/coco128/images/train2017"if os.path.exists(dataset_images_dir) and len(os.listdir(dataset_images_dir)) > 0:    print(f"\n[INFO] Dataset already exists at: {dataset_images_dir}")    print(f"   Found {len(os.listdir(dataset_images_dir))} images")    print(f"   Skipping download")else:    # Download and extract dataset    extraction_path = f"{project_location}data"    download_and_extract_dataset(data['download'], extraction_path)# Save updated configuration to UC Volumedata_yaml_path = f"{project_location}data/coco128.yaml"with open(data_yaml_path, 'w') as f:    yaml.dump(data, f, default_flow_style=False)print(f"\n[OK] Dataset configuration saved to UC Volume: {data_yaml_path}")print(f"   All dataset files in: {project_location}data/coco128/")

## Dataset splits[​](#dataset-splits "Direct link to Dataset splits")

Split COCO128 into train (62.5%), val (18.75%), and test (18.75%) sets with reproducible random seed.

Python

    # Split dataset into train/val/test with reproducible random seedsource_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"# Skip if splits already exist (deterministic with seed=42)train_dir = f"{base_images_dir}/train"val_dir = f"{base_images_dir}/val"test_dir = f"{base_images_dir}/test"if (os.path.isdir(train_dir) and len(os.listdir(train_dir)) > 0    and os.path.isdir(val_dir) and len(os.listdir(val_dir)) > 0    and os.path.isdir(test_dir) and len(os.listdir(test_dir)) > 0):    train_size = len(os.listdir(train_dir))    val_size = len(os.listdir(val_dir))    test_size = len(os.listdir(test_dir))    print(f"[INFO] Splits already exist \u2014 skipping re-split")else:    train_size, val_size, test_size = split_dataset(        source_images_dir=source_images_dir,        source_labels_dir=source_labels_dir,        base_images_dir=base_images_dir,        base_labels_dir=base_labels_dir,        train_ratio=0.625,  # 62.5%        val_ratio=0.1875,   # 18.75%        random_seed=42    )print(f"\nSplit summary:")print(f"  - Train: {train_size} images (62.5%)")print(f"  - Val: {val_size} images (18.75%)")print(f"  - Test: {test_size} images (18.75%)")print(f"  - Random seed: 42")

Python

    # Update data.yaml to use train/val/test splitswith open(data_yaml_path, 'r') as f:    yaml_content = yaml.safe_load(f)# Update pathsyaml_content['train'] = f"{project_location}data/coco128/images/train"yaml_content['val'] = f"{project_location}data/coco128/images/val"yaml_content['test'] = f"{project_location}data/coco128/images/test"# Save updated configurationwith open(data_yaml_path, 'w') as f:    yaml.dump(yaml_content, f, default_flow_style=False)print(f"[OK] data.yaml updated with train/val/test splits")print(f"   Train: {yaml_content['train']}")print(f"   Val: {yaml_content['val']}")print(f"   Test: {yaml_content['test']}")

## MLflow configuration[​](#mlflow-configuration "Direct link to MLflow configuration")

Infer model signature and configure experiment tracking with system metrics.

Python

    # Infer model signature from sample prediction# This defines the input/output schema for the serving endpoint# Input: base64-encoded images# Output: class, confidence, bounding boxes (11 columns)model_path = f"{project_location}raw_model/yolo11n.pt"# Find a sample image from training setsample_images = glob.glob(f"{project_location}data/coco128/images/train/*.jpg")if sample_images:    signature, input_example = infer_model_signature(model_path, sample_images[0])    print(f"\n[OK] Signature and input example ready for model registration")else:    print("[WARNING] No sample images found. Run dataset preparation first.")    signature = None    input_example = None

Python

    # Validate that signature and input_example are set before training/registeringif 'signature' not in dir() or signature is None:    raise ValueError(        "signature is None — re-run 'Infer MLflow Model Signature' cell before registering."    )if 'input_example' not in dir() or input_example is None:    raise ValueError(        "input_example is None — re-run 'Infer MLflow Model Signature' cell before registering."    )print(f"✓ signature: {signature}")print(f"✓ input_example shape: {input_example.shape}, columns: {list(input_example.columns)}")

Python

    # Configure YOLO to use MLflowfrom ultralytics import settingssettings.update({"mlflow": True})# Enable MLflow autologging for system metricsmlflow.autolog(disable=False)print(f"\n[MLflow Configuration]")print(f"   YOLO MLflow integration: Enabled")print(f"   MLflow autologging: Enabled")print(f"   System metrics: Enabled")

Python

    # Setup MLflow experiment with system metrics# Default: /Workspace/Shared/Experiments_YOLO_CoCo# To use /Workspace/Users/ path instead: setup_mlflow_experiment(use_workspaceUsers_path=True, expt_name_suffix="NightlyJob_YOLO_CoCo")experiment_name, experiment_id = setup_mlflow_experiment(use_workspaceUsers_path=False, expt_name_suffix="Experiments_YOLO_CoCo")print(f"\n[Ready for Training]")print(f"   Experiment: {experiment_name}")print(f"   Experiment ID: {experiment_id}")

## Model training[​](#model-training "Direct link to Model training")

Train YOLO11n with MLflow tracking and register to Unity Catalog.

Python

    from datetime import datetimeimport uuid# Close any active MLflow runsmlflow.end_run()# Create unique timestamp and temp directorytimestamp = datetime.now().strftime("%Y%m%d_%H%M%S")tmp_project_location_unique = f"/tmp/training_results_{uuid.uuid4().hex[:8]}/"os.makedirs(tmp_project_location_unique, exist_ok=True)# Model configurationmodel_path = f"{project_location}raw_model/yolo11n.pt"model_task = "detection"     # Task typemodel_arch = "yolo11n"       # Model architecturedataset_name = "coco128"     # Dataset name# ── GPU-aware batch size ──────────────────────────────────────# Auto-detect GPU and scale batch_size to available VRAM.# YOLO11n memory footprint at img_size=640:#   batch=8  → ~2.5 GB  |  batch=16 → ~4.5 GB  |  batch=32 → ~8 GB_gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu"_vram_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 if torch.cuda.is_available() else 0if _vram_gb >= 70:       # H100 (80 GB)    batch_size = 32elif _vram_gb >= 20:     # A10G (22.5 GB)    batch_size = 16else:                    # Fallback    batch_size = 8# Append GPU suffix to model/volume names for registry clarityif "H100" in _gpu_name:    gpu_suffix = "_h100"elif "A10" in _gpu_name:    gpu_suffix = "_a10"else:    gpu_suffix = ""model_name = dbutils.widgets.get("model_name") + gpu_suffixprint(f"[GPU] {_gpu_name} ({_vram_gb:.0f} GB VRAM) → batch_size={batch_size}")print(f"   Volume: {catalog_name}.{schema_name}.{volume_name}")print(f"   Model:  {catalog_name}.{schema_name}.{model_name}")# ── Training hyperparameters ──────────────────────────────────# Adjust these based on your dataset size and training goals:#   - epochs/patience: increase for larger datasets (e.g. epochs=300, patience=50)#   - learning_rate: lower (0.0001) for fine-tuning, higher (0.01) for scratch#   - dropout: increase (0.3-0.5) if overfitting, decrease (0.0) for large datasets#   - weight_decay: standard 0.0005 works for most casesepochs = 100learning_rate = 0.001        # Initial learning rate (YOLO lr0)patience = 5                 # Early stopping patience (epochs without improvement)dropout = 0.2                # Dropout rate for regularizationweight_decay = 0.0005        # L2 regularizationprint(f"\nTraining configuration:")print(f"  GPU: {_gpu_name} ({_vram_gb:.0f} GB)")print(f"  Task: {model_task}")print(f"  Model: {model_arch}")print(f"  Dataset: {dataset_name}")print(f"  Weights: {model_path}")print(f"  Data config: {data_yaml_path}")print(f"  Output: {tmp_project_location_unique}")print(f"  Epochs: {epochs}")print(f"  Batch size: {batch_size} (auto-scaled to GPU)")print(f"  Learning rate (lr0): {learning_rate}")print(f"  Patience: {patience}")print(f"  Dropout: {dropout}")print(f"  Weight decay: {weight_decay}")print()# Initialize and train modelprint(f"Loading YOLO model...")model = YOLO(model_path)print("Starting training...\n")results = model.train(    task="detect",    batch=batch_size,    device=0,                # Single GPU for Serverless    data=data_yaml_path,    epochs=epochs,    lr0=learning_rate,       # Initial learning rate    project=tmp_project_location_unique,    name=f"run_{timestamp}",    exist_ok=True,    fliplr=1,    flipud=1,    perspective=0.001,    degrees=0.45,    amp=True,    patience=patience,    dropout=dropout,    weight_decay=weight_decay,    save=True,    save_period=10)# Get MLflow run IDrun_id = mlflow.last_active_run().info.run_idprint(f"\n[OK] Training complete! MLflow Run ID: {run_id}")# Copy training results to Unity Catalog Volume with enhanced namingprint(f"\n[INFO] Copying training results to Unity Catalog Volume...")training_run_dir = os.path.join(tmp_project_location_unique, f"run_{timestamp}")# Enhanced folder naming: {task}_{model}_{dataset}_{datetime}_run_{mlflow_run_id}organized_run_name = f"{model_task}_{model_arch}_{dataset_name}_{timestamp}_run_{run_id}"volume_run_dir = os.path.join(project_location, "runs", organized_run_name)# Create train subfolder for MLflow training outputsvolume_train_dir = os.path.join(volume_run_dir, "train")if os.path.exists(training_run_dir):    shutil.copytree(training_run_dir, volume_train_dir, dirs_exist_ok=True)    print(f"   [OK] Training outputs copied to: {volume_train_dir}")# Validate model and save outputs to UC Volumeprint("\n[INFO] Validating model...")val_metrics_dir = os.path.join(tmp_project_location_unique, "validation_metrics")os.makedirs(val_metrics_dir, exist_ok=True)val_results = model.val(    project=tmp_project_location_unique,    name="validation_metrics")# Copy validation metrics to UC Volume (at run level, not in train/)if os.path.exists(val_metrics_dir):    volume_val_metrics_dir = os.path.join(volume_run_dir, "validation_metrics")    shutil.copytree(val_metrics_dir, volume_val_metrics_dir, dirs_exist_ok=True)    print(f"   [OK] Validation metrics copied to: {volume_val_metrics_dir}")# Save best model in YOLO native formatprint("\n[INFO] Saving best model...")best_model = YOLO(str(model.trainer.best))best_model_path = f"/tmp/best_yolo_model_{timestamp}.pt"best_model.save(best_model_path)print(f"   Saved to: {best_model_path}")# Register model to Unity Catalog using model_name widgetregistered_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,  # Use widget parameter    signature=signature,    input_example=input_example,    data_yaml_path=data_yaml_path)print(f"\n[OK] Training complete!")print(f"\n[Model Details]")print(f"   - Name: {registered_model_name}")print(f"   - Run ID: {run_id}")print(f"   - Location: Unity Catalog Model Registry")print(f"   - Format: Custom YOLO wrapper (base64 input, bbox output)")print(f"\n[Training Artifacts]")print(f"   - Volume location: {volume_run_dir}")print(f"   - Run name: {organized_run_name}")print(f"   - Structure: train/, validation_metrics/")print(f"\n[View Results]")print(f"   {mlflow.get_tracking_uri()}/#/experiments/{experiment_id}/runs/{run_id}")

## Model evaluations[​](#model-evaluations "Direct link to Model evaluations")

**Split Evaluation (Native YOLO):** Assess model accuracy using file paths from UC Volume. Validates model quality on validation/test sets.

**Local Serving Test (MLflow PyFunc):** Validate production serving format using base64-encoded images. Ensures endpoint compatibility before deployment.

### Split evaluation[​](#split-evaluation "Direct link to Split evaluation")

Evaluate model performance on validation and test sets before deployment.

Python

    # Load model from MLflow and evaluate on validation setmodel_uri = f"runs:/{run_id}/model"model_path = mlflow.artifacts.download_artifacts(model_uri)# Find the .pt fileimport glob as glob_modulept_files = glob_module.glob(f"{model_path}/**/*.pt", recursive=True)if pt_files:    loaded_model = YOLO(pt_files[0], task='detect')    print(f"[OK] Model loaded from MLflow\n")    # Evaluate on validation set    val_image_dir = f"{project_location}data/coco128/images/val"    val_output_dir = os.path.join(volume_run_dir, "validation_samples")    evaluate_model_on_split(        model=loaded_model,        image_dir=val_image_dir,        split_name="validation",        output_dir=val_output_dir,        run_id=run_id,        registered_model_name=registered_model_name,        organized_run_name=organized_run_name,        num_samples=3    )else:    print("[ERROR] Model file not found")

Python

    # Evaluate model on test set (uses loaded_model from validation cell)test_image_dir = f"{project_location}data/coco128/images/test"test_output_dir = os.path.join(volume_run_dir, "test_samples")evaluate_model_on_split(    model=loaded_model,    image_dir=test_image_dir,    split_name="test",    output_dir=test_output_dir,    run_id=run_id,    registered_model_name=registered_model_name,    organized_run_name=organized_run_name,    num_samples=3)

### Local serving test[​](#local-serving-test "Direct link to Local serving test")

Validate the registered model's serving format (base64 in → bbox out) before deployment.

Python

    import base64import matplotlib.pyplot as pltimport matplotlib.patches as patchesfrom PIL import Imageimport ioimport numpy as np# Test registered model locally with base64 input (serving format)print("=" * 60)print("LOCAL MODEL TEST - BASE64 INPUT")print("=" * 60)client = MlflowClient()# Ensure we have the registered model name (use widget parameter)if 'registered_model_name' not in dir():    registered_model_name = f"{catalog_name}.{schema_name}.{model_name}"print(f"\nTesting model: {registered_model_name}\n")try:    # Get latest model version    model_versions = client.search_model_versions(f"name='{registered_model_name}'")    if model_versions:        latest_version = model_versions[0].version        print(f"[OK] Found model version: {latest_version}")        print(f"   Status: {model_versions[0].status}")        # Load model using pyfunc (this is what serving endpoint uses)        model_uri = f"models:/{registered_model_name}/{latest_version}"        serving_model = mlflow.pyfunc.load_model(model_uri)        print(f"[OK] MLflow pyfunc model loaded successfully\n")        # Get test images (skip first 3 used in test_samples evaluation)        test_images = glob.glob(f"{project_location}data/coco128/images/test/*.jpg")        if test_images:            # Use images 10-12 (different from test_samples which uses 1-3)            num_samples = min(3, len(test_images) - 3)            sample_images = test_images[10:10+num_samples]  # Skip first 10 test images            print(f"Testing with {num_samples} sample images (different from test_samples)\n")            # Create color map for different classes            colors = plt.cm.tab20(np.linspace(0, 1, 20))  # 20 distinct colors            # Create visualization            fig, axes = plt.subplots(1, num_samples, figsize=(15, 5))            if num_samples == 1:                axes = [axes]            for i, test_image_path in enumerate(sample_images):                print(f"Sample {i+1}/{num_samples}: {test_image_path.split('/')[-1]}")                # Encode image as base64                with open(test_image_path, 'rb') as f:                    image_bytes = f.read()                image_base64 = base64.b64encode(image_bytes).decode('utf-8')                # Test pyfunc wrapper with base64 input                input_df = pd.DataFrame({"image_base64": [image_base64]})                predictions = serving_model.predict(input_df)                # Load and display image                image = Image.open(test_image_path)                axes[i].imshow(image)                axes[i].axis('off')                # Draw bounding boxes from pyfunc predictions                if len(predictions) > 0:                    num_detections = len(predictions)                    # Draw each bounding box with class-specific color                    for idx, row in predictions.iterrows():                        # Use xyxy coordinates                        x1, y1, x2, y2 = row['bbox_x1'], row['bbox_y1'], row['bbox_x2'], row['bbox_y2']                        width = x2 - x1                        height = y2 - y1                        # Get color based on class number                        color = colors[int(row['class_num']) % len(colors)]                        # Draw rectangle                        rect = patches.Rectangle(                            (x1, y1), width, height,                            linewidth=2, edgecolor=color, facecolor='none'                        )                        axes[i].add_patch(rect)                        # Add label with matching color                        label = f"{row['class_name']} {row['confidence']:.2f}"                        axes[i].text(                            x1, y1 - 5, label,                            color='white', fontsize=8,                            bbox=dict(facecolor=color, alpha=0.8, pad=2)                        )                    axes[i].set_title(f"{test_image_path.split('/')[-1]}\n{num_detections} objects", fontsize=10)                    print(f"   [OK] Detections: {num_detections} objects")                    for idx, row in predictions.head(3).iterrows():                        print(f"      - {row['class_name']}: {row['confidence']:.3f}")                else:                    axes[i].set_title(f"{test_image_path.split('/')[-1]}\nNo objects", fontsize=10)                    print(f"   [OK] No objects detected")                print()            plt.tight_layout()            plt.suptitle(f"Local Serving Test - MLflow PyFunc with Base64 (v{latest_version})", fontsize=14, y=1.02)            plt.show()        print("=" * 60)        print("[OK] MODEL READY FOR DEPLOYMENT")        print("=" * 60)        print(f"\n[Test Summary]")        print(f"   - Model: {registered_model_name} (v{latest_version})")        print(f"   - Input format: Base64-encoded images ✓")        print(f"   - MLflow pyfunc wrapper: ✓")        print(f"   - Bounding boxes: ✓ (color-coded by class)")        print(f"   - Test images: Different from test_samples evaluation")        print(f"   - Status: Validated and ready")        print(f"\n   Next: Deploy to serving endpoint")    else:        print(f"[ERROR] No versions found for: {registered_model_name}")        print(f"\nPlease run training cell to register the model first.")except Exception as e:    print(f"[ERROR] {e}")    import traceback    traceback.print_exc()

## Deployment checkpoint[​](#deployment-checkpoint "Direct link to Deployment checkpoint")

Manual gate before creating/updating the serving endpoint. Set the `Proceed with Deployment` widget to `true` to continue.

Python

    # ============================================================# DEPLOYMENT CHECKPOINT# ============================================================# This cell acts as a safety gate before deployment cells.# Set the 'Proceed with Deployment' widget to 'true' to continue.# ============================================================# Get deployment approval from widget (set in cell 11)PROCEED_WITH_DEPLOYMENT = dbutils.widgets.get("proceed_with_deployment") == "true"if not PROCEED_WITH_DEPLOYMENT:    message = """============================================================⚠️  DEPLOYMENT PAUSED - MANUAL CONFIRMATION REQUIRED============================================================This checkpoint prevents accidental execution of deployment cells.[To Proceed]   1. Review the model validation results above   2. Verify the model is ready for deployment   3. Set 'Proceed with Deployment' widget to 'true' (top of notebook)   4. Re-run this cell[What Happens Next]   - Cell 37: Create/update serving endpoint (AI Gateway enabled automatically)   - Cell 39: Test deployed endpoint[Safety Note]   This checkpoint ensures you don't accidentally deploy   an unvalidated model or overwrite a production endpoint.[For 'Run All']   Deployment cells will skip execution if not approved.   No errors will be raised.============================================================⏸️  DEPLOYMENT PAUSED - AWAITING APPROVAL============================================================"""    dbutils.notebook.exit(message)else:    message = """============================================================✓ DEPLOYMENT CHECKPOINT PASSED============================================================[Confirmation]   User has manually approved deployment   Execution will stop here for manual control[Next Steps - Run Manually]   1. Run cell 37: Create/update serving endpoint (AI Gateway enabled automatically)   2. Wait for endpoint to be ready (10-20 minutes)   3. Run cell 39: Test deployed endpoint[Why Manual Execution?]   - Endpoint provisioning takes 10-20 minutes   - You can monitor progress in the UI   - Each step requires verification before proceeding   - Prevents accidental 'Run All' through deployment============================================================⏸️  STOPPING HERE - RUN DEPLOYMENT CELLS MANUALLY============================================================"""    dbutils.notebook.exit(message)

## Model deployment[​](#model-deployment "Direct link to Model deployment")

Deploy to Model Serving endpoint with AI Gateway and inference table logging.

Python

    # Check deployment approval from cell 35if 'PROCEED_WITH_DEPLOYMENT' not in dir() or not PROCEED_WITH_DEPLOYMENT:    print("\n" + "="*60)    print("⚠️  DEPLOYMENT SKIPPED - NOT APPROVED")    print("="*60)    print("\n[Reason]")    print("   PROCEED_WITH_DEPLOYMENT flag is not set to True")    print("\n[To Enable Deployment]")    print("   1. Set 'Proceed with Deployment' widget to 'true' (top of notebook)")    print("   2. Re-run cell 35 (checkpoint) and this cell")    print("\n" + "="*60)else:    # Deployment approved - proceed with endpoint creation    from databricks.sdk import WorkspaceClient    from databricks.sdk.service.serving import (        ServedEntityInput,        EndpointCoreConfigInput,        AiGatewayConfig,        AiGatewayInferenceTableConfig    )    from mlflow.tracking import MlflowClient    import time    w = WorkspaceClient()    client = MlflowClient()    # Get latest model version (use widget parameter)    if 'registered_model_name' not in dir():        registered_model_name = f"{catalog_name}.{schema_name}.{model_name}"    model_versions = client.search_model_versions(        f"name='{registered_model_name}'"    )    if not model_versions:        raise ValueError(f"No model versions found for {registered_model_name}. Run training cell to register the model.")    # Unity Catalog does not support order_by; sort versions in Python instead    model_versions = sorted(model_versions, key=lambda mv: int(mv.version), reverse=True)    model_version = model_versions[0].version  # latest version (sorted DESC)    # Derive endpoint name from model name    model_name_only = registered_model_name.split('.')[-1]    endpoint_name = f"{model_name_only}_endpoint"    print("=" * 60)    print("CREATING MODEL SERVING ENDPOINT")    print("=" * 60)    print(f"\nEndpoint configuration:")    print(f"   Name: {endpoint_name}")    print(f"   Model: {registered_model_name}")    print(f"   Version: {model_version}")    print(f"   Workload size: Small")    print(f"   Scale to zero: Enabled")    print(f"   AI Gateway: Enabled with inference tables")    print(f"   Inference table: {catalog_name}.{schema_name}.{endpoint_name}_payload")    print()    try:        # Check if endpoint already exists        endpoint_exists = False        needs_update = True        needs_ai_gateway_update = False        try:            existing_endpoint = w.serving_endpoints.get(endpoint_name)            endpoint_exists = True            print(f"[INFO] Endpoint '{endpoint_name}' already exists")            # Check if endpoint is currently being updated            if existing_endpoint.state.config_update.value != "NOT_UPDATING":                print(f"[INFO] Endpoint is currently being updated (status: {existing_endpoint.state.config_update.value})")                print(f"   Checking status briefly (will timeout after 2 minutes)...\n")                # Brief check for current update status                max_wait_time = 120  # Only wait 2 minutes                poll_interval = 10                elapsed_time = 0                while elapsed_time < max_wait_time:                    endpoint = w.serving_endpoints.get(endpoint_name)                    if endpoint.state.config_update.value == "NOT_UPDATING":                        print(f"\n[OK] Current update completed (took {elapsed_time}s)")                        existing_endpoint = endpoint                        break                    elif endpoint.state.config_update.value == "UPDATE_FAILED":                        print(f"\n[WARNING] Current update failed")                        existing_endpoint = endpoint                        break                    else:                        if elapsed_time % 30 == 0:                            print(f"   Status: {endpoint.state.config_update.value} ({elapsed_time}s elapsed)")                        time.sleep(poll_interval)                        elapsed_time += poll_interval                if elapsed_time >= max_wait_time:                    print(f"\n[INFO] Endpoint update still in progress after {max_wait_time}s")                    print(f"   This cell will complete now to avoid blocking")                    print(f"\n[NEXT STEP]")                    print(f"   1. Wait a few minutes for the update to complete")                    print(f"   2. Re-run this cell to check status")                    print(f"   3. Once ready, proceed to testing")                    # Get final status and exit                    endpoint = w.serving_endpoints.get(endpoint_name)                    workspace_url = spark.conf.get("spark.databricks.workspaceUrl")                    endpoint_url = f"https://{workspace_url}/ml/endpoints/{endpoint_name}"                    print(f"\n[View Endpoint]")                    print(f"   {endpoint_url}")                    print(f"\n[Current Status]")                    print(f"   - Config State: {endpoint.state.config_update.value}")                    print(f"   - Ready State: {endpoint.state.ready.value}")                    # Exit early - don't proceed to update logic                    raise SystemExit("Endpoint update in progress - cell completed to avoid blocking")            # Check if it's already serving the same model version            current_config = existing_endpoint.config            if current_config and current_config.served_entities:                current_entity = current_config.served_entities[0]                current_model = current_entity.entity_name                current_version = current_entity.entity_version                if current_model == registered_model_name and current_version == str(model_version):                    print(f"   Already serving {registered_model_name} version {model_version}")                    print(f"   No model update needed")                    needs_update = False                    # Check if AI Gateway inference tables are enabled                    ai_gateway = existing_endpoint.ai_gateway                    if ai_gateway and ai_gateway.inference_table_config and ai_gateway.inference_table_config.enabled:                        print(f"   AI Gateway inference tables already enabled")                        print(f"   No AI Gateway update needed\n")                    else:                        print(f"   AI Gateway inference tables not enabled")                        print(f"   Will enable AI Gateway\n")                        needs_ai_gateway_update = True                else:                    print(f"   Currently serving: {current_model} v{current_version}")                    print(f"   Updating to: {registered_model_name} v{model_version}")                    print(f"   Note: AI Gateway must be configured separately for updates\n")                    needs_ai_gateway_update = True        except Exception as e:            if "does not exist" in str(e).lower() or "RESOURCE_DOES_NOT_EXIST" in str(e):                print(f"[INFO] Endpoint '{endpoint_name}' does not exist")                print(f"   Creating new endpoint with AI Gateway enabled...\n")            else:                raise        # Create/update endpoint if needed        if needs_update:            if endpoint_exists:                # Update existing endpoint using SDK method                # Note: update_config() doesn't support ai_gateway parameter                w.serving_endpoints.update_config(                    name=endpoint_name,                    served_entities=[                        ServedEntityInput(                            entity_name=registered_model_name,                            entity_version=str(model_version),                            workload_size="Small",                            scale_to_zero_enabled=True                        )                    ]                )                print(f"[OK] Endpoint update submitted")            else:                # Create new endpoint with AI Gateway enabled using SDK method                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                        )                    )                )                print(f"[OK] Endpoint creation submitted (with AI Gateway enabled)")            # Brief initial wait with shorter timeout to avoid stuck state            print(f"\n[INFO] Checking initial status (endpoint provisioning may take 10-20+ minutes)...")            max_wait_time = 120  # Only wait 2 minutes here            poll_interval = 10   # Check every 10 seconds            elapsed_time = 0            while elapsed_time < max_wait_time:                endpoint = w.serving_endpoints.get(endpoint_name)                if endpoint.state.config_update.value == "NOT_UPDATING" and endpoint.state.ready.value == "READY":                    print(f"\n[OK] Endpoint is ready! (took {elapsed_time}s)")                    break                elif endpoint.state.config_update.value == "UPDATE_FAILED":                    print(f"\n[ERROR] Endpoint update failed!")                    print(f"   Check the endpoint UI for error details")                    break                else:                    if elapsed_time % 30 == 0:  # Print status every 30 seconds                        print(f"   Status: {endpoint.state.config_update.value} ({elapsed_time}s elapsed)")                    time.sleep(poll_interval)                    elapsed_time += poll_interval            if elapsed_time >= max_wait_time:                print(f"\n[INFO] Endpoint is still initializing (this may take several more minutes)")                print(f"   This cell will complete now to avoid blocking")                print(f"\n[NEXT STEP]")                print(f"   1. Wait for endpoint to finish provisioning (check UI)")                print(f"   2. Re-run this cell to verify status")                if needs_ai_gateway_update:                    print(f"   3. Once ready, run 'Enable AI Gateway Inference Tables' cell")                else:                    print(f"   3. Once ready, proceed to testing")        # Enable AI Gateway if needed (for existing endpoints that were updated)        if needs_ai_gateway_update and endpoint_exists:            print(f"\n[INFO] Enabling AI Gateway inference tables...")            # First verify endpoint is ready            endpoint = w.serving_endpoints.get(endpoint_name)            if endpoint.state.ready.value != "READY":                print(f"[WARNING] Endpoint not ready yet (status: {endpoint.state.ready.value})")                print(f"   Run 'Enable AI Gateway Inference Tables' cell once endpoint is ready")            else:                # Enable AI Gateway (table will be created automatically by AI Gateway)                w.serving_endpoints.put_ai_gateway(                    name=endpoint_name,                    inference_table_config=AiGatewayInferenceTableConfig(                        catalog_name=catalog_name,                        schema_name=schema_name,                        table_name_prefix=endpoint_name,                        enabled=True                    )                )                print(f"[OK] AI Gateway configuration submitted")                # Brief wait for configuration                time.sleep(5)                max_wait = 60                elapsed = 0                while elapsed < max_wait:                    endpoint = w.serving_endpoints.get(endpoint_name)                    if endpoint.state.config_update.value == "NOT_UPDATING":                        break                    time.sleep(5)                    elapsed += 5        # Get final status        endpoint = w.serving_endpoints.get(endpoint_name)        # Get workspace URL for link        workspace_url = spark.conf.get("spark.databricks.workspaceUrl")        endpoint_url = f"https://{workspace_url}/ml/endpoints/{endpoint_name}"        print("\n" + "=" * 60)        if endpoint.state.config_update.value == "NOT_UPDATING" and endpoint.state.ready.value == "READY":            print("[OK] SERVING ENDPOINT READY")        else:            print("[INFO] SERVING ENDPOINT INITIALIZING")        print("=" * 60)        print(f"\n[Endpoint Details]")        print(f"   - Name: {endpoint_name}")        print(f"   - Model: {registered_model_name} (v{model_version})")        print(f"   - Config State: {endpoint.state.config_update.value}")        print(f"   - Ready State: {endpoint.state.ready.value}")        # Check AI Gateway status        if endpoint.ai_gateway and endpoint.ai_gateway.inference_table_config:            ai_config = endpoint.ai_gateway.inference_table_config            if ai_config.enabled:                print(f"   - AI Gateway: Enabled")                print(f"   - Inference Table: {ai_config.catalog_name}.{ai_config.schema_name}.{ai_config.table_name_prefix}_payload")            else:                print(f"   - AI Gateway: Disabled")        else:            print(f"   - AI Gateway: Not configured")        print(f"\n[View Endpoint]")        print(f"   {endpoint_url}")        if endpoint.state.config_update.value == "NOT_UPDATING" and endpoint.state.ready.value == "READY":            if endpoint.ai_gateway and endpoint.ai_gateway.inference_table_config and endpoint.ai_gateway.inference_table_config.enabled:                print(f"\n[Next Step]")                print(f"   Run 'Test Deployed Endpoint' cell")            else:                print(f"\n[Next Step]")                print(f"   Run 'Enable AI Gateway Inference Tables' cell")        else:            print(f"\n[Next Step]")            print(f"   Wait for endpoint to be ready, then re-run this cell")    except SystemExit as e:        # Clean exit when endpoint is still updating        print(f"\n[INFO] Cell completed (endpoint update in progress)")    except Exception as e:        print(f"[ERROR] Failed to create/update endpoint: {e}")        import traceback        traceback.print_exc()

## AI Gateway inference tables[​](#ai-gateway-inference-tables "Direct link to AI Gateway inference tables")

**AI Gateway is configured automatically** when creating new endpoints. For existing endpoints being updated, AI Gateway is enabled separately after the model update.

**Key behaviors:**

*   **Table Creation**: AI Gateway creates the inference table automatically AFTER the first request is made to the endpoint, not when AI Gateway is configured. The table structure is created immediately, but remains empty until requests are logged.
*   **Logging Delay**: There is typically a delay (usually 2-5 minutes) between when an inference request is made and when the request/response data appears in the payload table. This is normal behavior; the data is being processed and written asynchronously.
*   **Verification**: After running the test endpoint cell below, wait a few minutes then query the table to see logged requests:
    
    SQL
    
        SELECT * FROM `main`.`default`.`yolo11n_coco128_sgc_endpoint_payload`ORDER BY timestamp_ms DESC LIMIT 10
    

Python

    # Check deployment approvalif 'PROCEED_WITH_DEPLOYMENT' not in dir() or not PROCEED_WITH_DEPLOYMENT:    print("\n" + "="*60)    print("⚠️  DEPLOYMENT SKIPPED - NOT APPROVED")    print("="*60)    print("\n[Reason]")    print("   PROCEED_WITH_DEPLOYMENT flag is not set to True")    print("\n[To Enable Deployment]")    print("   1. Go to cell 35")    print("   2. Set PROCEED_WITH_DEPLOYMENT = True")    print("   3. Re-run cell 34 and deployment cells")    print("\n" + "="*60)else:    # Deployment approved - proceed with endpoint testing    import json    import base64    import glob    import random    print("=" * 60)    print("TESTING DEPLOYED ENDPOINT")    print("=" * 60)    print(f"\nEndpoint: {endpoint_name}\n")    try:        test_images = glob.glob(f"{project_location}data/coco128/images/test/*.jpg")        if test_images:            test_image_path = random.choice(test_images)            print(f"Test image: {os.path.basename(test_image_path)} (randomly selected from {len(test_images)} test images)")            # Encode image as base64            print(f"[INFO] Encoding image as base64...")            with open(test_image_path, 'rb') as f:                image_bytes = f.read()            image_base64 = base64.b64encode(image_bytes).decode('utf-8')            # Test endpoint with base64 input            print(f"[INFO] Testing endpoint with base64 input...\n")            input_data = {"dataframe_records": [{"image_base64": image_base64}]}            response = w.serving_endpoints.query(                name=endpoint_name,                dataframe_records=input_data["dataframe_records"]            )            print(f"[OK] Endpoint test successful!\n")            print(f"Response preview:")            response_dict = response.as_dict()            print(json.dumps(response_dict, indent=2)[:500])            print("\n" + "=" * 60)            print("[OK] DEPLOYMENT COMPLETE")            print("=" * 60)            print(f"\n[Deployment Summary]")            print(f"   - Endpoint: {endpoint_name}")            print(f"   - Model: {registered_model_name} (v{model_version})")            print(f"   - Status: Ready and tested")            print(f"   - AI Gateway: Enabled")            print(f"   - Input format: Base64-encoded images")            print(f"   - Inference table: {catalog_name}.{schema_name}.{endpoint_name}_payload")            # Get workspace URL for links            workspace_url = spark.conf.get("spark.databricks.workspaceUrl")            endpoint_url = f"https://{workspace_url}/ml/endpoints/{endpoint_name}"            table_url = f"https://{workspace_url}/explore/data/{catalog_name}/{schema_name}/{endpoint_name}_payload"            print(f"\n[Links]")            print(f"   - Endpoint: {endpoint_url}")            print(f"   - Inference table: {table_url}")            print(f"\n[Usage Example]")            print(f"   import base64")            print(f"   with open('image.jpg', 'rb') as f:")            print(f"       img_b64 = base64.b64encode(f.read()).decode('utf-8')")            print(f"   ")            print(f"   w.serving_endpoints.query(")            print(f"       name='{endpoint_name}',")            print(f"       dataframe_records=[{{'image_base64': img_b64}}]")            print(f"   )")            print(f"\n[Monitor Inference]")            print(f"   SELECT * FROM {catalog_name}.{schema_name}.{endpoint_name}_payload")            print(f"   ORDER BY timestamp_ms DESC LIMIT 10")        else:            print("[WARNING] No test images found")    except Exception as e:        print(f"[ERROR] Endpoint test failed: {e}")        import traceback        traceback.print_exc()        print(f"\n[INFO] Verify endpoint is ready and AI Gateway is configured")

Python

    # ============================================================# VISUALIZE ENDPOINT PREDICTIONS# ============================================================# Draws bounding boxes returned by the serving endpoint onto# the test image used in the previous cell.import matplotlib.pyplot as pltimport matplotlib.patches as patchesfrom PIL import Imageimport numpy as npif 'response_dict' not in dir() or 'test_image_path' not in dir():    print("[SKIP] Run the 'Test Deployed Endpoint' cell first.")else:    img = Image.open(test_image_path)    fig, (ax_orig, ax) = plt.subplots(1, 2, figsize=(20, 9))    # Left panel: original image    ax_orig.imshow(img)    ax_orig.set_title("Original", fontsize=13)    ax_orig.axis("off")    # Right panel: image with endpoint detections drawn    ax.imshow(img)    # Parse predictions from the endpoint response    predictions = response_dict.get("predictions", [])    if isinstance(predictions, list) and len(predictions) > 0:        # predictions is a list of dicts (one per detection)        # or a list with one element that is a list of dicts        if isinstance(predictions[0], dict) and "class_name" in predictions[0]:            detections = predictions        elif isinstance(predictions[0], list):            detections = predictions[0]        else:            detections = predictions    else:        detections = []    # Color map for distinct classes    unique_classes = list({d.get("class_name", "?") for d in detections})    cmap = plt.get_cmap("tab10", max(len(unique_classes), 1))    class_colors = {cls: cmap(i) for i, cls in enumerate(unique_classes)}    for det in detections:        x1 = det.get("bbox_x1", 0)        y1 = det.get("bbox_y1", 0)        x2 = det.get("bbox_x2", 0)        y2 = det.get("bbox_y2", 0)        cls_name = det.get("class_name", "?")        conf = det.get("confidence", 0)        color = class_colors.get(cls_name, "lime")        rect = patches.Rectangle(            (x1, y1), x2 - x1, y2 - y1,            linewidth=2, edgecolor=color, facecolor="none"        )        ax.add_patch(rect)        ax.text(            x1, y1 - 4,            f"{cls_name} {conf:.2f}",            fontsize=9, color="white",            bbox=dict(facecolor=color, alpha=0.7, pad=1, edgecolor="none")        )    ax.set_title(f"Endpoint Detections — {len(detections)} objects", fontsize=13)    ax.axis("off")    plt.tight_layout()    plt.show()    print(f"\n[INFO] {len(detections)} detections drawn from endpoint response")    print(f"   Classes: {', '.join(sorted(unique_classes))}")

Python

    # ============================================================# DETECTION SUMMARY TABLE# ============================================================# Tabular summary of the endpoint detections drawn in the cell above.import pandas as pdif 'detections' not in dir():    print("[SKIP] Run the 'Visualize Endpoint Predictions' cell first.")elif len(detections) == 0:    print("[INFO] No detections returned by the endpoint for this image.")else:    det_df = pd.DataFrame(detections)    # Per-class summary: count + confidence stats    summary_df = (        det_df.groupby("class_name")        .agg(            count=("confidence", "size"),            avg_confidence=("confidence", "mean"),            min_confidence=("confidence", "min"),            max_confidence=("confidence", "max"),        )        .round(3)        .reset_index()        .sort_values("count", ascending=False)        .reset_index(drop=True)    )    # Per-detection detail, sorted by confidence    detail_cols = [c for c in ["class_name", "confidence", "bbox_x1", "bbox_y1", "bbox_x2", "bbox_y2"] if c in det_df.columns]    detail_df = det_df[detail_cols].copy()    for c in detail_df.columns:        if c != "class_name":            detail_df[c] = detail_df[c].round(2)    detail_df = detail_df.sort_values("confidence", ascending=False).reset_index(drop=True)    print(f"[Per-class summary] {len(detections)} detections across {summary_df.shape[0]} class(es)")    print(summary_df.to_string(index=False))    print("\n[Per-detection detail]")    print(detail_df.to_string(index=False))

## Next steps[​](#next-steps "Direct link to Next steps")

*   Learn more about [Databricks AI Runtime](https://docs.databricks.com/aws/en/machine-learning/ai-runtime/) and the [AI Runtime environment versions](https://docs.databricks.com/aws/en/release-notes/serverless/environment-version/five-gpu).
*   Explore [Model Serving](https://docs.databricks.com/aws/en/machine-learning/model-serving/) to deploy and manage your model endpoint.
*   Scale this workflow to a larger dataset (1K+ images) by updating the data paths. See [NuInsSeg](https://github.com/databricks-industry-solutions/cv-playground/tree/main/projects/NuInsSeg) for a real-world example.

## Example notebook[​](#example-notebook "Direct link to Example notebook")

#### YOLO11n object detection on Databricks AI Runtime
