maxframe 0.1.0b5__cp39-cp39-macosx_11_0_arm64.whl

maxframe/extension.py

@@ -0,0 +1,81 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import abc
+import itertools
+from typing import TYPE_CHECKING, Iterable, Tuple, Type
+
+if TYPE_CHECKING:
+    from maxframe.codegen import BigDagCodeGenerator
+
+    try:
+        import maxframe_framedriver.runners
+    except ImportError:
+        pass
+
+
+class MaxFrameExtension(metaclass=abc.ABCMeta):
+    @classmethod
+    @abc.abstractmethod
+    def get_codegen(cls) -> Type["BigDagCodeGenerator"]:
+        raise NotImplementedError
+
+    @classmethod
+    @abc.abstractmethod
+    def get_runner(cls) -> Type["maxframe_framedriver.runners.SubDagRunner"]:
+        raise NotImplementedError
+
+    @classmethod
+    async def destroy_session(cls, session_id: str) -> None:
+        """
+        This interface will be called once a session in the framedriver is ended.
+
+        Parameters
+        ----------
+        session_id : str
+            The session id.
+        """
+        pass
+
+    @classmethod
+    def init_service_extension(cls) -> None:
+        """
+        Init the services of the extension before the app is actually run and will be
+        called only once.
+        """
+        pass
+
+    @classmethod
+    def destroy_service_extension(cls) -> None:
+        """
+        Destroy the services of the extension before the app is actually stopped and
+        will be called only once.
+        """
+        pass
+
+
+def iter_extensions() -> Iterable[Tuple[str, MaxFrameExtension]]:
+    try:
+        from importlib.metadata import entry_points
+    except (ImportError, TypeError):
+        from importlib_metadata import entry_points
+
+    eps = entry_points(group="maxframe.extension")
+    if callable(getattr(eps, "values", None)):
+        _it = itertools.chain(*eps.values())
+    else:
+        _it = eps
+
+    for entry_point in _it:
+        yield entry_point.name, entry_point.load()

maxframe/learn/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from . import contrib
+
+del contrib

maxframe/learn/contrib/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from . import pytorch
+
+del pytorch

maxframe/learn/contrib/pytorch/__init__.py

@@ -0,0 +1,16 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .run_function import run_pytorch_function
+from .run_script import run_pytorch_script

maxframe/learn/contrib/pytorch/run_function.py

@@ -0,0 +1,110 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import io
+from typing import Any, Callable, Dict
+
+from ....core import TILEABLE_TYPE
+from ....lib import wrapped_pickle as pickle
+from ....typing_ import SessionType
+from ....utils import find_objects, replace_objects
+from .run_script import run_pytorch_script
+
+_script_template = """
+import cloudpickle
+from maxframe.utils import replace_objects
+
+
+def main(**kwargs):
+    func = cloudpickle.loads(%(pickled_func)r)
+    nested = cloudpickle.loads(%(nested)r)
+    args, kw = replace_objects(nested, kwargs)
+    return func(*args, **kw)
+
+
+if __name__ == "__main__":
+    vars = dict()
+    for var_name in cloudpickle.loads(%(var_names)r):
+        vars[var_name] = globals()[var_name]
+    main(**vars)
+"""
+
+
+def run_pytorch_function(
+    func: Callable,
+    args: tuple = (),
+    kwargs: dict = None,
+    *,
+    n_workers: int = 1,
+    retry_when_fail: bool = False,
+    session: SessionType = None,
+    run_kwargs: Dict[str, Any] = None,
+    port: int = None,
+    execute: bool = True,
+):
+    """
+    Run PyTorch function in MaxFrame cluster.
+
+    Besides args and kwargs, the function will receive extra
+    environment variables from the caller:
+
+    * MASTER_ADDR, MASTER_PORT: the endpoint of the master node
+    * RANK: the index of the current worker
+
+    Parameters
+    ----------
+    func: Callable
+        Function or callable object to run
+    args: tuple
+        Args tuple to pass to the function
+    kwargs: dict
+    n_workers : int
+        Number of PyTorch workers
+    retry_when_fail : bool
+        If True, retry when function failed.
+    session
+        MaxFrame session, if not provided, will use default one.
+    run_kwargs : dict
+        Extra kwargs for `session.run`.
+    port : int
+        Port of PyTorch worker or ps, will automatically increase for the same worker
+
+    Returns
+    -------
+    status
+        return {'status': 'ok'} if succeeded, or error raised
+    """
+    kwargs = kwargs or {}
+    packed = [args, kwargs]
+    inputs = find_objects(packed, TILEABLE_TYPE)
+    input_to_var = {t: f"torch_var_{i}" for i, t in enumerate(inputs)}
+    var_to_input = {v: k for k, v in input_to_var.items()}
+    structure = replace_objects(packed, input_to_var)
+
+    replaces = {
+        "pickled_func": pickle.dumps(func),
+        "nested": pickle.dumps(structure),
+        "var_names": pickle.dumps(list(input_to_var.values())),
+    }
+    code = _script_template % replaces
+    return run_pytorch_script(
+        io.StringIO(code),
+        n_workers,
+        data=var_to_input,
+        retry_when_fail=retry_when_fail,
+        session=session,
+        run_kwargs=run_kwargs,
+        port=port,
+        execute=execute,
+    )

