Repository: mlflow/mlflow
Pull request: [#7921] Modify pytorch predict to use the GPU if it's available
<!-- 🚨 We recommend pull requests be filed from a non-master branch on a repository fork (e.g. <username>:fix-xxx). 🚨 -->

## Related Issues/PRs

Resolve #7921 

## What changes are proposed in this pull request?

This PR proposes a modification in the PyTorch flavor's scoring methodology to use the GPU by default. This way, the scoring server will automatically use the GPU when it is available.

In order to still allow use cases to force using the CPU even when the GPU is available, the PR also introduces the environment variable `MLFLOW_DEFAULT_PREDICTION_DEVICE`, which will be respected if device isn't explicitly passed in.

## How is this patch tested?

Manual testing

Verified on a machine with a V100 GPU available:

inference with CUDA not available -> device is CPU
inference with CUDA available -> device is GPU
inference with CUDA available and MLFLOW_DEFAULT_PREDICTION_DEVICE="cpu"  -> device is CPU

## Does this PR change the documentation?

- [x] No. You can skip the rest of this section.
- [ ] Yes. Make sure the changed pages / sections render correctly in the documentation preview.

## Release Notes

### Is this a user-facing change?

- [ ] No. You can skip the rest of this section.
- [x] Yes. Give a description of this change to be included in the release notes for MLflow users.

MLflow's PyTorch flavor now uses the GPU by default when it is available, and GPU inference should now work for tensor inputs, in addition to dataframe inputs.

### What component(s), interfaces, languages, and integrations does this PR affect?

Components

- [ ] `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

Interface

- [ ] `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

Language

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

Integrations

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations

<!--
Insert an empty named anchor here to allow jumping to this section with a fragment URL
(e.g. https://github.com/mlflow/mlflow/pull/123#user-content-release-note-category).
Note that GitHub prefixes anchor names in markdown with "user-content-".
-->

<a name="release-note-category"></a>

### How should the PR be classified in the release notes? Choose one:

- [ ] `rn/breaking-change` - The PR will be mentioned in the "Breaking Changes" section
- [ ] `rn/none` - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
- [x] `rn/feature` - A new user-facing feature worth mentioning in the release notes
- [ ] `rn/bug-fix` - A user-facing bug fix worth mentioning in the release notes
- [ ] `rn/documentation` - A user-facing documentation change worth mentioning in the release notes


Problem statement:
[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

Hints from issue discussion:

Commit message: modify pytorch predict to use the GPU if it's available

Signed-off-by: Ankit Mathur <ankit.mathur@databricks.com>

Patch to review:
diff --git a/mlflow/environment_variables.py b/mlflow/environment_variables.py
index aa7ac4017a509..c8cb7f7ce5d3f 100644
--- a/mlflow/environment_variables.py
+++ b/mlflow/environment_variables.py
@@ -194,3 +194,7 @@ def get(self):
 MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT = _EnvironmentVariable(
     "MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT", int, None
 )
+
+MLFLOW_PYTORCH_DISABLE_AUTOMATIC_GPU_USAGE = _BooleanEnvironmentVariable(
+    "MLFLOW_PYTORCH_DISABLE_AUTOMATIC_GPU_USAGE", False
+)
diff --git a/mlflow/pytorch/__init__.py b/mlflow/pytorch/__init__.py
index ed16eb0ebf83a..68f0798842c0b 100644
--- 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_PYTORCH_DISABLE_AUTOMATIC_GPU_USAGE
 from mlflow.exceptions import MlflowException
 from mlflow.ml_package_versions import _ML_PACKAGE_VERSIONS
 from mlflow.models import Model, ModelSignature
@@ -63,6 +64,7 @@
 _PICKLE_MODULE_INFO_FILE_NAME = "pickle_module_info.txt"
 _EXTRA_FILES_KEY = "extra_files"
 _REQUIREMENTS_FILE_KEY = "requirements_file"
+_DEFAULT_PYTORCH_PREDICTION_DEVICE = "cpu"
 
 _logger = logging.getLogger(__name__)
 
@@ -769,9 +771,15 @@ class _PyTorchWrapper:
     def __init__(self, pytorch_model):
         self.pytorch_model = pytorch_model
 
-    def predict(self, data, device="cpu"):
+    def predict(self, data, device=_DEFAULT_PYTORCH_PREDICTION_DEVICE):
         import torch
-
+        if not MLFLOW_PYTORCH_DISABLE_AUTOMATIC_GPU_USAGE.get() and torch.cuda.is_available():
+            # if CUDA is available, we use the default CUDA device.
+            # To force inference to the CPU when the GPU is available, please set
+            # MLFLOW_PYTORCH_DISABLE_AUTOMATIC_GPU_USAGE to false
+            # If a specific non-default device is passed in, we continue to respect that.
+            if device == _DEFAULT_PYTORCH_PREDICTION_DEVICE:
+                device = "cuda"
         if isinstance(data, pd.DataFrame):
             inp_data = data.values.astype(np.float32)
         elif isinstance(data, np.ndarray):


Code review comments:
diff_hunk: @@ -769,9 +771,15 @@ class _PyTorchWrapper:
     def __init__(self, pytorch_model):
         self.pytorch_model = pytorch_model
 
-    def predict(self, data, device="cpu"):
+    def predict(self, data, device=_DEFAULT_PYTORCH_PREDICTION_DEVICE):
         import torch
-
+        if not MLFLOW_PYTORCH_DISABLE_AUTOMATIC_GPU_USAGE.get() and torch.cuda.is_available():
+            # if CUDA is available, we use the default CUDA device.
+            # To force inference to the CPU when the GPU is available, please set
+            # MLFLOW_PYTORCH_DISABLE_AUTOMATIC_GPU_USAGE to false
+            # If a specific non-default device is passed in, we continue to respect that.
+            if device == _DEFAULT_PYTORCH_PREDICTION_DEVICE:
path: mlflow/pytorch/__init__.py
text: @user1:
If I passed in device="cpu" and I had a cuda enabled bit of hardware on my machine, this wouldn't respect my declaration, would it?

@author:
That's right - you'd have to use the env variable to disable the GPU placement.

@author:
I'm not aware of a way to distinguish though - do you know how we might do that?

@user1:
could we have the `_DEFAULT_PYTORCH_PREDICTION_DEVICE` be a placeholder that gets overridden to 'cpu' if the default isn't overridden or the flag is `True` or cuda isn't available?

@author:
That's a great point - we could do that

@user2:
I would rather we define an environment variable called `MLFLOW_DEFAULT_PREDICTION_DEVICE` with default value `None` and then set:

```
_DEFAULT_PYTORCH_PREDICTION_DEVICE = MLFLOW_DEFAULT_PREDICTION_DEVICE.get() or (
    "gpu" if torch.cuda.is_available() else "cpu"
)
```

@author:
Updated to something along these lines - please lmk what you guys think