maxframe 1.3.1__cp310-cp310-macosx_10_9_universal2.whl → 2.0.0b2__cp310-cp310-macosx_10_9_universal2.whl

maxframe/codegen/spe/tests/test_spe_codegen.py

@@ -0,0 +1,141 @@
+# Copyright 1999-2025 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 base64
+import hashlib
+import textwrap
+
+import pytest
+
+from .... import dataframe as md
+from .... import tensor as mt
+from ....dataframe.utils import bind_func_args_from_pos
+from ....lib import wrapped_pickle as pickle
+from ..core import SPECodeGenerator
+
+
+@pytest.fixture
+def codegen():
+    return SPECodeGenerator("session_1")
+
+
+def test_simple_codegen(codegen):
+    df = md.DataFrame(mt.random.rand(1000, 5), columns=list("ABCDE"))
+    df["F"] = df["A"] + df["B"]
+    df["G"] = df["C"] * 2 + df["D"]
+    result = df.cumsum()
+    dag = result.build_graph()
+    generated = codegen.generate(dag)
+
+    local_vars = generated.constants.copy()
+    exec(generated.code, {}, local_vars)
+    df_result = local_vars[generated.output_key_to_variables[result.key]]
+    assert len(df_result.columns) == 7
+
+
+def test_codegen_with_udf(codegen):
+    df = md.DataFrame(mt.random.rand(1, 3), columns=list("ABC"))
+
+    def f1(x) -> int:
+        return x + 1
+
+    def f2(x) -> int:
+        return x - 1
+
+    result = df.transform({"A": f1, "B": [f1, f2]})
+    dag = result.build_graph()
+    generated = codegen.generate(dag)
+
+    udf_1_body = base64.b64encode(pickle.dumps(f1, protocol=pickle.DEFAULT_PROTOCOL))
+    udf_2_body = base64.b64encode(pickle.dumps(f2, protocol=pickle.DEFAULT_PROTOCOL))
+    udf_1_value = (
+        f"pickled_data = cloudpickle.loads(base64.b64decode({udf_1_body}), buffers=[])"
+    )
+    udf_2_value = (
+        f"pickled_data = cloudpickle.loads(base64.b64decode({udf_2_body}), buffers=[])"
+    )
+    udf_1 = f"user_udf_f1_{hashlib.md5(udf_1_value.encode('utf-8')).hexdigest()}"
+    udf_2 = f"user_udf_f2_{hashlib.md5(udf_2_value.encode('utf-8')).hexdigest()}"
+
+    expected_contents = f"""
+    import numpy as np
+    import pandas as pd
+    import base64
+    import cloudpickle
+    {udf_1_value}
+    {udf_1} = pickled_data
+    {udf_2_value}
+    {udf_2} = pickled_data
+    if not running:
+        raise RuntimeError('CANCELLED')
+    var_0 = np.random.rand(1, 3)
+
+    if not running:
+        raise RuntimeError('CANCELLED')
+    var_1 = pd.DataFrame(var_0, index=const_0, columns=['A', 'B', 'C'])
+
+    del var_0
+    if not running:
+        raise RuntimeError('CANCELLED')
+    var_2 = var_1.transform({{'A': {udf_1}, 'B': [{udf_1}, {udf_2}]}}, axis=0)
+
+    del var_1
+    """
+    assert generated.code == textwrap.dedent(expected_contents).strip()
+
+
+def test_codegen_with_udf_and_args(codegen):
+    df = md.DataFrame(mt.random.rand(1, 3), columns=list("ABC"))
+
+    def f1(x, y, a, b) -> int:
+        return x + y + a + b
+
+    result = df.transform({"A": f1}, 0, 2, a=3, b=4)
+    dag = result.build_graph()
+    generated = codegen.generate(dag)
+
+    udf_1_body = base64.b64encode(
+        pickle.dumps(
+            bind_func_args_from_pos(f1, 1, 2, a=3, b=4),
+            protocol=pickle.DEFAULT_PROTOCOL,
+        )
+    )
+    udf_1_value = (
+        f"pickled_data = cloudpickle.loads(base64.b64decode({udf_1_body}), buffers=[])"
+    )
+    udf_1 = f"user_udf_f1_{hashlib.md5(udf_1_value.encode('utf-8')).hexdigest()}"
+
+    expected_contents = f"""
+    import numpy as np
+    import pandas as pd
+    import base64
+    import cloudpickle
+    {udf_1_value}
+    {udf_1} = pickled_data
+    if not running:
+        raise RuntimeError('CANCELLED')
+    var_0 = np.random.rand(1, 3)
+
+    if not running:
+        raise RuntimeError('CANCELLED')
+    var_1 = pd.DataFrame(var_0, index=const_0, columns=['A', 'B', 'C'])
+
+    del var_0
+    if not running:
+        raise RuntimeError('CANCELLED')
+    var_2 = var_1.transform({{'A': {udf_1}}}, axis=0)
+
+    del var_1
+    """
+    assert generated.code == textwrap.dedent(expected_contents).strip()

