diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 2304c0e5..221ff40d 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -32,10 +32,13 @@ jobs: uses: actions/setup-python@v4 with: python-version: ${{ matrix.python-version }} + - name: Install OpenMP runtime for XGBoost (macOS) + if: runner.os == 'macOS' + run: brew install libomp - name: Installation run: | python -m pip install --upgrade pip - pip install .["streaming"] + pip install ".[streaming,xgboost]" - name: Test requirements Installation run: | python otherfiles/requirements-splitter.py diff --git a/.gitignore b/.gitignore index bf47b471..356777b7 100644 --- a/.gitignore +++ b/.gitignore @@ -44,9 +44,13 @@ htmlcov/ .cache nosetests.xml coverage.xml +coverage_*.xml *,cover .hypothesis/ +# Generated by otherfiles/requirements-splitter.py from dev-requirements.txt (CI only) +test-requirements.txt + # Translations *.mo *.pot @@ -82,9 +86,13 @@ celerybeat-schedule # virtualenv .venv/ +.venv*/ venv/ ENV/ +# Editor / IDE +.vscode/ + # Spyder project settings .spyderproject diff --git a/CHANGELOG.md b/CHANGELOG.md index fdadfe54..45e1af34 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,11 +7,42 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0. ## [Unreleased] ### Added - `.codacy.yaml` +- `XGBoostChain` class in `pymilo/chains/xgboost_chain.py` +- `BoosterTransporter` class in `pymilo/transporters/booster_transporter.py` +- `XGBOOST_SUPPORTED_TABLE`, `XGBOOST_SUPPORTED_CATEGORIES` and `CLI_XGBOOST_NOT_INSTALLED` in `pymilo/pymilo_param.py` +- `get_xgboost_class` and `get_ml_model_class` functions in `pymilo/utils/util.py` +- `get_xgboost_version` and `get_framework_versions` functions in `pymilo/pymilo_func.py` +- `get_ml_model_data` function in `pymilo/pymilo_func.py` +- `to_ml_model` function in `pymilo/pymilo_func.py` +- `pymilo_xgb_booster_test` function in `pymilo/utils/test_pymilo.py` +- `xgboost-requirements.txt` +- XGBoost models test runner +- `scenario5` XGBoost test for streaming +- `PYMILO_CLIENT_DEFAULT_CLIENT_ID` and `PYMILO_CLIENT_DEFAULT_ML_MODEL_ID` in `pymilo/streaming/param.py` +- `MANIFEST.in` ### Changed - `get_requires` function in `setup.py` - `meta.yaml` - `dev-requirements.txt` - `requirements.txt` +- `Export` class in `pymilo/pymilo_obj.py` +- `Import` class in `pymilo/pymilo_obj.py` +- `print_supported_ml_models` function in `pymilo/pymilo_func.py` +- `pymilo/__main__.py` +- `.github/workflows/test.yml` +- `scenario3` in `tests/test_ml_streaming/scenarios/scenario3.py` +- `serialize_spline` function in `PreprocessingTransporter` +- `deserialize_spline` function in `PreprocessingTransporter` +- `SUPPORTED_MODELS.md` updated +- `README.md` updated +### Fixed +- PyMilo CLI server launch +- `generate_dockerfile` `--init` CMD +- Source distribution packaging +- `Export.save` directory creation +### Removed +- `get_sklearn_data` function in `pymilo/pymilo_func.py` +- `to_sklearn_model` function in `pymilo/pymilo_func.py` ## [1.6] - 2026-04-03 ### Added - `_check_response_error` function in `WebSocketClientCommunicator` diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 00000000..5123ac1c --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,5 @@ +# setup.py reads these at build time, so they must ship inside the sdist. +include requirements.txt +include streaming-requirements.txt +include xgboost-requirements.txt +include CHANGELOG.md diff --git a/README.md b/README.md index a3f29dc4..f310c418 100644 --- a/README.md +++ b/README.md @@ -192,6 +192,27 @@ The `Import` class facilitates importing of serialized models from JSON files, J This loaded model is exactly the same as the original trained model. +### XGBoost +PyMilo transparently exports and imports XGBoost models, just like scikit-learn ones. + +⚠️ In order to use XGBoost support, make sure you've installed the `xgboost` mode of PyMilo: + +```console +pip install pymilo[xgboost] +``` + +Both the scikit-learn-style wrappers (`XGBRegressor`, `XGBClassifier`, `XGBRanker`, `XGBRFRegressor`, `XGBRFClassifier`) and the native `xgboost.Booster` are supported. + +```python +from xgboost import XGBRegressor +from pymilo import Export, Import + +model = XGBRegressor(n_estimators=20).fit(X_train, y_train) +Export(model).save("xgb.json") + +restored = Import("xgb.json").to_model() +``` + ### ML streaming You can easily serve your ML model from a remote server using `ML streaming` feature of PyMilo. @@ -314,24 +335,7 @@ result = pymilo_client.predict(x_test) ## Supported ML models -| scikit-learn | PyTorch | -| ---------------- | ---------------- | -| Linear Models ✅ | - | -| Neural Networks ✅ | - | -| Trees ✅ | - | -| Clustering ✅ | - | -| Naïve Bayes ✅ | - | -| Support Vector Machines (SVMs) ✅ | - | -| Nearest Neighbors ✅ | - | -| Ensemble Models ✅ | - | -| Pipeline Model ✅ | - | -| Preprocessing Models ✅ | - | -| Cross Decomposition Models ✅ | - | -| Feature Extractor Models ✅ | - | -| Composite Models ✅ | - | - - -Details are available in [Supported Models](https://github.com/openscilab/pymilo/blob/main/SUPPORTED_MODELS.md). +See [Supported Models](https://github.com/openscilab/pymilo/blob/main/SUPPORTED_MODELS.md). ## Issues & bug reports diff --git a/SUPPORTED_MODELS.md b/SUPPORTED_MODELS.md index fe3eb177..7c0dac69 100644 --- a/SUPPORTED_MODELS.md +++ b/SUPPORTED_MODELS.md @@ -1,6 +1,6 @@ # Supported Models -**Last Update: 2025-12-30** +**Last Update: 2026-09-07**

Scikit-Learn

@@ -798,3 +798,47 @@ >=1.5 + +

XGBoost

