Move psutil to test requirements and do optimistic load (mlflow#10342)

Signed-off-by: B-Step62 <[email protected]>
Signed-off-by: Yuki Watanabe <[email protected]>
Co-authored-by: Harutaka Kawamura <[email protected]>

Author: B-Step62

.github/workflows/gateway.yml

@@ -34,7 +34,7 @@ jobs:
         run: |
           pip install .[gateway] \
             pytest pytest-timeout pytest-asyncio \
-            httpx sentence-transformers transformers
+            httpx psutil sentence-transformers transformers
           # Ensure pydantic v2 is installed
           pip freeze | tail -n +1 | grep -q '^pydantic==2\.'
       - name: Run tests with pydantic v2

dev/Dockerfile.test

@@ -20,7 +20,7 @@ ADD ./requirements /app/requirements
 RUN \
     # install required python packages
     pip install -r requirements/test-requirements.txt --no-cache-dir \
-        -r requirements/lint-requirements.txt psutil
+        -r requirements/lint-requirements.txt
 
 ADD . /app
 

docker/Dockerfile

@@ -1,12 +1,3 @@
 FROM python:3.10-slim-bullseye
 ARG VERSION
-# NB: installing gcc is required in order to build a wheel file 
-# for psutil, which is a dependency of mlflow. 
-# The binaries for psutil do not support Python 3.10, requiring 
-# a build from source. psutil has a large amount of cython within 
-# its code base, necessitating gcc being present in the image.
-RUN apt-get update && \
-    apt-get install -y gcc && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists/*
 RUN pip install --no-cache mlflow==$VERSION

mlflow/__init__.py

@@ -209,7 +209,7 @@
 ]
 
 
-# `mlflow.gateway` depends on optional dependencies such as pydantic and has version
+# `mlflow.gateway` depends on optional dependencies such as pydantic, psutil, and has version
 # restrictions for dependencies. Importing this module fails if they are not installed or
 # if invalid versions of these required packages are installed.
 with contextlib.suppress(Exception):

mlflow/tracking/fluent.py

@@ -380,6 +380,18 @@ def start_run(
             # if `log_system_metrics` is not specified, we will check environment variable.
             log_system_metrics = MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING.get()
         if log_system_metrics:
+            # Ensure psutil is installed. It was moved to an optional dependency, as it doesn't
+            # have binary for Arm64 Linux and requires build from CPython which is a headache.
+            # https://github.com/giampaolo/psutil/issues/1972
+            if importlib.util.find_spec("psutil") is None:
+                raise MlflowException(
+                    "Failed to start system metrics monitoring as package `psutil` is not "
+                    "installed. Run `pip install psutil` to resolve the issue, "
+                    "otherwise you can disable system metrics logging by passing "
+                    "`log_system_metrics=False` to `start_run()` or setting environment "
+                    f"variable {MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING} to False."
+                )
+
             try:
                 from mlflow.system_metrics.system_metrics_monitor import SystemMetricsMonitor
 

requirements/core-requirements.txt

@@ -18,4 +18,3 @@ markdown<4,>=3.3
 Jinja2<4,>=2.11; platform_system != 'Windows'
 Jinja2<4,>=3.0; platform_system == 'Windows'
 matplotlib<4
-psutil<6

requirements/core-requirements.yaml

@@ -81,7 +81,3 @@ Jinja2-windows:
 matplotlib:
   pip_release: matplotlib
   max_major_version: 3
-
-psutil:
-  pip_release: psutil
-  max_major_version: 5

requirements/test-requirements.txt

@@ -42,3 +42,5 @@ databricks-sdk
 # Required for LLM eval in `mlflow.evaluate`
 # TODO: Unpin once we support openai 1.0
 openai<1.0
+# Required for showing pytest stats
+psutil