maxframe/codegen/spe/utils.py

@@ -0,0 +1,54 @@
+# Copyright 1999-2025 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 typing import List, Type
+
+from ...typing_ import OperatorType
+from .core import SPECodeContext, SPEOperatorAdapter, register_op_adapter
+
+_import_aliases = {
+    "np": "numpy",
+    "npl": "numpy.linalg",
+    "pd": "pandas",
+    "spl": "scipy.linalg",
+}
+
+
+def build_method_call_adapter(
+    op_cls: Type[OperatorType],
+    method_name: str,
+    *call_args,
+    skip_none: bool = True,
+    source_module: str = None,
+    **call_kwargs,
+) -> Type[SPEOperatorAdapter]:
+    class MethodCallAdapter(SPEOperatorAdapter):
+        def generate_code(self, op: OperatorType, context: SPECodeContext) -> List[str]:
+            if source_module in _import_aliases:
+                context.register_import(_import_aliases[source_module], source_module)
+
+            input_var_name = source_module or context.get_input_tileable_variable(
+                op.inputs[0]
+            )
+            res_var_names = [
+                context.get_output_tileable_variable(out) for out in op.outputs
+            ]
+            args_str = self.generate_call_args_with_attributes(
+                op, context, *call_args, skip_none=skip_none, **call_kwargs
+            )
+            res_var_names_str = ", ".join(res_var_names)
+            return [f"{res_var_names_str} = {input_var_name}.{method_name}({args_str})"]
+
+    MethodCallAdapter.__name__ = op_cls.__name__ + "Adapter"
+    return register_op_adapter(op_cls)(MethodCallAdapter)

maxframe/codegen/tests/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 1999-2025 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/{tests → codegen/tests}/test_codegen.py

@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 # Copyright 1999-2025 Alibaba Group Holding Ltd.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -16,12 +15,11 @@
 import base64
 from typing import List, Tuple
 
-# 使用pytest生成单元测试
 import pytest
 