+ +Available when PyMilo is installed with the optional `xgboost` extra (`pip install pymilo[xgboost]`). + +📚 XGBoost Documentation + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
IDModel NamePyMilo Version
1XGBRegressor>=1.7
2XGBClassifier>=1.7
3XGBRanker>=1.7
4XGBRFRegressor>=1.7
5XGBRFClassifier>=1.7
6Booster (native)>=1.7
diff --git a/dev-requirements.txt b/dev-requirements.txt index 510574ef..74e06bd8 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -7,6 +7,7 @@ requests==2.33.0; python_full_version >= "3.10" requests>=2.0.0; python_full_version < "3.10" websockets==15.0.1 pydantic==2.13.2 +xgboost==3.4.1 setuptools>=40.8.0 vulture>=1.0 bandit>=1.5.1 diff --git a/otherfiles/meta.yaml b/otherfiles/meta.yaml index 37cf81e8..ea7243f1 100644 --- a/otherfiles/meta.yaml +++ b/otherfiles/meta.yaml @@ -26,6 +26,7 @@ requirements: - fastapi>=0.68.0 - pydantic>=1.5.0 - websockets>=9.0 + - xgboost>=1.7.0 about: home: https://github.com/openscilab/pymilo diff --git a/pymilo/__main__.py b/pymilo/__main__.py index a2c5ce6c..165055d3 100644 --- a/pymilo/__main__.py +++ b/pymilo/__main__.py @@ -9,14 +9,17 @@ CLI_MORE_INFO, CLI_UNKNOWN_MODEL, CLI_ML_STREAMING_NOT_INSTALLED, + CLI_XGBOOST_NOT_INSTALLED, + XGBOOST_SUPPORTED_TABLE, ) from .pymilo_func import print_supported_ml_models, pymilo_help from .pymilo_obj import Import -from .utils.util import get_sklearn_class +from .utils.util import get_ml_model_class ml_streaming_support = True try: from .streaming import PymiloServer, Compression, CommunicationProtocol + from .streaming.param import PYMILO_CLIENT_DEFAULT_CLIENT_ID, PYMILO_CLIENT_DEFAULT_ML_MODEL_ID except BaseException: ml_streaming_support = False @@ -95,9 +98,12 @@ def main(): _model = _model.to_model() elif args.init: model_name = args.init - model_class = get_sklearn_class(model_name) + model_class = get_ml_model_class(model_name) if model_class is None: - print(f"{CLI_UNKNOWN_MODEL}\n{print_supported_ml_models()}") + if model_name in XGBOOST_SUPPORTED_TABLE: + print(f"{CLI_XGBOOST_NOT_INSTALLED}\n{CLI_MORE_INFO}") + else: + print(f"{CLI_UNKNOWN_MODEL}\n{print_supported_ml_models()}") return run_ps = True _model = model_class() @@ -109,12 +115,16 @@ def main(): pymilo_help() parser.print_help() else: - PymiloServer( - model=_model, + pymilo_server = PymiloServer( port=_port, compressor=_compressor, communication_protocol=_communication_protocol, - ).communicator.run() + ) + if _model is not None: + pymilo_server.init_client(PYMILO_CLIENT_DEFAULT_CLIENT_ID) + pymilo_server.init_ml_model(PYMILO_CLIENT_DEFAULT_CLIENT_ID, PYMILO_CLIENT_DEFAULT_ML_MODEL_ID) + pymilo_server.set_ml_model(PYMILO_CLIENT_DEFAULT_CLIENT_ID, PYMILO_CLIENT_DEFAULT_ML_MODEL_ID, _model) + pymilo_server.communicator.run() if __name__ == '__main__': diff --git a/pymilo/chains/util.py b/pymilo/chains/util.py index fad8ff92..8542b53a 100644 --- a/pymilo/chains/util.py +++ b/pymilo/chains/util.py @@ -11,6 +11,11 @@ from .cross_decomposition_chain import cross_decomposition_chain from ..utils.util import get_sklearn_type, check_str_in_iterable +try: + from .xgboost_chain import xgboost_chain as _xgboost_chain +except BaseException: + _xgboost_chain = None + MODEL_TYPE_TRANSPORTER = { "LINEAR_MODEL": linear_chain.transport, @@ -23,6 +28,9 @@ "CROSS_DECOMPOSITION": cross_decomposition_chain.transport, } +if _xgboost_chain is not None: + MODEL_TYPE_TRANSPORTER["XGBOOST"] = _xgboost_chain.transport + def get_concrete_transporter(model): """ @@ -37,6 +45,9 @@ def get_concrete_transporter(model): if upper_model in MODEL_TYPE_TRANSPORTER.keys(): return upper_model, MODEL_TYPE_TRANSPORTER[upper_model] + if _xgboost_chain is not None and _xgboost_chain.is_supported(model): + return "XGBOOST", _xgboost_chain.transport + if linear_chain.is_supported(model): return "LINEAR_MODEL", linear_chain.transport elif neural_network_chain.is_supported(model): diff --git a/pymilo/chains/xgboost_chain.py b/pymilo/chains/xgboost_chain.py new file mode 100644 index 00000000..16e2a603 --- /dev/null +++ b/pymilo/chains/xgboost_chain.py @@ -0,0 +1,126 @@ +# -*- coding: utf-8 -*- +"""PyMilo chain for XGBoost models.""" + +from .chain import AbstractChain +from ..transporters.transporter import Command +from ..transporters.booster_transporter import BoosterTransporter +from ..transporters.general_data_structure_transporter import GeneralDataStructureTransporter + +from ..pymilo_param import XGBOOST_SUPPORTED_TABLE, NOT_SUPPORTED +from ..utils.util import get_sklearn_type + +_xgb = None +_xgboost_available = False +try: + import xgboost as _xgb_imported + _xgb = _xgb_imported + _xgboost_available = True +except BaseException: + _xgb = None + _xgboost_available = False + + +XGBOOST_CHAIN = { + "BoosterTransporter": BoosterTransporter(), + "GeneralDataStructureTransporter": GeneralDataStructureTransporter(), +} + + +class XGBoostChain(AbstractChain): + """XGBoostChain handles export/import of XGBoost ML model objects. + + Two model shapes are supported and routed via isinstance checks (no hardcoded keys): + + 1. The scikit-learn-style wrappers (XGBRegressor, XGBClassifier, XGBRanker, + XGBRFRegressor, XGBRFClassifier). These have a regular Python __dict__ that is + walked through the transporter chain. The inner xgboost.Booster (typically stored + under the '_Booster' attribute) is handled by BoosterTransporter via + isinstance(value, xgboost.Booster); everything else (ndarrays, classes_, + evals_result_, scalar hyperparameters, etc.) is handled by + GeneralDataStructureTransporter. + + 2. The native xgboost.Booster (returned by xgboost.train). It has no usable + Python __dict__, so it is wrapped in a synthetic single-key dict and serialized + directly via BoosterTransporter. + """ + + def is_supported(self, model): + """ + Check if the given model is an XGBoost ML model supported by this chain. + + :param model: a string name of an ML model or an XGBoost object of it + :type model: any object + :return: check result as bool + """ + if not _xgboost_available: + return False + + if isinstance(model, str): + return model in self._supported_models and self._supported_models[model] != NOT_SUPPORTED + + if isinstance(model, _xgb.Booster): + return self._supported_models.get("Booster", NOT_SUPPORTED) != NOT_SUPPORTED + + try: + from xgboost.sklearn import XGBModel + except BaseException: + XGBModel = None + + if XGBModel is not None and isinstance(model, XGBModel): + model_name = get_sklearn_type(model) + return model_name in self._supported_models and self._supported_models[model_name] != NOT_SUPPORTED + + return False + + def serialize(self, model): + """ + Return the serialized data dict of the given XGBoost model. + + :param model: given XGBoost model to be serialized + :type model: xgboost sklearn-style estimator or xgboost.Booster + :return: dict + """ + if _xgboost_available and isinstance(model, _xgb.Booster): + shell = {"_booster": model} + shell["_booster"] = self._transporters["BoosterTransporter"].serialize( + shell, "_booster", "Booster") + return shell + + for transporter in self._transporters: + self._transporters[transporter].transport(model, Command.SERIALIZE) + return model.__dict__ + + def deserialize(self, serialized_model, is_inner_model=False): + """ + Return the associated XGBoost model of the previously serialized model. + + :param serialized_model: pymilo Import object or inner-model dict + :type serialized_model: obj + :param is_inner_model: whether this is being deserialized as a nested model + :type is_inner_model: boolean + :return: associated XGBoost model + """ + if is_inner_model: + data = serialized_model["data"] + model_type = serialized_model["type"] + else: + data = serialized_model.data + model_type = serialized_model.type + + if model_type == "Booster": + booster_value = data.get("_booster", None) + return self._transporters["BoosterTransporter"].deserialize( + {"_booster": booster_value}, "_booster", model_type) + + raw_model = self._supported_models[model_type]() + + for transporter in self._transporters: + self._transporters[transporter].transport( + serialized_model, Command.DESERIALIZE, is_inner_model) + + for item in data: + setattr(raw_model, item, data[item]) + return raw_model + + +xgboost_chain = XGBoostChain(XGBOOST_CHAIN, XGBOOST_SUPPORTED_TABLE) diff --git a/pymilo/pymilo_func.py b/pymilo/pymilo_func.py index d62c7335..664d5128 100644 --- a/pymilo/pymilo_func.py +++ b/pymilo/pymilo_func.py @@ -5,7 +5,17 @@ from .chains.util import get_transporter from .transporters.transporter import Command -from .pymilo_param import SKLEARN_SUPPORTED_CATEGORIES, NOT_SUPPORTED, OVERVIEW +from .pymilo_param import SKLEARN_SUPPORTED_CATEGORIES, XGBOOST_SUPPORTED_CATEGORIES, NOT_SUPPORTED, OVERVIEW + +_xgb = None +_xgboost_available = False +try: + import xgboost as _xgb_imported + _xgb = _xgb_imported + _xgboost_available = True +except BaseException: + _xgb = None + _xgboost_available = False def get_sklearn_version(): @@ -17,30 +27,65 @@ def get_sklearn_version(): return sklearn.__version__ -def get_sklearn_data(model): +def get_xgboost_version(): + """ + Return xgboost version when available. + + :return: xgboost version as a str, or None when xgboost is not installed """ - Return sklearn data by serializing given model. + if not _xgboost_available: + return None + return _xgb.__version__ + + +def get_ml_model_data(model): + """ + Return the serialized data of the given ML model. :param model: given model - :type model: any sklearn's model class - :return: sklearn data + :type model: any supported ML model class + :return: serialized model data """ _, transporter = get_transporter(model) return transporter(model, Command.SERIALIZE) -def to_sklearn_model(import_obj): +def to_ml_model(import_obj): """ - Deserialize the imported object as a sklearn model. + Deserialize the imported object as an ML model. :param import_obj: given object :type import_obj: pymilo.Import - :return: sklearn model + :return: ML model """ _, transporter = get_transporter(import_obj.type) return transporter(import_obj, Command.DESERIALIZE) +def get_framework_versions(model): + """ + Return top-level framework version fields for the given model. + + Pure sklearn models emit `sklearn_version`. Pure native xgboost.Booster models emit + `xgboost_version`. XGBoost sklearn-API wrappers emit both, since they sit on both stacks. + + :param model: given model + :type model: any supported ML model + :return: dict of framework version fields + """ + category, _ = get_transporter(model) + versions = {} + if category == "XGBOOST": + xgb_version = get_xgboost_version() + if xgb_version is not None: + versions["xgboost_version"] = xgb_version + if not (_xgboost_available and isinstance(model, _xgb.Booster)): + versions["sklearn_version"] = get_sklearn_version() + else: + versions["sklearn_version"] = get_sklearn_version() + return versions + + def compare_model_outputs(exported_output, imported_output, epsilon_error=10**(-8)): @@ -62,17 +107,20 @@ def compare_model_outputs(exported_output, if key not in imported_output: return False # TODO: throw exception total_error += np.abs(imported_output[key] - exported_output[key]) - return np.abs(total_error) < epsilon_error + return bool(np.abs(total_error) < epsilon_error) def print_supported_ml_models(): """ - Print the supported sklearn ML models categorized by type. + Print the supported ML models categorized by type. :return: None """ print("Supported Machine Learning Models:") - for category, table in SKLEARN_SUPPORTED_CATEGORIES.items(): + all_categories = {} + all_categories.update(SKLEARN_SUPPORTED_CATEGORIES) + all_categories.update(XGBOOST_SUPPORTED_CATEGORIES) + for category, table in all_categories.items(): print(f"**{category}**:") for model_name in table: if table[model_name] != NOT_SUPPORTED: diff --git a/pymilo/pymilo_obj.py b/pymilo/pymilo_obj.py index 7598b8ed..b4ca1269 100644 --- a/pymilo/pymilo_obj.py +++ b/pymilo/pymilo_obj.py @@ -8,10 +8,10 @@ from traceback import format_exc from .utils.util import get_sklearn_type, download_model from concurrent.futures import ThreadPoolExecutor, as_completed -from .pymilo_func import get_sklearn_data, get_sklearn_version, to_sklearn_model +from .pymilo_func import get_ml_model_data, get_sklearn_version, get_xgboost_version, get_framework_versions, to_ml_model from .exceptions.serialize_exception import PymiloSerializationException, SerializationErrorTypes from .exceptions.deserialize_exception import PymiloDeserializationException, DeserializationErrorTypes -from .pymilo_param import PYMILO_VERSION, UNEQUAL_PYMILO_VERSIONS, UNEQUAL_SKLEARN_VERSIONS +from .pymilo_param import PYMILO_VERSION, UNEQUAL_PYMILO_VERSIONS, UNEQUAL_SKLEARN_VERSIONS, UNEQUAL_XGBOOST_VERSIONS from .pymilo_param import INVALID_IMPORT_INIT_PARAMS, BATCH_IMPORT_INVALID_DIRECTORY @@ -28,13 +28,16 @@ def __init__(self, model): """ Initialize the Pymilo Export instance. - :param model: given model(any sklearn linear model) - :type model: any class of the sklearn's linear models + :param model: given model(any supported ML model) + :type model: any class of the supported ML models :return: an instance of the Pymilo Export class """ - self.data = get_sklearn_data(deepcopy(model)) - self.version = get_sklearn_version() + self.data = get_ml_model_data(deepcopy(model)) self.type = get_sklearn_type(model) + self.framework_versions = get_framework_versions(model) + # Keep `version` as sklearn_version for backward compatibility with callers + # that read Export.version; it is None for pure native xgboost.Booster models. + self.version = self.framework_versions.get("sklearn_version") def save(self, file_adr): """ @@ -44,6 +47,9 @@ def save(self, file_adr): :type file_adr: str :return: None """ + directory = os.path.dirname(file_adr) + if directory: + os.makedirs(directory, exist_ok=True) with open(file_adr, 'w') as fp: fp.write(self.to_json()) @@ -51,18 +57,19 @@ def to_json(self): """ Return a json-like representation of model. + Top-level keys always include `data`, `pymilo_version`, and `model_type`. + Framework-specific version fields (`sklearn_version` and/or `xgboost_version`) + are also placed at the first depth, right after `data`, which keeps the key + order of sklearn exports byte-identical to previous PyMilo releases. + :return: model's representation as str """ + payload = {"data": self.data} + payload.update(self.framework_versions) + payload["pymilo_version"] = PYMILO_VERSION + payload["model_type"] = self.type try: - return json.dumps( - { - "data": self.data, - "sklearn_version": self.version, - "pymilo_version": PYMILO_VERSION, - "model_type": self.type - }, - indent=4 - ) + return json.dumps(payload, indent=4) except Exception as e: raise PymiloSerializationException( { @@ -70,11 +77,7 @@ def to_json(self): 'error': { 'Exception': repr(e), 'Traceback': format_exc()}, - 'object': { - "data": self.data, - "sklearn_version": self.version, - "pymilo_version": PYMILO_VERSION, - "model_type": self.type}, + 'object': payload, }) @staticmethod @@ -151,10 +154,18 @@ def __init__(self, file_adr=None, json_dump=None, url=None): try: if not serialized_model_obj["pymilo_version"] == PYMILO_VERSION: warn(UNEQUAL_PYMILO_VERSIONS, category=Warning) - if not serialized_model_obj["sklearn_version"] == get_sklearn_version(): - warn(UNEQUAL_SKLEARN_VERSIONS, category=Warning) + if "sklearn_version" in serialized_model_obj: + if not serialized_model_obj["sklearn_version"] == get_sklearn_version(): + warn(UNEQUAL_SKLEARN_VERSIONS, category=Warning) + self.version = serialized_model_obj["sklearn_version"] + else: + # Pure native Booster exports omit sklearn_version; keep BC attribute present. + self.version = None + if "xgboost_version" in serialized_model_obj: + installed_xgboost = get_xgboost_version() + if installed_xgboost is not None and serialized_model_obj["xgboost_version"] != installed_xgboost: + warn(UNEQUAL_XGBOOST_VERSIONS, category=Warning) self.data = serialized_model_obj["data"] - self.version = serialized_model_obj["sklearn_version"] self.type = serialized_model_obj["model_type"] except Exception as e: json_content = None @@ -176,11 +187,11 @@ def __init__(self, file_adr=None, json_dump=None, url=None): def to_model(self): """ - Convert imported model to sklearn model. + Convert imported model to its ML model. - :return: sklearn model + :return: ML model """ - return to_sklearn_model(self) + return to_ml_model(self) @staticmethod def batch_import(file_addr, run_parallel=False): diff --git a/pymilo/pymilo_param.py b/pymilo/pymilo_param.py index 0c456768..692750ab 100644 --- a/pymilo/pymilo_param.py +++ b/pymilo/pymilo_param.py @@ -84,6 +84,13 @@ except BaseException: pass +xgboost_support = False +try: + import xgboost as _xgb + xgboost_support = True +except BaseException: + _xgb = None + OVERVIEW = """ PyMilo is an open source Python package that provides a simple, efficient, and safe way for users to export pre-trained machine learning models in a transparent way. """ @@ -92,6 +99,7 @@ PYMILO_VERSION_DOES_NOT_EXIST = "Corrupted JSON file, `pymilo_version` doesn't exist in this file." UNEQUAL_PYMILO_VERSIONS = "warning: Installed PyMilo version differs from the PyMilo version used to create the JSON file." UNEQUAL_SKLEARN_VERSIONS = "warning: Installed Scikit version differs from the Scikit version used to create the JSON file and it may prevent PyMilo from transporting seamlessly." +UNEQUAL_XGBOOST_VERSIONS = "warning: Installed XGBoost version differs from the XGBoost version used to create the JSON file and it may prevent PyMilo from transporting seamlessly." INVALID_IMPORT_INIT_PARAMS = "Invalid input parameters, you should either pass a valid file_adr or a json_dump or a url to initiate Import class." URL_REGEX = r'^(http|https)://[a-zA-Z0-9.-_]+\.[a-zA-Z]{2,}(/\S*)?$' DOWNLOAD_MODEL_FAILED = "Failed to download the JSON file, Server didn't respond." @@ -101,6 +109,9 @@ CLI_ML_STREAMING_NOT_INSTALLED = """ML Streaming is not installed. To install ML Streaming, run the following command: pip install pymilo[streaming]""" +CLI_XGBOOST_NOT_INSTALLED = """XGBoost is not installed. +To install XGBoost support, run the following command: +pip install pymilo[xgboost]""" CLI_MORE_INFO = "For more information, visit the PyMilo README at https://github.com/openscilab/pymilo" CLI_UNKNOWN_MODEL = "The provided ML model name is either invalid or unsupported." @@ -308,6 +319,7 @@ "ENSEMBLE": "exported_ensembles", "CROSS_DECOMPOSITION": "exported_cross_decomposition", "COMPOSE": "exported_composes", + "XGBOOST": "exported_xgboost", } SKLEARN_SUPPORTED_CATEGORIES = { @@ -322,3 +334,16 @@ "CROSS_DECOMPOSITION": SKLEARN_CROSS_DECOMPOSITION_TABLE, "COMPOSE": SKLEARN_COMPOSE_TABLE, } + +XGBOOST_SUPPORTED_TABLE = { + "XGBRegressor": _xgb.XGBRegressor if xgboost_support else NOT_SUPPORTED, + "XGBClassifier": _xgb.XGBClassifier if xgboost_support else NOT_SUPPORTED, + "XGBRanker": _xgb.XGBRanker if xgboost_support else NOT_SUPPORTED, + "XGBRFRegressor": _xgb.XGBRFRegressor if xgboost_support else NOT_SUPPORTED, + "XGBRFClassifier": _xgb.XGBRFClassifier if xgboost_support else NOT_SUPPORTED, + "Booster": _xgb.Booster if xgboost_support else NOT_SUPPORTED, +} + +XGBOOST_SUPPORTED_CATEGORIES = { + "XGBOOST": XGBOOST_SUPPORTED_TABLE, +} diff --git a/pymilo/streaming/param.py b/pymilo/streaming/param.py index 8b99b3b5..69f46e97 100644 --- a/pymilo/streaming/param.py +++ b/pymilo/streaming/param.py @@ -7,6 +7,11 @@ PYMILO_CLIENT_INVALID_ATTRIBUTE = "This attribute doesn't exist in either PymiloClient or the inner ML model." PYMILO_CLIENT_FAILED_TO_DOWNLOAD_REMOTE_MODEL = "PyMiloClient failed to download the remote ML model." +# Identifiers a PymiloClient uses before it registers itself; a PyMilo Server started from the +# CLI with a preloaded ML model publishes that model under them, so an unregistered client works. +PYMILO_CLIENT_DEFAULT_CLIENT_ID = "0x_client_id" +PYMILO_CLIENT_DEFAULT_ML_MODEL_ID = "0x_ml_model_id" + PYMILO_SERVER_NON_EXISTENT_ATTRIBUTE = "The requested attribute doesn't exist in this model." PYMILO_INVALID_URL = "The given URL is not valid." PYMILO_CLIENT_WEBSOCKET_NOT_CONNECTED = "WebSocket is not connected." diff --git a/pymilo/streaming/pymilo_client.py b/pymilo/streaming/pymilo_client.py index 5a9ddc66..544a6309 100644 --- a/pymilo/streaming/pymilo_client.py +++ b/pymilo/streaming/pymilo_client.py @@ -6,7 +6,8 @@ from ..pymilo_obj import Export, Import from .param import PYMILO_CLIENT_INVALID_MODE, PYMILO_CLIENT_MODEL_SYNCHED, \ PYMILO_CLIENT_LOCAL_MODEL_UPLOADED, PYMILO_CLIENT_LOCAL_MODEL_UPLOAD_FAILED, \ - PYMILO_CLIENT_INVALID_ATTRIBUTE, PYMILO_CLIENT_FAILED_TO_DOWNLOAD_REMOTE_MODEL + PYMILO_CLIENT_INVALID_ATTRIBUTE, PYMILO_CLIENT_FAILED_TO_DOWNLOAD_REMOTE_MODEL, \ + PYMILO_CLIENT_DEFAULT_CLIENT_ID, PYMILO_CLIENT_DEFAULT_ML_MODEL_ID from .communicator import CommunicationProtocol from ..transporters.general_data_structure_transporter import GeneralDataStructureTransporter @@ -44,8 +45,8 @@ def __init__( :return: an instance of the Pymilo PymiloClient class """ self.model = model - self.client_id = "0x_client_id" - self.ml_model_id = "0x_ml_model_id" + self.client_id = PYMILO_CLIENT_DEFAULT_CLIENT_ID + self.ml_model_id = PYMILO_CLIENT_DEFAULT_ML_MODEL_ID self._mode = mode self._compressor = compressor.value self._encryptor = DummyEncryptor() @@ -121,7 +122,7 @@ def deregister(self): :return: None """ self._communicator.remove_client(self.client_id) - self.client_id = "0x_client_id" + self.client_id = PYMILO_CLIENT_DEFAULT_CLIENT_ID def register_ml_model(self): """ @@ -138,7 +139,7 @@ def deregister_ml_model(self): :return: None """ self._communicator.remove_model(self.client_id, self.ml_model_id) - self.ml_model_id = "0x_ml_model_id" + self.ml_model_id = PYMILO_CLIENT_DEFAULT_ML_MODEL_ID def get_ml_models(self): """ diff --git a/pymilo/streaming/pymilo_server.py b/pymilo/streaming/pymilo_server.py index 130cd0a1..2e149fd8 100644 --- a/pymilo/streaming/pymilo_server.py +++ b/pymilo/streaming/pymilo_server.py @@ -21,8 +21,6 @@ def __init__( """ Initialize the Pymilo PymiloServer instance. - :param model: the ML model which will be streamed - :type model: any :param port: the port to which PyMiloServer listens :type port: int :param host: the url to which PyMilo Server listens diff --git a/pymilo/streaming/util.py b/pymilo/streaming/util.py index aaf7ed9e..671ff34c 100644 --- a/pymilo/streaming/util.py +++ b/pymilo/streaming/util.py @@ -65,8 +65,8 @@ def generate_dockerfile( :type protocol: str :param port: Port for the PyMilo server (default: 8000). :type port: int - :param init_model: The model that the server initialized with. - :type init_model: boolean + :param init_model: Name of the ML model to initialize the PyMilo server with. + :type init_model: str :param bare: A flag that sets if the server runs without an internal ML model. :type bare: boolean """ @@ -100,7 +100,7 @@ def generate_dockerfile( else: cmd += f", \"--load\", \"/app/model.json\"" elif init_model: - cmd += f", \"--init\" \"{init_model}\"" + cmd += f", \"--init\", \"{init_model}\"" elif bare: cmd += f", \"--bare\"" diff --git a/pymilo/transporters/booster_transporter.py b/pymilo/transporters/booster_transporter.py new file mode 100644 index 00000000..18526cfd --- /dev/null +++ b/pymilo/transporters/booster_transporter.py @@ -0,0 +1,79 @@ +# -*- coding: utf-8 -*- +"""PyMilo XGBoost Booster transporter.""" +import json + +from .transporter import AbstractTransporter +from ..utils.util import check_str_in_iterable + +_xgb = None +_xgboost_available = False +try: + import xgboost as _xgb_imported + _xgb = _xgb_imported + _xgboost_available = True +except BaseException: + _xgb = None + _xgboost_available = False + + +class BoosterTransporter(AbstractTransporter): + """Customized PyMilo Transporter developed to handle xgboost.Booster fields. + + The Booster's model state lives behind a C handle and cannot be reconstructed by walking + its Python __dict__. Therefore this transporter routes via isinstance(value, xgboost.Booster) + only (no key-name match) and relies on XGBoost's officially documented, version-stable JSON + model format. The booster JSON is parsed and inlined as a regular nested dict so the + resulting PyMilo JSON is fully transparent: every tree, leaf, and split condition stays + human-readable and editable. + """ + + PYMILO_BOOSTER_KEY = "pymilo-xgb-booster" + + def serialize(self, data, key, model_type): + """ + Serialize the value at data[key] if it is an xgboost.Booster instance. + + :param data: the internal data dictionary of the given model + :type data: dict + :param key: the special key of the data param, which we're going to serialize its value(data[key]) + :type key: object + :param model_type: the model type of the ML model + :type model_type: str + :return: pymilo serialized output of data[key] + """ + value = data[key] + if not _xgboost_available: + return value + if not isinstance(value, _xgb.Booster): + return value + raw = value.save_raw(raw_format='json') + if isinstance(raw, (bytes, bytearray)): + raw = bytes(raw).decode('utf-8') + booster_dict = json.loads(raw) + return { + self.PYMILO_BOOSTER_KEY: True, + "pymilo-bypass": True, + "data": booster_dict, + } + + def deserialize(self, data, key, model_type): + """ + Deserialize a previously serialized xgboost.Booster value. + + :param data: the internal data dictionary of the associated json file of the ML model + :type data: dict + :param key: the special key of the data param, which we're going to deserialize its value(data[key]) + :type key: object + :param model_type: the model type of the ML model + :type model_type: str + :return: pymilo deserialized output of data[key] + """ + content = data[key] + if not (isinstance(content, dict) and check_str_in_iterable(self.PYMILO_BOOSTER_KEY, content)): + return content + if not _xgboost_available: + return content + booster_payload = json.dumps(content["data"]).encode('utf-8') + booster = _xgb.Booster() + booster.load_model(bytearray(booster_payload)) + return booster diff --git a/pymilo/transporters/preprocessing_transporter.py b/pymilo/transporters/preprocessing_transporter.py index 9c3d2ae7..26015891 100644 --- a/pymilo/transporters/preprocessing_transporter.py +++ b/pymilo/transporters/preprocessing_transporter.py @@ -7,6 +7,15 @@ from .function_transporter import FunctionTransporter from scipy.interpolate._bsplines import BSpline +# scipy>=1.18 wraps a Cython/internal `_BSpline` behind the public `BSpline` API. +# Keep isinstance checks working across both generations. +_BSPLINE_TYPES = (BSpline,) +try: + from scipy.interpolate._bsplines import _BSpline + _BSPLINE_TYPES = (BSpline, _BSpline) +except ImportError: + pass + PREPROCESSING_CHAIN = { "GeneralDataStructureTransporter": GeneralDataStructureTransporter(), "FunctionTransporter": FunctionTransporter(), @@ -84,10 +93,10 @@ def serialize_pre_module(self, pre_module): for key, value in pre_module.__dict__.items(): if self.is_preprocessing_module(value): pre_module.__dict__[key] = self.serialize_pre_module(value) - elif isinstance(value, BSpline): + elif isinstance(value, _BSPLINE_TYPES): pre_module.__dict__[key] = self.serialize_spline(value) elif isinstance(value, list): - if len(value) > 0 and isinstance(value[0], BSpline): + if len(value) > 0 and isinstance(value[0], _BSPLINE_TYPES): pre_module.__dict__[key] = [self.serialize_spline(bspline) for bspline in value] for transporter in PREPROCESSING_CHAIN: @@ -129,7 +138,7 @@ def deserialize_pre_module(self, serialized_pre_module): def is_bspline(self, bspline): """ - Check whether the given module is a scipy.interpolate._bsplines.BSpline or not. + Check whether the given module is a scipy BSpline (or its internal _BSpline) or not. :param bspline: given object :type bspline: any @@ -139,22 +148,39 @@ def is_bspline(self, bspline): return check_str_in_iterable( "pymilo-preprocessing-type", bspline) and bspline["pymilo-preprocessing-type"] == "BSpline" - return get_sklearn_type(bspline) == "BSpline" + if isinstance(bspline, _BSPLINE_TYPES): + return True + # Fall back to type-name check for older wire formats / unexpected wrappers. + return get_sklearn_type(bspline) in ("BSpline", "_BSpline") def serialize_spline(self, bspline): """ Serialize scipy.interpolate._bsplines.BSpline object. + Only the public reconstructible state (t, c, k, extrapolate, axis) is persisted. + Dumping `bspline.__dict__` is unsafe on scipy>=1.18 because it contains a + non-JSON-serializable `_BSpline` delegate (`_delegate_to`) and array-api modules. + :param bspline: given scipy.interpolate._bsplines.BSpline module :type bspline: scipy.interpolate._bsplines.BSpline :return: pymilo serialized bspline """ - for transporter in PREPROCESSING_CHAIN: - PREPROCESSING_CHAIN[transporter].transport(bspline, Command.SERIALIZE) + gdst = PREPROCESSING_CHAIN["GeneralDataStructureTransporter"] + # Read via public properties — works for both legacy BSpline and scipy>=1.18 + # wrappers that store state on an internal `_BSpline` delegate. + raw = { + "t": bspline.t, + "c": bspline.c, + "k": bspline.k, + "extrapolate": bspline.extrapolate, + "axis": bspline.axis, + } + for key in list(raw.keys()): + raw[key] = gdst.serialize(raw, key, "BSpline") return { "pymilo-bypass": True, - "pymilo-preprocessing-type": get_sklearn_type(bspline), - "pymilo-preprocessing-data": bspline.__dict__ + "pymilo-preprocessing-type": "BSpline", + "pymilo-preprocessing-data": raw, } def deserialize_spline(self, serialized_bspline): @@ -166,15 +192,14 @@ def deserialize_spline(self, serialized_bspline): :return: retrieved associated scipy.interpolate._bsplines.BSpline object """ data = serialized_bspline["pymilo-preprocessing-data"] - associated_type = BSpline for key in list(data.keys()): for transporter in PREPROCESSING_CHAIN: data[key] = PREPROCESSING_CHAIN[transporter].deserialize(data, key, "") # Handle both old scipy (t, c, k) and new scipy (_t, _c, _k) attribute names + # from previously serialized payloads that dumped __dict__ directly. t = data.get("t", data.get("_t")) c = data.get("c", data.get("_c")) k = data.get("k", data.get("_k")) - retrieved_pre_module = associated_type(t=t, k=k, c=c) - for key in data: - setattr(retrieved_pre_module, key, data[key]) - return retrieved_pre_module + extrapolate = data.get("extrapolate", True) + axis = data.get("axis", 0) + return BSpline(t=t, c=c, k=k, extrapolate=extrapolate, axis=axis) diff --git a/pymilo/utils/test_pymilo.py b/pymilo/utils/test_pymilo.py index d06b66ad..dc458983 100644 --- a/pymilo/utils/test_pymilo.py +++ b/pymilo/utils/test_pymilo.py @@ -154,6 +154,38 @@ def pymilo_nearest_neighbor_test(nearest_neighbor, model_name, test_data): report_status(array_equal(pre_pymilo_kneighbors, post_pymilo_kneighbors), model_name) +def pymilo_xgb_booster_test(booster, model_name, test_data): + """ + Test native xgboost.Booster round-trip by comparing prediction metrics pre/post export. + + :param booster: trained xgboost.Booster + :type booster: xgboost.Booster + :param model_name: model name + :type model_name: str + :param test_data: data for testing as (x_test, y_test) + :type test_data: tuple + :return: True if the test succeed + """ + import xgboost as xgb + x_test, y_test = test_data + dtest = xgb.DMatrix(x_test) + pre_pymilo_model_y_pred = booster.predict(dtest) + pre_pymilo_model_prediction_output = { + "mean-error": mean_squared_error(y_test, pre_pymilo_model_y_pred), + "r2-score": r2_score(y_test, pre_pymilo_model_y_pred) + } + post_pymilo_model_y_pred = pymilo_test(booster, model_name).predict(xgb.DMatrix(x_test)) + post_pymilo_model_prediction_outputs = { + "mean-error": mean_squared_error(y_test, post_pymilo_model_y_pred), + "r2-score": r2_score(y_test, post_pymilo_model_y_pred) + } + comparison_result = compare_model_outputs( + pre_pymilo_model_prediction_output, + post_pymilo_model_prediction_outputs) + report_status(comparison_result, model_name) + return comparison_result + + def report_status(result, model_name): """ Print status for each model. diff --git a/pymilo/utils/util.py b/pymilo/utils/util.py index e43d0741..752ca614 100644 --- a/pymilo/utils/util.py +++ b/pymilo/utils/util.py @@ -3,7 +3,8 @@ import requests import importlib from inspect import signature -from ..pymilo_param import DOWNLOAD_MODEL_FAILED, INVALID_DOWNLOADED_MODEL, SKLEARN_SUPPORTED_CATEGORIES +from ..pymilo_param import DOWNLOAD_MODEL_FAILED, INVALID_DOWNLOADED_MODEL, NOT_SUPPORTED +from ..pymilo_param import SKLEARN_SUPPORTED_CATEGORIES, XGBOOST_SUPPORTED_CATEGORIES def get_sklearn_type(model): @@ -174,9 +175,47 @@ def get_sklearn_class(model_name): :param model_name: model name :type model_name: str - :return: sklearn ML model class + :return: sklearn ML model class or None """ for _, category_models in SKLEARN_SUPPORTED_CATEGORIES.items(): if model_name in category_models: - return category_models[model_name] + klass = category_models[model_name] + if klass == NOT_SUPPORTED: + return None + return klass return None + + +def get_xgboost_class(model_name): + """ + Return the xgboost class of the requested model name. + + :param model_name: model name + :type model_name: str + + :return: xgboost ML model class or None + """ + for _, category_models in XGBOOST_SUPPORTED_CATEGORIES.items(): + if model_name in category_models: + klass = category_models[model_name] + if klass == NOT_SUPPORTED: + return None + return klass + return None + + +def get_ml_model_class(model_name): + """ + Return the ML model class of the requested model name. + + Looks up sklearn-supported categories first, then XGBoost-supported categories. + + :param model_name: model name + :type model_name: str + + :return: ML model class or None when not found / not supported in this environment + """ + model_class = get_sklearn_class(model_name) + if model_class is not None: + return model_class + return get_xgboost_class(model_name) diff --git a/setup.py b/setup.py index 0477a561..c5ee6c01 100644 --- a/setup.py +++ b/setup.py @@ -8,6 +8,7 @@ INSTALLATION_MODES = { 'core': 'requirements.txt', 'streaming': 'streaming-requirements.txt', + 'xgboost': 'xgboost-requirements.txt', } @@ -55,6 +56,7 @@ def read_description(): install_requires=get_requires(), extras_require={ 'streaming': get_requires(mode='streaming'), + 'xgboost': get_requires(mode='xgboost'), }, python_requires='>=3.7', classifiers=[ diff --git a/tests/test_ml_streaming/run_server.py b/tests/test_ml_streaming/run_server.py index 335052e8..7fe7b18d 100644 --- a/tests/test_ml_streaming/run_server.py +++ b/tests/test_ml_streaming/run_server.py @@ -1,6 +1,7 @@ import argparse from sklearn.linear_model import LinearRegression from pymilo.streaming import PymiloServer, Compression, CommunicationProtocol +from pymilo.streaming.param import PYMILO_CLIENT_DEFAULT_CLIENT_ID, PYMILO_CLIENT_DEFAULT_ML_MODEL_ID from pymilo.utils.data_exporter import prepare_simple_regression_datasets @@ -44,8 +45,8 @@ def main(): compressor=Compression[args.compression], communication_protocol=CommunicationProtocol[args.protocol], ) - sample_client_id = "0x_demo_client_id" - sample_ml_model_id = "0x_demo_ml_model_id" + sample_client_id = PYMILO_CLIENT_DEFAULT_CLIENT_ID + sample_ml_model_id = PYMILO_CLIENT_DEFAULT_ML_MODEL_ID ps.init_client(sample_client_id) ps.init_ml_model(sample_client_id, sample_ml_model_id) ps.set_ml_model(sample_client_id, sample_ml_model_id, linear_regression) diff --git a/tests/test_ml_streaming/scenarios/scenario3.py b/tests/test_ml_streaming/scenarios/scenario3.py index cf555e43..4cb9cf11 100644 --- a/tests/test_ml_streaming/scenarios/scenario3.py +++ b/tests/test_ml_streaming/scenarios/scenario3.py @@ -4,7 +4,7 @@ from pymilo.utils.data_exporter import prepare_simple_regression_datasets -def scenario3(compression_method, communication_protocol): +def scenario3(compression_method, communication_protocol, server_url="127.0.0.1:9000"): # [PyMilo Server is initialized with ML Model] # 1. calculate mse in server # 2. download model in local @@ -16,11 +16,9 @@ def scenario3(compression_method, communication_protocol): client = PymiloClient( mode=PymiloClient.Mode.LOCAL, compressor=Compression[compression_method], - server_url="127.0.0.1:9000", + server_url=server_url, communication_protocol=CommunicationProtocol[communication_protocol], ) - client.client_id = "0x_demo_client_id" - client.ml_model_id = "0x_demo_ml_model_id" client.toggle_mode(PymiloClient.Mode.DELEGATE) result = client.predict(x_test) diff --git a/tests/test_ml_streaming/scenarios/scenario5.py b/tests/test_ml_streaming/scenarios/scenario5.py new file mode 100644 index 00000000..8c601f71 --- /dev/null +++ b/tests/test_ml_streaming/scenarios/scenario5.py @@ -0,0 +1,47 @@ +import numpy as np +import pytest + +xgboost = pytest.importorskip("xgboost") +from xgboost import XGBRegressor + +from pymilo.streaming import PymiloClient, Compression, CommunicationProtocol +from sklearn.metrics import mean_squared_error +from pymilo.utils.data_exporter import prepare_simple_regression_datasets + + +def scenario5(compression_method, communication_protocol): + # XGBoost model through ML streaming (same flow as scenario1): + # 1. create + train XGBRegressor locally + # 2. register on server + # 3. measure mse before streaming + # 4. upload / download + # 5. measure mse after streaming — must match + + x_train, y_train, x_test, y_test = prepare_simple_regression_datasets() + model = XGBRegressor( + n_estimators=8, + max_depth=3, + random_state=0, + verbosity=0, + ) + model.fit(x_train, y_train) + client = PymiloClient( + model=model, + mode=PymiloClient.Mode.LOCAL, + compressor=Compression[compression_method], + communication_protocol=CommunicationProtocol[communication_protocol], + ) + + client.register() + client.register_ml_model() + + result = client.predict(x_test) + mse_before = mean_squared_error(y_test, result) + + client.upload() + client.download() + + result = client.predict(x_test) + mse_after = mean_squared_error(y_test, result) + + return np.abs(mse_after - mse_before) diff --git a/tests/test_ml_streaming/test_streaming.py b/tests/test_ml_streaming/test_streaming.py index 45dc7585..05b316ff 100644 --- a/tests/test_ml_streaming/test_streaming.py +++ b/tests/test_ml_streaming/test_streaming.py @@ -1,4 +1,5 @@ import os +import json import time import pytest import subprocess @@ -15,17 +16,6 @@ params=["NULL", "GZIP", "ZLIB", "LZMA", "BZ2"]) def prepare_bare_server(request): compression_method = request.param - # Using PyMilo direct CLI - # server_proc = subprocess.Popen( - # [ - # executable, - # "-m", "pymilo", - # "--compression", compression_method, - # "--protocol", "REST", - # "--port", "8000", - # "--bare", - # ], - # ) path = os.path.join( os.getcwd(), "tests", @@ -51,18 +41,6 @@ def prepare_bare_server(request): def prepare_ml_server(request): communication_protocol = request.param compression_method = "ZLIB" - # Using PyMilo direct CLI - # server_proc = subprocess.Popen( - # [ - # executable, - # "-m", "pymilo", - # "--compression", compression_method, - # "--protocol", communication_protocol, - # "--port", "9000", - # "--load", os.path.join(os.getcwd(), "tests", "test_exceptions", "valid_jsons", "linear_regression.json") - # # "--load", "https://raw.githubusercontent.com/openscilab/pymilo/main/tests/test_exceptions/valid_jsons/linear_regression.json", - # ], - # ) path = os.path.join( os.getcwd(), "tests", @@ -110,6 +88,27 @@ def prepare_access_control_server(request): time.sleep(2) +@pytest.fixture(scope="session") +def prepare_cli_server(): + # The PyMilo CLI is what the generated Dockerfiles run, so it needs coverage of its own. + compression_method, communication_protocol = "ZLIB", "REST" + server_proc = subprocess.Popen( + [ + executable, + "-m", "pymilo", + "--compression", compression_method, + "--protocol", communication_protocol, + "--port", "8600", + "--load", os.path.join( + os.getcwd(), "tests", "test_exceptions", "valid_jsons", "linear_regression.json"), + ], + ) + time.sleep(10) + yield (compression_method, communication_protocol) + server_proc.terminate() + time.sleep(2) + + def test1(prepare_bare_server): compression_method, communication_protocol = prepare_bare_server assert scenario1(compression_method, communication_protocol) == 0 @@ -130,6 +129,29 @@ def test4(prepare_access_control_server): assert scenario4(compression_method, communication_protocol) == 0 +def test5(prepare_bare_server): + pytest.importorskip("xgboost") + from scenarios.scenario5 import scenario5 + compression_method, communication_protocol = prepare_bare_server + assert scenario5(compression_method, communication_protocol) == 0 + + +def test6(prepare_cli_server): + # same scenario as test3, against a server the PyMilo CLI preloaded through `--load` + compression_method, communication_protocol = prepare_cli_server + assert scenario3(compression_method, communication_protocol, server_url="127.0.0.1:8600") == 0 + + +@pytest.mark.parametrize("cli_args, port", [(["--init", "LinearRegression"], "8400"), (["--bare"], "8401")]) +def test_cli_server_boots(cli_args, port): + server_proc = subprocess.Popen([executable, "-m", "pymilo", "--port", port] + cli_args) # nosec B603 + time.sleep(8) + is_running = server_proc.poll() is None + server_proc.terminate() + time.sleep(2) + assert is_running + + def test_dockerfile(): docker_files_folder = os.path.join( os.getcwd(), @@ -154,5 +176,13 @@ def test_dockerfile(): "Dockerfile2" ) ) + generate_dockerfile( + dockerfile_name="Dockerfile", + init_model="LinearRegression", + ) + # the `--init` branch of generate_dockerfile used to emit a malformed (non-JSON) CMD + with open('Dockerfile') as f: + cmd_line = [line for line in f.read().splitlines() if line.startswith("CMD ")][0] + assert json.loads(cmd_line[len("CMD "):])[-2:] == ["--init", "LinearRegression"] os.remove(path='Dockerfile') assert r1 and r2 diff --git a/tests/test_xgboost/native_booster.py b/tests/test_xgboost/native_booster.py new file mode 100644 index 00000000..7c2ce4f6 --- /dev/null +++ b/tests/test_xgboost/native_booster.py @@ -0,0 +1,32 @@ +import pytest + +xgboost = pytest.importorskip("xgboost") + +from pymilo.utils.data_exporter import prepare_simple_regression_datasets, prepare_simple_classification_datasets +from pymilo.utils.test_pymilo import pymilo_xgb_booster_test + + +def native_booster_regression(): + x_train, y_train, x_test, y_test = prepare_simple_regression_datasets() + dtrain = xgboost.DMatrix(x_train, label=y_train) + params = { + "objective": "reg:squarederror", + "max_depth": 3, + "eta": 0.1, + "verbosity": 0, + } + booster = xgboost.train(params, dtrain, num_boost_round=8) + assert pymilo_xgb_booster_test(booster, "NativeBooster-Regression", (x_test, y_test)) == True + + +def native_booster_classification(): + x_train, y_train, x_test, y_test = prepare_simple_classification_datasets() + dtrain = xgboost.DMatrix(x_train, label=y_train) + params = { + "objective": "binary:logistic", + "max_depth": 3, + "eta": 0.1, + "verbosity": 0, + } + booster = xgboost.train(params, dtrain, num_boost_round=8) + assert pymilo_xgb_booster_test(booster, "NativeBooster-Binary", (x_test, y_test)) == True diff --git a/tests/test_xgboost/test_xgboost_models.py b/tests/test_xgboost/test_xgboost_models.py new file mode 100644 index 00000000..92f5b864 --- /dev/null +++ b/tests/test_xgboost/test_xgboost_models.py @@ -0,0 +1,55 @@ +import os +import pytest + +xgboost = pytest.importorskip("xgboost") + +from xgb_regressor import xgb_regressor_all_objectives +from xgb_classifier_binary import xgb_classifier_binary_all_objectives +from xgb_classifier_multiclass import xgb_classifier_multiclass_all_objectives +from xgb_ranker import xgb_ranker_all_objectives +from xgb_rf_regressor import xgb_rf_regressor +from xgb_rf_classifier import xgb_rf_classifier +from xgb_survival import xgb_survival_cox +from native_booster import native_booster_regression, native_booster_classification +from xgb_nested import ( + xgb_inside_pipeline, + xgb_inside_voting_regressor, + xgb_untrained_model, + xgb_json_dump_round_trip, +) + + +XGBOOST_MODELS = { + "REGRESSOR": [xgb_regressor_all_objectives], + "CLASSIFIER_BINARY": [xgb_classifier_binary_all_objectives], + "CLASSIFIER_MULTICLASS": [xgb_classifier_multiclass_all_objectives], + "RANKER": [xgb_ranker_all_objectives], + "RF_REGRESSOR": [xgb_rf_regressor], + "RF_CLASSIFIER": [xgb_rf_classifier], + "SURVIVAL": [xgb_survival_cox], + "NATIVE_BOOSTER": [native_booster_regression, native_booster_classification], + "NESTED": [ + xgb_inside_pipeline, + xgb_inside_voting_regressor, + xgb_untrained_model, + xgb_json_dump_round_trip, + ], +} + + +@pytest.fixture(scope="session", autouse=True) +def reset_exported_xgboost_directory(): + exported_dir = os.path.join(os.getcwd(), "tests", "exported_xgboost") + if not os.path.isdir(exported_dir): + os.makedirs(exported_dir, exist_ok=True) + return + for file_name in os.listdir(exported_dir): + json_file = os.path.join(exported_dir, file_name) + if os.path.isfile(json_file): + os.remove(json_file) + + +def test_full(): + for category in XGBOOST_MODELS: + for func in XGBOOST_MODELS[category]: + func() diff --git a/tests/test_xgboost/xgb_classifier_binary.py b/tests/test_xgboost/xgb_classifier_binary.py new file mode 100644 index 00000000..6a086a80 --- /dev/null +++ b/tests/test_xgboost/xgb_classifier_binary.py @@ -0,0 +1,31 @@ +import pytest + +xgboost = pytest.importorskip("xgboost") +from xgboost import XGBClassifier + +from pymilo.utils.data_exporter import prepare_simple_classification_datasets +from pymilo.utils.test_pymilo import pymilo_classification_test, pymilo_regression_test + + +BINARY_OBJECTIVES = [ + ("binary:logistic", "XGBClassifier-BinaryLogistic"), + ("binary:hinge", "XGBClassifier-BinaryHinge"), +] + + +def xgb_classifier_binary_all_objectives(): + x_train, y_train, x_test, y_test = prepare_simple_classification_datasets() + for objective, name in BINARY_OBJECTIVES: + model = XGBClassifier( + objective=objective, + n_estimators=8, + max_depth=3, + random_state=0, + ) + model.fit(x_train, y_train) + # binary:hinge yields discrete labels; hinge_loss metrics still apply via classification helper. + # Prefer classification metrics for logistic; hinge can use either path. + if objective == "binary:logistic": + assert pymilo_classification_test(model, name, (x_test, y_test)) == True + else: + assert pymilo_regression_test(model, name, (x_test, y_test)) == True diff --git a/tests/test_xgboost/xgb_classifier_multiclass.py b/tests/test_xgboost/xgb_classifier_multiclass.py new file mode 100644 index 00000000..538dc0f3 --- /dev/null +++ b/tests/test_xgboost/xgb_classifier_multiclass.py @@ -0,0 +1,37 @@ +import pytest +from sklearn.datasets import load_iris + +xgboost = pytest.importorskip("xgboost") +from xgboost import XGBClassifier + +from pymilo.utils.test_pymilo import pymilo_regression_test + + +MULTICLASS_OBJECTIVES = [ + ("multi:softmax", "XGBClassifier-MultiSoftmax"), + ("multi:softprob", "XGBClassifier-MultiSoftprob"), +] + + +def _split_iris(threshold=20): + X, y = load_iris(return_X_y=True) + rng_idx = list(range(len(y))) + interleaved = sorted(rng_idx, key=lambda i: (i % 3, i)) + X = X[interleaved] + y = y[interleaved] + return X[:-threshold], y[:-threshold], X[-threshold:], y[-threshold:] + + +def xgb_classifier_multiclass_all_objectives(): + x_train, y_train, x_test, y_test = _split_iris() + for objective, name in MULTICLASS_OBJECTIVES: + model = XGBClassifier( + objective=objective, + num_class=3, + n_estimators=8, + max_depth=3, + random_state=0, + ) + model.fit(x_train, y_train) + # Use regression metrics on class labels (hinge_loss is binary-oriented). + assert pymilo_regression_test(model, name, (x_test, y_test)) == True diff --git a/tests/test_xgboost/xgb_nested.py b/tests/test_xgboost/xgb_nested.py new file mode 100644 index 00000000..f0c14eac --- /dev/null +++ b/tests/test_xgboost/xgb_nested.py @@ -0,0 +1,67 @@ +import numpy as np +import pytest + +xgboost = pytest.importorskip("xgboost") +from xgboost import XGBRegressor + +from sklearn.pipeline import Pipeline +from sklearn.preprocessing import StandardScaler +from sklearn.ensemble import VotingRegressor +from sklearn.linear_model import LinearRegression + +from pymilo import Export, Import +from pymilo.utils.data_exporter import prepare_simple_regression_datasets +from pymilo.utils.test_pymilo import pymilo_regression_test, report_status + + +def xgb_inside_pipeline(): + """An XGBoost model nested in a sklearn Pipeline is routed through the compose chain.""" + x_train, y_train, x_test, y_test = prepare_simple_regression_datasets() + model = Pipeline([ + ("scaler", StandardScaler()), + ("xgb", XGBRegressor(n_estimators=8, max_depth=3, random_state=0)), + ]) + model.fit(x_train, y_train) + assert pymilo_regression_test(model, "Pipeline-XGBRegressor", (x_test, y_test)) == True + + +def xgb_inside_voting_regressor(): + """An XGBoost model nested in a sklearn ensemble is routed through the ensemble chain.""" + x_train, y_train, x_test, y_test = prepare_simple_regression_datasets() + model = VotingRegressor([ + ("xgb", XGBRegressor(n_estimators=8, max_depth=3, random_state=0)), + ("linear", LinearRegression()), + ]) + model.fit(x_train, y_train) + assert pymilo_regression_test(model, "VotingRegressor-XGBRegressor", (x_test, y_test)) == True + + +def _same_param(left, right): + """XGBoost's default `missing` hyperparameter is NaN, which never equals itself.""" + if isinstance(left, float) and isinstance(right, float) and np.isnan(left) and np.isnan(right): + return True + return left == right + + +def xgb_untrained_model(): + """An unfitted wrapper carries no inner Booster yet, its hyperparameters must still survive.""" + model = XGBRegressor(n_estimators=7, max_depth=4, random_state=0) + retrieved = Import(json_dump=Export(model).to_json()).to_model() + expected, retrieved_params = model.get_params(), retrieved.get_params() + result = set(expected) == set(retrieved_params) and all( + _same_param(expected[key], retrieved_params[key]) for key in expected) + # NaN survives the JSON round-trip rather than degrading to None + result = result and np.isnan(retrieved.missing) + report_status(result, "XGBRegressor-Untrained") + assert result == True + + +def xgb_json_dump_round_trip(): + """Importing from an in-memory JSON dump must be as lossless as importing from a file.""" + x_train, y_train, x_test, _ = prepare_simple_regression_datasets() + model = XGBRegressor(n_estimators=8, max_depth=3, random_state=0) + model.fit(x_train, y_train) + retrieved = Import(json_dump=Export(model).to_json()).to_model() + result = np.array_equal(model.predict(x_test), retrieved.predict(x_test)) + report_status(result, "XGBRegressor-JsonDump") + assert result == True diff --git a/tests/test_xgboost/xgb_ranker.py b/tests/test_xgboost/xgb_ranker.py new file mode 100644 index 00000000..078a4a06 --- /dev/null +++ b/tests/test_xgboost/xgb_ranker.py @@ -0,0 +1,47 @@ +import numpy as np +import pytest + +xgboost = pytest.importorskip("xgboost") +from xgboost import XGBRanker + +from pymilo.utils.test_pymilo import pymilo_regression_test + + +RANKER_OBJECTIVES = [ + ("rank:pairwise", "XGBRanker-Pairwise"), + ("rank:ndcg", "XGBRanker-NDCG"), + ("rank:map", "XGBRanker-MAP"), +] + + +def _synthetic_ranking_dataset(n_queries=20, items_per_query=10, n_features=6, seed=0): + rng = np.random.default_rng(seed) + X = rng.normal(size=(n_queries * items_per_query, n_features)) + y = (X[:, 0] + rng.normal(scale=0.1, size=X.shape[0])).round().clip(0, 4).astype(int) + group = np.full(n_queries, items_per_query, dtype=int) + return X, y, group + + +def xgb_ranker_all_objectives(): + X, y, group = _synthetic_ranking_dataset() + n_train_queries = 16 + items = 10 + cut = n_train_queries * items + x_train, y_train, group_train = X[:cut], y[:cut], group[:n_train_queries] + x_test, y_test = X[cut:], y[cut:] + + # rank:map requires binary relevance labels (0/1 only). + y_train_binary = (y_train > 0).astype(int) + y_test_binary = (y_test > 0).astype(int) + + for objective, name in RANKER_OBJECTIVES: + labels_train = y_train_binary if objective == "rank:map" else y_train + labels_test = y_test_binary if objective == "rank:map" else y_test + model = XGBRanker( + objective=objective, + n_estimators=8, + max_depth=3, + random_state=0, + ) + model.fit(x_train, labels_train, group=group_train) + assert pymilo_regression_test(model, name, (x_test, labels_test)) == True diff --git a/tests/test_xgboost/xgb_regressor.py b/tests/test_xgboost/xgb_regressor.py new file mode 100644 index 00000000..799040e8 --- /dev/null +++ b/tests/test_xgboost/xgb_regressor.py @@ -0,0 +1,41 @@ +import numpy as np +import pytest + +xgboost = pytest.importorskip("xgboost") +from xgboost import XGBRegressor + +from pymilo.utils.data_exporter import prepare_simple_regression_datasets +from pymilo.utils.test_pymilo import pymilo_regression_test + + +REGRESSION_OBJECTIVES = [ + ("reg:squarederror", "XGBRegressor-SquaredError"), + ("reg:absoluteerror", "XGBRegressor-AbsoluteError"), + ("reg:gamma", "XGBRegressor-Gamma"), + ("reg:tweedie", "XGBRegressor-Tweedie"), +] + + +def _positive_targets(y): + """Gamma and Tweedie require strictly positive targets.""" + y = np.asarray(y, dtype=float) + return np.where(y > 0, y, 1e-3) + + +def xgb_regressor_all_objectives(): + x_train, y_train, x_test, y_test = prepare_simple_regression_datasets() + for objective, name in REGRESSION_OBJECTIVES: + if objective in ("reg:gamma", "reg:tweedie"): + y_train_use = _positive_targets(y_train) + y_test_use = _positive_targets(y_test) + else: + y_train_use = y_train + y_test_use = y_test + model = XGBRegressor( + objective=objective, + n_estimators=8, + max_depth=3, + random_state=0, + ) + model.fit(x_train, y_train_use) + assert pymilo_regression_test(model, name, (x_test, y_test_use)) == True diff --git a/tests/test_xgboost/xgb_rf_classifier.py b/tests/test_xgboost/xgb_rf_classifier.py new file mode 100644 index 00000000..31f03f4f --- /dev/null +++ b/tests/test_xgboost/xgb_rf_classifier.py @@ -0,0 +1,18 @@ +import pytest + +xgboost = pytest.importorskip("xgboost") +from xgboost import XGBRFClassifier + +from pymilo.utils.data_exporter import prepare_simple_classification_datasets +from pymilo.utils.test_pymilo import pymilo_classification_test + + +def xgb_rf_classifier(): + x_train, y_train, x_test, y_test = prepare_simple_classification_datasets() + model = XGBRFClassifier( + n_estimators=10, + max_depth=4, + random_state=0, + ) + model.fit(x_train, y_train) + assert pymilo_classification_test(model, "XGBRFClassifier", (x_test, y_test)) == True diff --git a/tests/test_xgboost/xgb_rf_regressor.py b/tests/test_xgboost/xgb_rf_regressor.py new file mode 100644 index 00000000..de4e00ce --- /dev/null +++ b/tests/test_xgboost/xgb_rf_regressor.py @@ -0,0 +1,18 @@ +import pytest + +xgboost = pytest.importorskip("xgboost") +from xgboost import XGBRFRegressor + +from pymilo.utils.data_exporter import prepare_simple_regression_datasets +from pymilo.utils.test_pymilo import pymilo_regression_test + + +def xgb_rf_regressor(): + x_train, y_train, x_test, y_test = prepare_simple_regression_datasets() + model = XGBRFRegressor( + n_estimators=10, + max_depth=4, + random_state=0, + ) + model.fit(x_train, y_train) + assert pymilo_regression_test(model, "XGBRFRegressor", (x_test, y_test)) == True diff --git a/tests/test_xgboost/xgb_survival.py b/tests/test_xgboost/xgb_survival.py new file mode 100644 index 00000000..b7aa1143 --- /dev/null +++ b/tests/test_xgboost/xgb_survival.py @@ -0,0 +1,43 @@ +import numpy as np +import pytest + +xgboost = pytest.importorskip("xgboost") +from xgboost import XGBRegressor + +from pymilo.utils.test_pymilo import pymilo_regression_test + + +def _synthetic_survival_dataset(n=200, n_features=5, censor_fraction=0.3, seed=0): + """ + Build a small Cox-style dataset. + + XGBoost's `survival:cox` objective encodes labels as positive (observed event) + or negative (right-censored at |y|) survival times. + """ + rng = np.random.default_rng(seed) + X = rng.normal(size=(n, n_features)) + linpred = X[:, 0] - 0.5 * X[:, 1] + times = rng.exponential(scale=np.exp(-linpred)) + censored_mask = rng.uniform(size=n) < censor_fraction + y = np.where(censored_mask, -times, times) + return X, y + + +def xgb_survival_cox(): + try: + X, y = _synthetic_survival_dataset() + model = XGBRegressor( + objective="survival:cox", + n_estimators=8, + max_depth=3, + random_state=0, + ) + model.fit(X[:-30], y[:-30]) + except Exception as exc: + pytest.skip(f"survival:cox is unavailable in this xgboost build: {exc}") + return + + x_test = X[-30:] + y_test = y[-30:] + assert pymilo_regression_test( + model, "XGBRegressor-SurvivalCox", (x_test, y_test)) == True diff --git a/xgboost-requirements.txt b/xgboost-requirements.txt new file mode 100644 index 00000000..7d98309e --- /dev/null +++ b/xgboost-requirements.txt @@ -0,0 +1 @@ +xgboost>=1.7.0; python_version >= "3.8"