maxframe/learn/contrib/pytorch/run_script.py

@@ -0,0 +1,102 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import os
+from typing import Any, BinaryIO, Dict, List, Optional, TextIO, Union
+
+from .... import opcodes
+from ....remote.run_script import RunScript, _extract_inputs
+from ....serialization.serializables import Int32Field, StringField
+from ....typing_ import SessionType, TileableType
+from ....utils import to_binary
+
+logger = logging.getLogger(__name__)
+
+
+class RunPyTorch(RunScript):
+    _op_type_ = opcodes.RUN_PYTORCH
+
+    # used for chunk op
+    master_port = Int32Field("master_port", default=None)
+    master_addr = StringField("master_addr", default=None)
+    init_method = StringField("init_method", default=None)
+    master_waiter_name = StringField("master_waiter_name", default=None)
+
+
+def run_pytorch_script(
+    script: Union[bytes, str, BinaryIO, TextIO],
+    n_workers: int,
+    data: Dict[str, TileableType] = None,
+    gpu: Optional[bool] = None,
+    command_argv: List[str] = None,
+    retry_when_fail: bool = False,
+    session: SessionType = None,
+    run_kwargs: Dict[str, Any] = None,
+    port: int = None,
+    execute: bool = True,
+):
+    """
+    Run PyTorch script in MaxFrame cluster.
+
+    Parameters
+    ----------
+    script: str or file-like object
+        Script to run
+    n_workers : int
+        Number of PyTorch workers
+    data : dict
+        Variable name to data.
+    gpu : bool
+        Run PyTorch script on GPU
+    command_argv : list
+        Extra command args for script
+    retry_when_fail : bool
+        If True, retry when function failed.
+    session
+        MaxFrame session, if not provided, will use default one.
+    run_kwargs : dict
+        Extra kwargs for `session.run`.
+    port : int
+        Port of PyTorch worker or ps, will automatically increase for the same worker
+
+    Returns
+    -------
+    status
+        return {'status': 'ok'} if succeeded, or error raised
+    """
+    if int(n_workers) <= 0:
+        raise ValueError("n_workers should be at least 1")
+    if hasattr(script, "read"):
+        code = script.read()
+    else:
+        with open(os.path.abspath(script), "rb") as f:
+            code = f.read()
+
+    inputs = _extract_inputs(data)
+    port = 29500 if port is None else port
+    op = RunPyTorch(
+        data=data,
+        code=to_binary(code),
+        world_size=int(n_workers),
+        retry_when_fail=retry_when_fail,
+        gpu=gpu,
+        master_port=port,
+        command_args=command_argv,
+    )
+    t = op(inputs)
+    if execute:
+        t.execute(session=session, **(run_kwargs or {}))
+    else:
+        return t

maxframe/learn/contrib/pytorch/tests/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

maxframe/learn/contrib/pytorch/tests/test_pytorch.py

@@ -0,0 +1,42 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ..... import dataframe as md
+from ..... import tensor as mt
+from .....core import TILEABLE_TYPE
+from ..run_function import run_pytorch_function
+
+
+def test_run_function():
+    def test_func(a, b):
+        raise ValueError(a + b)
+
+    df = md.DataFrame(mt.random.rand(10, 10))
+    t = run_pytorch_function(test_func, args=(df, 1), execute=False)
+
+    global_dict = {"__name__": "__main__"}
+    global_dict.update(t.op.data)
+
+    try:
+        exec(t.op.code, global_dict)
+    except ValueError as ex:
+        result = ex.args[0]
+        assert isinstance(result, TILEABLE_TYPE)
+        assert result.op.lhs.key == df.key
+        assert result.op.rhs == 1
+    else:
+        raise AssertionError("Error not raised")
+
+    assert global_dict["torch_var_0"] is df
+    assert "torch_var_1" not in global_dict