-from maxframe.codegen import UserCodeMixin
-from maxframe.lib import wrapped_pickle
-from maxframe.serialization.core import PickleContainer
+from ...lib import wrapped_pickle
+from ...serialization.core import PickleContainer
+from ..core import UserCodeMixin
 
 
 @pytest.mark.parametrize(

maxframe/config/__init__.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .config import AttributeDict, option_context, options, update_wlm_quota_settings
+from .config import AttributeDict, option_context, options

maxframe/config/config.py

@@ -14,6 +14,7 @@
 
 import contextlib
 import contextvars
+import os
 import traceback
 import warnings
 from copy import deepcopy
@@ -28,14 +29,14 @@ except ImportError:
 
     available_timezones = lambda: all_timezones
 
-import logging
-
 from ..utils import get_python_tag
 from .validators import (
     ValidatorType,
     all_validator,
+    is_all_dict_keys_in,
     is_bool,
     is_dict,
+    is_float,
     is_in,
     is_integer,
     is_non_negative_integer,
@@ -45,11 +46,11 @@ from .validators import (
     is_valid_cache_path,
 )
 
-logger = logging.getLogger(__name__)
-
 _DEFAULT_REDIRECT_WARN = "Option {source} has been replaced by {target} and might be removed in a future release."
 _DEFAULT_MAX_ALIVE_SECONDS = 3 * 24 * 3600
 _DEFAULT_MAX_IDLE_SECONDS = 3600
+_DEFAULT_RETRY_DELAY = 0.1
+_DEFAULT_RETRY_TIMES = 4
 _DEFAULT_SPE_OPERATION_TIMEOUT_SECONDS = 120
 _DEFAULT_SPE_FAILURE_RETRY_TIMES = 5
 _DEFAULT_UPLOAD_BATCH_SIZE = 4096
@@ -57,6 +58,7 @@ _DEFAULT_TEMP_LIFECYCLE = 1
 _DEFAULT_TASK_START_TIMEOUT = 60
 _DEFAULT_TASK_RESTART_TIMEOUT = 300
 _DEFAULT_LOGVIEW_HOURS = 24 * 30
+_DEFAULT_FUNCTION_RUNNING_OPTIONS = {"cpu": 1, "memory": "4GiB", "gpu": 0}
 
 
 class OptionError(Exception):
@@ -321,6 +323,12 @@ def _get_legal_local_tz_name() -> Optional[str]:
 
 
 default_options = Config()
+default_options.register_option(
+    "retry_times", _DEFAULT_RETRY_TIMES, validator=is_integer
+)
+default_options.register_option(
+    "retry_delay", _DEFAULT_RETRY_DELAY, validator=is_integer | is_float
+)
 default_options.register_option(
     "execution_mode", "trigger", validator=is_in(["trigger", "eager"])
 )
@@ -368,6 +376,14 @@ default_options.register_option(
 default_options.register_option(
     "chunk_size", None, validator=is_null | is_integer, remote=True
 )
+default_options.register_option(
+    "chunk_store_limit", 128 * 1024**2, validator=is_numeric
+)
+
+default_options.register_option("local_execution.enabled", False, validator=is_bool)
+default_options.register_option(
+    "local_execution.size_limit", 5 * 1024**2, validator=is_integer
+)
 
 default_options.register_option(
     "session.max_alive_seconds",
@@ -425,6 +441,9 @@ default_options.register_option(
     validator=is_null | is_integer,
     remote=True,
 )
+default_options.register_option(
+    "session.client_version", None, validator=is_null | is_string, remote=True
+)
 
 default_options.register_option("warn_duplicated_execution", False, validator=is_bool)
 default_options.register_option("dataframe.use_arrow_dtype", True, validator=is_bool)
@@ -439,6 +458,23 @@ default_options.register_option(
     "dag.settings", value=dict(), validator=is_dict, remote=True
 )
 
+default_options.register_option(
+    "function.default_running_options",
+    _DEFAULT_FUNCTION_RUNNING_OPTIONS,
+    validator=is_dict | is_all_dict_keys_in("cpu", "memory", "gpu"),
+)
+
+################
+# DPE Settings #
+################
+default_options.register_option("dpe.settings", dict(), validator=is_dict, remote=True)
+default_options.register_option(
+    "dpe.task.settings", dict(), validator=is_dict, remote=True
+)
+default_options.register_option(
+    "dpe.reduction.combine_size", 4, validator=is_integer, remote=True
+)
+
 ################
 # SPE Settings #
 ################
@@ -463,6 +499,16 @@ default_options.register_option(
     "pythonpack.task.settings", {}, validator=is_dict, remote=True
 )
 
+##################
+# Learn settings #
+##################
+assume_finite = os.environ.get("SKLEARN_ASSUME_FINITE")
+if assume_finite is not None:
+    assume_finite = bool(assume_finite)
+default_options.register_option(
+    "learn.assume_finite", assume_finite, validator=is_null | is_bool
+)
+
 _options_ctx_var = contextvars.ContextVar("_options_ctx_var")
 
 
@@ -518,22 +564,3 @@ class OptionsProxy:
 
 
 options = OptionsProxy()
-
-
-def update_wlm_quota_settings(session_id: str, engine_settings: Dict[str, Any]):
-    engine_quota = engine_settings.get("odps.task.wlm.quota", None)
-    session_quota = options.session.quota_name or None
-    if engine_quota != session_quota and engine_quota:
-        logger.warning(
-            "[Session=%s] Session quota (%s) is different to SubDag engine quota (%s)",
-            session_id,
-            session_quota,
-            engine_quota,
-        )
-        # TODO(renxiang): overwrite or not overwrite
-        return
-
-    if session_quota:
-        engine_settings["odps.task.wlm.quota"] = session_quota
-    elif "odps.task.wlm.quota" in engine_settings:
-        engine_settings.pop("odps.task.wlm.quota")

maxframe/config/tests/test_config.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright 1999-2025 Alibaba Group Holding Ltd.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -18,14 +16,8 @@ import threading
 
 import pytest
 
-from ..config import (
-    Config,
-    is_integer,
-    is_string,
-    option_context,
-    options,
-    update_wlm_quota_settings,
-)
+from ...utils import update_wlm_quota_settings
+from ..config import Config, is_integer, is_string, option_context, options
 
 
 def test_config_context():
@@ -117,6 +109,6 @@ def test_update_wlm_quota_settings():
         update_wlm_quota_settings("session_id", engine_settings)
         assert engine_settings["odps.task.wlm.quota"] == "quota1"
         options.session.quota_name = None
-        update_wlm_quota_settings("session_id", engine_settings)
-        # TODO(renxiang): overwrite or not overwrite
+        with pytest.raises(ValueError):
+            update_wlm_quota_settings("session_id", engine_settings)
         assert "odps.task.wlm.quota" in engine_settings

maxframe/config/validators.py

@@ -64,6 +64,11 @@ def is_in(vals):
     return Validator(vals.__contains__)
 
 
+def is_all_dict_keys_in(*keys):
+    keys_set = set(keys)
+    return Validator(lambda x: x in keys_set)
+
+
 def _is_valid_cache_path(path: str) -> bool:
     """
     path should look like oss://oss_endpoint/oss_bucket/path

maxframe/conftest.py

@@ -35,7 +35,7 @@ def test_config():
     return config
 
 
-def _get_odps_env(test_config: ConfigParser, section_name: str) -> ODPS:
+def _get_account_env(test_config: ConfigParser, section_name: str) -> ODPS:
     try:
         access_id = test_config.get(section_name, "access_id")
     except NoOptionError:
@@ -60,10 +60,18 @@ def _get_odps_env(test_config: ConfigParser, section_name: str) -> ODPS:
         tunnel_endpoint = test_config.get("odps", "tunnel_endpoint")
     except NoOptionError:
         tunnel_endpoint = None
-
-    entry = ODPS(
-        access_id, secret_access_key, project, endpoint, overwrite_global=False
+    return ODPS(
+        access_id,
+        secret_access_key,
+        project,
+        endpoint,
+        tunnel_endpoint=tunnel_endpoint,
+        overwrite_global=False,
     )
+
+
+def _get_bearer_token_env(test_config: ConfigParser, section_name: str) -> ODPS:
+    entry = _get_account_env(test_config, section_name)
     policy = {
         "Version": "1",
         "Statement": [
@@ -73,27 +81,28 @@ def _get_odps_env(test_config: ConfigParser, section_name: str) -> ODPS:
     token = entry.get_project().generate_auth_token(policy, "bearer", 5)
     return ODPS(
         account=BearerTokenAccount(token, 5),
-        project=project,
-        endpoint=endpoint,
-        tunnel_endpoint=tunnel_endpoint,
+        project=entry.project,
+        endpoint=entry.endpoint,
+        tunnel_endpoint=entry.tunnel_endpoint,
     )
 
 
 @pytest.fixture(scope="session")
 def odps_with_schema(test_config):
     try:
-        return _get_odps_env(test_config, "odps_with_schema")
+        return _get_bearer_token_env(test_config, "odps_with_schema")
     except NoSectionError:
         pytest.skip("Need to specify odps_with_schema section in test.conf")
 
 
 @pytest.fixture(scope="session", autouse=True)
 def odps_envs(test_config):
-    entry = _get_odps_env(test_config, "odps")
+    entry = _get_bearer_token_env(test_config, "odps")
 
     os.environ["ODPS_BEARER_TOKEN"] = entry.account.token
     os.environ["ODPS_PROJECT_NAME"] = entry.project
     os.environ["ODPS_ENDPOINT"] = entry.endpoint
+    os.environ["RAY_ISOLATION_UT_ENV"] = "UT"
     if entry.tunnel_endpoint:
         os.environ["ODPS_TUNNEL_ENDPOINT"] = entry.tunnel_endpoint
 
@@ -104,6 +113,7 @@ def odps_envs(test_config):
         os.environ.pop("ODPS_PROJECT_NAME", None)
         os.environ.pop("ODPS_ENDPOINT", None)
         os.environ.pop("ODPS_TUNNEL_ENDPOINT", None)
+        os.environ.pop("RAY_ISOLATION_UT_ENV", None)
 
         from .tests.utils import _test_tables_to_drop
 
@@ -114,6 +124,11 @@ def odps_envs(test_config):
                 pass
 
 
+@pytest.fixture(scope="session")
+def odps_account(test_config):
+    return _get_account_env(test_config, "odps")
+
+
 @pytest.fixture(scope="session")
 def oss_config():
     config = ConfigParser()
@@ -155,7 +170,7 @@ def oss_config():
         config.oss_rolearn = oss_rolearn
         yield config
     except (NoSectionError, NoOptionError, ImportError):
-        return None
+        yield None
     finally:
         options.service_role_arn = old_role_arn
         options.object_cache_url = old_cache_url
@@ -194,3 +209,16 @@ def local_test_envs():
         os.environ[spe_launcher_env] = old_value
     else:
         del os.environ[spe_launcher_env]
+
+
+@pytest.fixture
+def enable_local_execution(request):
+    old_enabled = options.local_execution.enabled
+    old_limit = options.local_execution.size_limit
+    try:
+        options.local_execution.enabled = True
+        options.local_execution.size_limit = getattr(request, "param", 0)
+        yield
+    finally:
+        options.local_execution.enabled = old_enabled
+        options.local_execution.size_limit = old_limit

maxframe/core/__init__.py

@@ -32,6 +32,7 @@ from .entity import (
     Tileable,
     TileableData,
     _ExecuteAndFetchMixin,
+    df_output_types,
     get_fetch_class,
     get_output_types,
     get_tileable_types,

maxframe/core/context.py

@@ -0,0 +1,110 @@
+# Copyright 1999-2025 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 abc import ABC, abstractmethod
+from typing import Any, List
+
+from ..utils import classproperty
+
+
+class Context(ABC):
+    """
+    Context that providing API that can be
+    used inside `tile` and `execute`.
+    """
+
+    all_contexts = []
+
+    @abstractmethod
+    def get_session_id(self) -> str:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_subdag_id(self) -> str:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_local_host_ip(self) -> str:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_worker_cores(self) -> int:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_chunks_result(self, data_keys: List[str], copy: bool = True) -> List[Any]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def create_operator_controller(self, name: str, object_cls, *args, **kwargs):
+        """
+        Create operator controller.
+
+        Parameters
+        ----------
+        name : str
+            Object name.
+        object_cls
+            Object class.
+        args
+        kwargs
+
+        Returns
+        -------
+        ref
+        """
+
+    @abstractmethod
+    def get_operator_controller(self, name: str):
+        """
+        Get remote object
+
+        Parameters
+        ----------
+        name : str
+            Object name.
+
+        Returns
+        -------
+        ref
+        """
+
+    @abstractmethod
+    def destroy_operator_controller(self, name: str):
+        """
+        Destroy remote object.
+
+        Parameters
+        ----------
+        name : str
+            Object name.
+        """
+
+    def __enter__(self):
+        Context.all_contexts.append(self)
+
+    def __exit__(self, *_):
+        Context.all_contexts.pop()
+
+    @classproperty
+    def current(cls):
+        return cls.all_contexts[-1] if cls.all_contexts else None
+
+
+def set_context(context: Context):
+    Context.all_contexts.append(context)
+
+
+def get_context() -> Context:
+    return Context.current

maxframe/core/entity/__init__.py

@@ -17,6 +17,7 @@ from .executable import ExecutableTuple, _ExecuteAndFetchMixin
 from .objects import OBJECT_TYPE, Object, ObjectData
 from .output_types import (
     OutputType,
+    df_output_types,
     get_fetch_class,
     get_output_types,
     get_tileable_types,

maxframe/core/entity/core.py

@@ -129,13 +129,6 @@ class Entity(Serializable):
     def copy_from(self, obj):
         self.data = obj.data
 
-    def tiles(self):
-        from .tileables import handler
-
-        new_entity = self.copy()
-        new_entity.data = handler.tiles(self.data)
-        return new_entity
-
     def __getattr__(self, attr):
         return getattr(self._data, attr)