Repository: mlflow/mlflow

Issue:
[FR] PyTorch predictions should use the GPU when available
### Willingness to contribute

Yes. I can contribute this feature independently.

### Proposal Summary

When running MLflow's Pytorch predict function, if you want to use the GPU, it's required to pass in the device flag to the predict function. However, PyTorch offers functionality that lets us see if the GPU is available, so it would be nice if the predict function just used the CUDA default device when CUDA is available.

This has the added benefit that mlflow models serve (the scoring server) would use the GPU for pytorch models. Currently, it does not because it loads the model as pyfunc, which does not allow passing custom parameters to the underlying flavor's predict function, meaning that the default device of CPU is used.

### Motivation

Described in the Proposal summary

### Details

_No response_

### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [X] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:

Resolving patch:
diff --git a/mlflow/environment_variables.py b/mlflow/environment_variables.py
--- a/mlflow/environment_variables.py
+++ b/mlflow/environment_variables.py
@@ -194,3 +194,11 @@ def get(self):
 MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT = _EnvironmentVariable(
     "MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT", int, None
 )
+
+#: Specifies the device intended for use in the predict function - can be used
+#: to override behavior where the GPU is used by default when available by
+#: setting this environment variable to be ``cpu``. Currently, this
+#: variable is only supported for the MLflow PyTorch flavor.
+MLFLOW_DEFAULT_PREDICTION_DEVICE = _EnvironmentVariable(
+    "MLFLOW_DEFAULT_PREDICTION_DEVICE", str, None
+)
diff --git a/mlflow/pytorch/__init__.py b/mlflow/pytorch/__init__.py
--- a/mlflow/pytorch/__init__.py
+++ b/mlflow/pytorch/__init__.py
@@ -22,6 +22,7 @@
 import mlflow
 import shutil
 from mlflow import pyfunc
+from mlflow.environment_variables import MLFLOW_DEFAULT_PREDICTION_DEVICE
 from mlflow.exceptions import MlflowException
 from mlflow.ml_package_versions import _ML_PACKAGE_VERSIONS
 from mlflow.models import Model, ModelSignature
@@ -63,6 +64,8 @@
 _PICKLE_MODULE_INFO_FILE_NAME = "pickle_module_info.txt"
 _EXTRA_FILES_KEY = "extra_files"
 _REQUIREMENTS_FILE_KEY = "requirements_file"
+_TORCH_CPU_DEVICE_NAME = "cpu"
+_TORCH_DEFAULT_GPU_DEVICE_NAME = "cuda"
 
 _logger = logging.getLogger(__name__)
 
@@ -769,9 +772,20 @@ class _PyTorchWrapper:
     def __init__(self, pytorch_model):
         self.pytorch_model = pytorch_model
 
-    def predict(self, data, device="cpu"):
+    def predict(self, data, device=None):
         import torch
 
+        # if CUDA is available, we use the default CUDA device.
+        # To force inference to the CPU when the GPU is available, please set
+        # MLFLOW_DEFAULT_PREDICTION_DEVICE to "cpu"
+        # If a specific non-default device is passed in, we continue to respect that.
+        if device is None:
+            if MLFLOW_DEFAULT_PREDICTION_DEVICE.get():
+                device = MLFLOW_DEFAULT_PREDICTION_DEVICE.get()
+            elif torch.cuda.is_available():
+                device = _TORCH_DEFAULT_GPU_DEVICE_NAME
+            else:
+                device = _TORCH_CPU_DEVICE_NAME
         if isinstance(data, pd.DataFrame):
             inp_data = data.values.astype(np.float32)
         elif isinstance(data, np.ndarray):
@@ -789,8 +803,10 @@ def predict(self, data, device="cpu"):
         with torch.no_grad():
             input_tensor = torch.from_numpy(inp_data).to(device)
             preds = self.pytorch_model(input_tensor)
-            if device != "cpu":
-                preds = preds.to("cpu")
+            # if the predictions happened on a remote device, copy them back to
+            # the host CPU for processing
+            if device != _TORCH_CPU_DEVICE_NAME:
+                preds = preds.to(_TORCH_CPU_DEVICE_NAME)
             if not isinstance(preds, torch.Tensor):
                 raise TypeError(
                     "Expected PyTorch model to output a single output tensor, "


Test patch: