pkgs/development/python-modules/mlflow/default.nix

   1 {
   2   lib,
   3   fetchFromGitHub,
   4
   5   # build-system
   6   setuptools,
   7
   8   # dependencies
   9   alembic,
  10   buildPythonPackage,
  11   cachetools,
  12   click,
  13   cloudpickle,
  14   databricks-sdk,
  15   docker,
  16   flask,
  17   gitpython,
  18   graphene,
  19   gunicorn,
  20   importlib-metadata,
  21   jinja2,
  22   markdown,
  23   matplotlib,
  24   numpy,
  25   opentelemetry-api,
  26   opentelemetry-sdk,
  27   packaging,
  28   pandas,
  29   protobuf,
  30   pyarrow,
  31   pyyaml,
  32   requests,
  33   scikit-learn,
  34   scipy,
  35   sqlalchemy,
  36   sqlparse,
  37
  38   # tests
  39   aiohttp,
  40   azure-core,
  41   azure-storage-blob,
  42   azure-storage-file,
  43   boto3,
  44   botocore,
  45   catboost,
  46   datasets,
  47   fastapi,
  48   google-cloud-storage,
  49   httpx,
  50   jwt,
  51   keras,
  52   langchain,
  53   librosa,
  54   moto,
  55   opentelemetry-exporter-otlp,
  56   optuna,
  57   pydantic,
  58   pyspark,
  59   pytestCheckHook,
  60   pytorch-lightning,
  61   sentence-transformers,
  62   starlette,
  63   statsmodels,
  64   tensorflow,
  65   torch,
  66   transformers,
  67   uvicorn,
  68   xgboost,
  69 }:
  70
  71 buildPythonPackage rec {
  72   pname = "mlflow";
  73   version = "2.17.2";
  74   pyproject = true;
  75
  76   src = fetchFromGitHub {
  77     owner = "mlflow";
  78     repo = "mlflow";
  79     rev = "refs/tags/v${version}";
  80     hash = "sha256-s3t6RAJh129d5XJKtMNxS0wgGO4mKbAfMCXDBXEKBxM=";
  81   };
  82
  83   # Remove currently broken dependency `shap`, a model explainability package.
  84   # This seems quite unprincipled especially with tests not being enabled,
  85   # but not mlflow has a 'skinny' install option which does not require `shap`.
  86   pythonRemoveDeps = [ "shap" ];
  87   pythonRelaxDeps = [
  88     "gunicorn"
  89     "importlib-metadata"
  90     "packaging"
  91     "protobuf"
  92     "pytz"
  93     "pyarrow"
  94   ];
  95
  96   build-system = [ setuptools ];
  97
  98   dependencies = [
  99     alembic
 100     cachetools
 101     click
 102     cloudpickle
 103     databricks-sdk
 104     docker
 105     flask
 106     gitpython
 107     graphene
 108     gunicorn
 109     importlib-metadata
 110     jinja2
 111     markdown
 112     matplotlib
 113     numpy
 114     opentelemetry-api
 115     opentelemetry-sdk
 116     packaging
 117     pandas
 118     protobuf
 119     pyarrow
 120     pyyaml
 121     requests
 122     scikit-learn
 123     scipy
 124     sqlalchemy
 125     sqlparse
 126   ];
 127
 128   pythonImportsCheck = [ "mlflow" ];
 129
 130   nativeCheckInputs = [
 131     aiohttp
 132     azure-core
 133     azure-storage-blob
 134     azure-storage-file
 135     boto3
 136     botocore
 137     catboost
 138     datasets
 139     fastapi
 140     google-cloud-storage
 141     httpx
 142     jwt
 143     keras
 144     langchain
 145     librosa
 146     moto
 147     opentelemetry-exporter-otlp
 148     optuna
 149     pydantic
 150     pyspark
 151     pytestCheckHook
 152     pytorch-lightning
 153     sentence-transformers
 154     starlette
 155     statsmodels
 156     tensorflow
 157     torch
 158     transformers
 159     uvicorn
 160     xgboost
 161   ];
 162
 163   disabledTestPaths = [
 164     # Requires unpackaged `autogen`
 165     "tests/autogen/test_autogen_autolog.py"
 166
 167     # Requires unpackaged `diviner`
 168     "tests/diviner/test_diviner_model_export.py"
 169
 170     # Requires unpackaged `sktime`
 171     "examples/sktime/test_sktime_model_export.py"
 172
 173     # Requires `fastai` which would cause a circular dependency
 174     "tests/fastai/test_fastai_autolog.py"
 175     "tests/fastai/test_fastai_model_export.py"
 176
 177     # Requires `spacy` which would cause a circular dependency
 178     "tests/spacy/test_spacy_model_export.py"
 179
 180     # Requires `tensorflow.keras` which is not included in our outdated version of `tensorflow` (2.13.0)
 181     "tests/gateway/providers/test_ai21labs.py"
 182     "tests/tensorflow/test_keras_model_export.py"
 183     "tests/tensorflow/test_keras_pyfunc_model_works_with_all_input_types.py"
 184     "tests/tensorflow/test_mlflow_callback.py"
 185   ];
 186
 187   # I (@GaetanLepage) gave up at enabling tests:
 188   # - They require a lot of dependencies (some unpackaged);
 189   # - Many errors occur at collection time;
 190   # - Most (all ?) tests require internet access anyway.
 191   doCheck = false;
 192
 193   meta = {
 194     description = "Open source platform for the machine learning lifecycle";
 195     mainProgram = "mlflow";
 196     homepage = "https://github.com/mlflow/mlflow";
 197     changelog = "https://github.com/mlflow/mlflow/blob/v${version}/CHANGELOG.md";
 198     license = lib.licenses.asl20;
 199     maintainers = with lib.maintainers; [ tbenst ];
 200   };
 201 }