chore: import upstream snapshot with attribution

doc/source/serve/doc_code/mlflow_model_registry_integration.py

@@ -0,0 +1,139 @@
+# __train_model_start__
+from sklearn.datasets import make_regression
+from sklearn.ensemble import RandomForestRegressor
+from sklearn.metrics import mean_squared_error
+from sklearn.model_selection import train_test_split
+from sklearn.pipeline import Pipeline
+from sklearn.preprocessing import StandardScaler
+import mlflow
+import mlflow.sklearn
+import mlflow.pyfunc
+from mlflow.entities import LoggedModelStatus
+from mlflow.models import infer_signature
+import numpy as np
+
+
+def train_and_register_model():
+    # Initialize model in PENDING state
+    logged_model = mlflow.initialize_logged_model(
+        name="sk-learn-random-forest-reg-model",
+        model_type="sklearn",
+        tags={"model_type": "random_forest"},
+    )
+
+    try:
+        with mlflow.start_run() as run:
+            X, y = make_regression(n_features=4, n_informative=2, random_state=0, shuffle=False)
+            X_train, X_test, y_train, y_test = train_test_split(
+                X, y, test_size=0.2, random_state=42
+            )
+
+            params = {"max_depth": 2, "random_state": 42}
+            
+            # Best Practice: Use sklearn Pipeline to persist preprocessing
+            # This ensures training and serving transformations stay aligned
+            pipeline = Pipeline([
+                ("scaler", StandardScaler()),
+                ("regressor", RandomForestRegressor(**params))
+            ])
+            pipeline.fit(X_train, y_train)
+
+            # Log parameters and metrics
+            mlflow.log_params(params)
+
+            y_pred = pipeline.predict(X_test)
+            mlflow.log_metrics({"mse": mean_squared_error(y_test, y_pred)})
+
+            # Best Practice: Infer model signature for input validation
+            # Prevents silent failures from mismatched feature order or missing columns
+            signature = infer_signature(X_train, y_pred)
+
+            # Best Practice: Pin dependency versions explicitly
+            # Ensures identical behavior across training, evaluation, and serving
+            pip_requirements = [
+                f"scikit-learn=={__import__('sklearn').__version__}",
+                f"numpy=={np.__version__}",
+            ]
+
+            # Log the sklearn pipeline with signature and dependencies
+            mlflow.sklearn.log_model(
+                sk_model=pipeline,
+                name="sklearn-model",
+                input_example=X_train[:1],
+                signature=signature,
+                pip_requirements=pip_requirements,
+                registered_model_name="sk-learn-random-forest-reg-model",
+                model_id=logged_model.model_id,
+            )
+
+        # Finalize model as READY
+        mlflow.finalize_logged_model(logged_model.model_id, LoggedModelStatus.READY)
+
+        mlflow.set_logged_model_tags(
+            logged_model.model_id,
+            tags={"production": "true"},
+        )
+
+    except Exception as e:
+        # Mark model as FAILED if issues occur
+        mlflow.finalize_logged_model(logged_model.model_id, LoggedModelStatus.FAILED)
+        raise
+
+    # Retrieve and work with the logged model
+    final_model = mlflow.get_logged_model(logged_model.model_id)
+    print(f"Model {final_model.name} is {final_model.status}")
+# __train_model_end__
+
+
+# __deployment_start__
+from ray import serve
+import mlflow.pyfunc
+import numpy as np
+
+
+@serve.deployment
+class MLflowModelDeployment:
+    def __init__(self):
+        # Search for models with production tag
+        models = mlflow.search_logged_models(
+            filter_string="tags.production='true' AND name='sk-learn-random-forest-reg-model'",
+            order_by=[{"field_name": "creation_time", "ascending": False}],
+        )
+        if models.empty:
+            raise ValueError("No model with production tag found")
+        
+        # Get the most recent production model
+        model_row = models.iloc[0]
+        artifact_location = model_row["artifact_location"]
+        
+        # Best Practice: Load model once during initialization (warm-start)
+        # This eliminates first-request latency spikes
+        self.model = mlflow.pyfunc.load_model(artifact_location)
+        
+        # Pre-warm the model with a dummy prediction
+        dummy_input = np.zeros((1, 4))
+        _ = self.model.predict(dummy_input)
+
+    async def __call__(self, request):
+        data = await request.json()
+        features = np.array(data["features"])
+        
+        # MLflow validates input against the logged signature automatically
+        prediction = self.model.predict(features)
+        return {"prediction": prediction.tolist()}
+
+
+app = MLflowModelDeployment.bind()
+# __deployment_end__
+
+
+if __name__ == "__main__":
+    import requests
+    from ray import serve
+
+    train_and_register_model()
+    serve.run(app)
+
+    # Test prediction
+    response = requests.post("http://localhost:8000/", json={"features": [[0.1, 0.2, 0.3, 0.4]]})
+    print(response.json())