maxframe/learn/contrib/utils.py

@@ -0,0 +1,52 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import sys
+
+
+def make_import_error_func(package_name):
+    def _func(*_, **__):  # pragma: no cover
+        raise ImportError(
+            f"Cannot import {package_name}, please reinstall that package."
+        )
+
+    return _func
+
+
+def config_mod_getattr(mod_dict, globals_):
+    def __getattr__(name):
+        import importlib
+
+        if name in mod_dict:
+            mod_name, cls_name = mod_dict[name].rsplit(".", 1)
+            mod = importlib.import_module(mod_name, globals_["__name__"])
+            cls = globals_[name] = getattr(mod, cls_name)
+            return cls
+        else:  # pragma: no cover
+            raise AttributeError(name)
+
+    if sys.version_info[:2] < (3, 7):
+        for _mod in mod_dict.keys():
+            __getattr__(_mod)
+
+    def __dir__():
+        return sorted([n for n in globals_ if not n.startswith("_")] + list(mod_dict))
+
+    globals_.update(
+        {
+            "__getattr__": __getattr__,
+            "__dir__": __dir__,
+            "__all__": list(__dir__()),
+            "__warningregistry__": dict(),
+        }
+    )

maxframe/learn/contrib/xgboost/__init__.py

@@ -0,0 +1,26 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ..utils import config_mod_getattr as _config_mod_getattr
+from .dmatrix import DMatrix
+from .predict import predict
+from .train import train
+
+_config_mod_getattr(
+    {
+        "XGBClassifier": ".classifier.XGBClassifier",
+        "XGBRegressor": ".regressor.XGBRegressor",
+    },
+    globals(),
+)

maxframe/learn/contrib/xgboost/classifier.py

@@ -0,0 +1,86 @@
+# Copyright 1999-2024 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+
+from ....tensor import argmax
+from ..utils import make_import_error_func
+from .core import XGBScikitLearnBase, xgboost
+
+if not xgboost:
+    XGBClassifier = make_import_error_func("xgboost")
+else:
+    from xgboost.sklearn import XGBClassifierBase
+
+    from .core import wrap_evaluation_matrices
+    from .predict import predict
+    from .train import train
+
+    class XGBClassifier(XGBScikitLearnBase, XGBClassifierBase):
+        """
+        Implementation of the scikit-learn API for XGBoost classification.
+        """
+
+        def fit(
+            self,
+            X,
+            y,
+            sample_weight=None,
+            base_margin=None,
+            eval_set=None,
+            sample_weight_eval_set=None,
+            base_margin_eval_set=None,
+            num_class=None,
+        ):
+            dtrain, evals = wrap_evaluation_matrices(
+                None,
+                X,
+                y,
+                sample_weight,
+                base_margin,
+                eval_set,
+                sample_weight_eval_set,
+                base_margin_eval_set,
+            )
+            params = self.get_xgb_params()
+            self.n_classes_ = num_class or 1
+            if self.n_classes_ > 2:
+                params["objective"] = "multi:softprob"
+                params["num_class"] = self.n_classes_
+            else:
+                params["objective"] = "binary:logistic"
+            self.evals_result_ = dict()
+            result = train(
+                params,
+                dtrain,
+                num_boost_round=self.get_num_boosting_rounds(),
+                evals=evals,
+                evals_result=self.evals_result_,
+                num_class=num_class,
+            )
+            self._Booster = result
+            return self
+
+        def predict(self, data, **kw):
+            prob = self.predict_proba(data, flag=True, **kw)
+            if prob.ndim > 1:
+                prediction = argmax(prob, axis=1)
+            else:
+                prediction = (prob > 0.5).astype(np.int64)
+            return prediction
+
+        def predict_proba(self, data, ntree_limit=None, flag=False, **kw):
+            if ntree_limit is not None:
+                raise NotImplementedError("ntree_limit is not currently supported")
+            return predict(self.get_booster(), data, flag=flag, **kw)