maxframe 0.1.0b4__cp311-cp311-win_amd64.whl → 1.0.0__cp311-cp311-win_amd64.whl

maxframe/__init__.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from . import dataframe, learn, remote, tensor
+from .config import options
 from .session import execute, fetch, new_session, stop_server
 
 

maxframe/codegen.py

@@ -16,6 +16,7 @@ import abc
 import base64
 import dataclasses
 import logging
+from collections import defaultdict
 from enum import Enum
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
 
@@ -25,14 +26,14 @@ from odps.utils import camel_to_underline
 from .core import OperatorType, Tileable, TileableGraph
 from .core.operator import Fetch
 from .extension import iter_extensions
+from .io.odpsio import build_dataframe_table_meta
+from .io.odpsio.schema import pandas_to_odps_schema
 from .lib import wrapped_pickle as pickle
-from .odpsio import build_dataframe_table_meta
-from .odpsio.schema import pandas_to_odps_schema
 from .protocol import DataFrameTableMeta, ResultInfo
 from .serialization import PickleContainer
 from .serialization.serializables import Serializable, StringField
 from .typing_ import PandasObjectTypes
-from .udf import MarkedFunction
+from .udf import MarkedFunction, PythonPackOptions
 
 if TYPE_CHECKING:
     from odpsctx import ODPSSessionContext
@@ -75,8 +76,18 @@ class AbstractUDF(Serializable):
     def unregister(self, odps: "ODPSSessionContext"):
         raise NotImplementedError
 
+    @abc.abstractmethod
+    def collect_pythonpack(self) -> List[PythonPackOptions]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def load_pythonpack_resources(self, odps_ctx: "ODPSSessionContext") -> None:
+        raise NotImplementedError
+
 
 class UserCodeMixin:
+    __slots__ = ()
+
     @classmethod
     def obj_to_python_expr(cls, obj: Any = None) -> str:
         """
@@ -335,6 +346,9 @@ def register_engine_codegen(type_: Type["BigDagCodeGenerator"]):
 BUILTIN_ENGINE_SPE = "SPE"
 BUILTIN_ENGINE_MCSQL = "MCSQL"
 
+FAST_RANGE_INDEX_ENABLED = "codegen.fast_range_index_enabled"
+ROW_NUMBER_WINDOW_INDEX_ENABLED = "codegen.row_number_window_index_enabled"
+
 
 class BigDagCodeGenerator(metaclass=abc.ABCMeta):
     _context: BigDagCodeContext
@@ -469,14 +483,51 @@ class BigDagCodeGenerator(metaclass=abc.ABCMeta):
             output_key_to_result_infos=self._context.get_tileable_result_infos(),
         )
 
+    def run_pythonpacks(
+        self,
+        odps_ctx: "ODPSSessionContext",
+        python_tag: str,
+        is_production: bool = False,
+        schedule_id: Optional[str] = None,
+        hints: Optional[dict] = None,
+        priority: Optional[int] = None,
+    ) -> Dict[str, PythonPackOptions]:
+        key_to_packs = defaultdict(list)
+        for udf in self._context.get_udfs():
+            for pack in udf.collect_pythonpack():
+                key_to_packs[pack.key].append(pack)
+        distinct_packs = []
+        for packs in key_to_packs.values():
+            distinct_packs.append(packs[0])
+
+        inst_id_to_req = {}
+        for pack in distinct_packs:
+            inst = odps_ctx.run_pythonpack(
+                requirements=pack.requirements,
+                prefer_binary=pack.prefer_binary,
+                pre_release=pack.pre_release,
+                force_rebuild=pack.force_rebuild,
+                no_audit_wheel=pack.no_audit_wheel,
+                python_tag=python_tag,
+                is_production=is_production,
+                schedule_id=schedule_id,
+                hints=hints,
+                priority=priority,
+            )
+            # fulfill instance id of pythonpacks with same keys
+            for same_pack in key_to_packs[pack.key]:
+                same_pack.pack_instance_id = inst.id
+            inst_id_to_req[inst.id] = pack
+        return inst_id_to_req
+
     def register_udfs(self, odps_ctx: "ODPSSessionContext"):
         for udf in self._context.get_udfs():
-            logger.info("[Session %s] Registering UDF %s", self._session_id, udf.name)
+            logger.info("[Session=%s] Registering UDF %s", self._session_id, udf.name)
             udf.register(odps_ctx, True)
 
     def unregister_udfs(self, odps_ctx: "ODPSSessionContext"):
         for udf in self._context.get_udfs():
-            logger.info("[Session %s] Unregistering UDF %s", self._session_id, udf.name)
+            logger.info("[Session=%s] Unregistering UDF %s", self._session_id, udf.name)
             udf.unregister(odps_ctx)
 
     def get_udfs(self) -> List[AbstractUDF]:

maxframe/config/config.py

@@ -19,27 +19,40 @@ import warnings
 from copy import deepcopy
 from typing import Any, Dict, Optional, Union
 
+from odps.lib import tzlocal
+
+try:
+    from zoneinfo import available_timezones
+except ImportError:
+    from pytz import all_timezones
+
+    available_timezones = lambda: all_timezones
+
 from ..utils import get_python_tag
 from .validators import (
     ValidatorType,
     all_validator,
-    any_validator,
     is_bool,
     is_dict,
     is_in,
     is_integer,
+    is_non_negative_integer,
     is_null,
     is_numeric,
     is_string,
+    is_valid_cache_path,
 )
 
 _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_SPE_OPERATION_TIMEOUT_SECONDS = 120
+_DEFAULT_SPE_FAILURE_RETRY_TIMES = 5
 _DEFAULT_UPLOAD_BATCH_SIZE = 4096
 _DEFAULT_TEMP_LIFECYCLE = 1
 _DEFAULT_TASK_START_TIMEOUT = 60
+_DEFAULT_TASK_RESTART_TIMEOUT = 300
+_DEFAULT_LOGVIEW_HOURS = 24 * 30
 
 
 class OptionError(Exception):
@@ -295,23 +308,60 @@ class Config:
         return {k: v for k, v in res.items() if k in self._remote_options}
 
 
-default_options = Config()
+def _get_legal_local_tz_name() -> Optional[str]:
+    """Sometimes we may get illegal tz name from tzlocal.get_localzone()"""
+    tz_name = str(tzlocal.get_localzone())
+    if tz_name not in available_timezones():
+        return None
+    return tz_name
+
 
+default_options = Config()
 default_options.register_option(
     "execution_mode", "trigger", validator=is_in(["trigger", "eager"])
 )
+default_options.register_option("use_common_table", False, validator=is_bool)
 default_options.register_option(
     "python_tag", get_python_tag(), validator=is_string, remote=True
 )
+default_options.register_option(
+    "local_timezone",
+    _get_legal_local_tz_name(),
+    validator=is_null | is_in(set(available_timezones())),
+    remote=True,
+)
+default_options.register_option(
+    "session.logview_hours", _DEFAULT_LOGVIEW_HOURS, validator=is_integer, remote=True
+)
 default_options.register_option(
     "client.task_start_timeout", _DEFAULT_TASK_START_TIMEOUT, validator=is_integer
 )
+default_options.register_option(
+    "client.task_restart_timeout", _DEFAULT_TASK_RESTART_TIMEOUT, validator=is_integer
+)
 default_options.register_option("sql.enable_mcqa", True, validator=is_bool, remote=True)
 default_options.register_option(
     "sql.generate_comments", True, validator=is_bool, remote=True
 )
+default_options.register_option(
+    "sql.auto_use_common_image", True, validator=is_bool, remote=True
+)
 default_options.register_option("sql.settings", {}, validator=is_dict, remote=True)
 
+default_options.register_option("is_production", False, validator=is_bool, remote=True)
+default_options.register_option("schedule_id", "", validator=is_string, remote=True)
+
+default_options.register_option(
+    "service_role_arn", None, validator=is_null | is_string, remote=True
+)
+default_options.register_option(
+    "object_cache_url", None, validator=is_null | is_valid_cache_path, remote=True
+)
+
+default_options.register_option(
+    "chunk_size", None, validator=is_null | is_integer, remote=True
+)
+
 default_options.register_option(
     "session.max_alive_seconds",
     _DEFAULT_MAX_ALIVE_SECONDS,
@@ -324,15 +374,25 @@ default_options.register_option(
     validator=is_numeric,
     remote=True,
 )
+default_options.register_option(
+    "session.quota_name", None, validator=is_null | is_string, remote=True
+)
+default_options.register_option(
+    "session.enable_schema", None, validator=is_null | is_bool, remote=True
+)
+default_options.register_option(
+    "session.enable_high_availability", None, validator=is_null | is_bool, remote=True
+)
+default_options.register_option(
+    "session.default_schema", None, validator=is_null | is_string, remote=True
+)
 default_options.register_option(
     "session.upload_batch_size",
     _DEFAULT_UPLOAD_BATCH_SIZE,
     validator=is_integer,
 )
 default_options.register_option(
-    "session.table_lifecycle",
-    None,
-    validator=any_validator(is_null, is_integer),
+    "session.table_lifecycle", None, validator=is_null | is_integer, remote=True
 )
 default_options.register_option(
     "session.temp_table_lifecycle",
@@ -343,7 +403,7 @@ default_options.register_option(
 default_options.register_option(
     "session.subinstance_priority",
     None,
-    validator=any_validator(is_null, is_integer),
+    validator=is_null | is_integer,
     remote=True,
 )
 
@@ -355,9 +415,7 @@ default_options.register_option(
 default_options.register_option(
     "optimize.head_optimize_threshold", 1000, validator=is_integer
 )
-default_options.register_option(
-    "show_progress", "auto", validator=any_validator(is_bool, is_string)
-)
+default_options.register_option("show_progress", "auto", validator=is_bool | is_string)
 default_options.register_option(
     "dag.settings", value=dict(), validator=is_dict, remote=True
 )
@@ -368,7 +426,13 @@ default_options.register_option(
 default_options.register_option(
     "spe.operation_timeout_seconds",
     _DEFAULT_SPE_OPERATION_TIMEOUT_SECONDS,
-    validator=is_integer,
+    validator=is_non_negative_integer,
+    remote=True,
+)
+default_options.register_option(
+    "spe.failure_retry_times",
+    _DEFAULT_SPE_FAILURE_RETRY_TIMES,
+    validator=is_non_negative_integer,
     remote=True,
 )
 
@@ -376,6 +440,10 @@ default_options.register_option(
     "spe.task.settings", dict(), validator=is_dict, remote=True
 )
 
+default_options.register_option(
+    "pythonpack.task.settings", {}, validator=is_dict, remote=True
+)
+
 _options_ctx_var = contextvars.ContextVar("_options_ctx_var")
 
 

maxframe/config/validators.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from typing import Callable
+from urllib.parse import urlparse
 
 ValidatorType = Callable[..., bool]
 
@@ -32,21 +33,51 @@ def all_validator(*validators: ValidatorType):
     return validate
 
 
-is_null = lambda x: x is None
-is_bool = lambda x: isinstance(x, bool)
-is_float = lambda x: isinstance(x, float)
-is_integer = lambda x: isinstance(x, int)
-is_numeric = lambda x: isinstance(x, (int, float))
-is_string = lambda x: isinstance(x, str)
-is_dict = lambda x: isinstance(x, dict)
-is_positive_integer = lambda x: is_integer(x) and x > 0
+class Validator:
+    def __init__(self, func: ValidatorType):
+        self._func = func
+
+    def __call__(self, arg) -> bool:
+        return self._func(arg)
+
+    def __or__(self, other):
+        return OrValidator(self, other)
+
+
+class OrValidator(Validator):
+    def __init__(self, lhs: Validator, rhs: Validator):
+        super().__init__(lambda x: lhs(x) or rhs(x))
+
+
+is_null = Validator(lambda x: x is None)
+is_bool = Validator(lambda x: isinstance(x, bool))
+is_float = Validator(lambda x: isinstance(x, float))
+is_integer = Validator(lambda x: isinstance(x, int))
+is_numeric = Validator(lambda x: isinstance(x, (int, float)))
+is_string = Validator(lambda x: isinstance(x, str))
+is_dict = Validator(lambda x: isinstance(x, dict))
+is_positive_integer = Validator(lambda x: is_integer(x) and x > 0)
+is_non_negative_integer = Validator(lambda x: is_integer(x) and x >= 0)
 
 
 def is_in(vals):
-    def validate(x):
-        return x in vals
+    return Validator(vals.__contains__)
 
-    return validate
+
+def _is_valid_cache_path(path: str) -> bool:
+    """
+    path should look like oss://oss_endpoint/oss_bucket/path
+    """
+    parsed_url = urlparse(path)
+    return (
+        parsed_url.scheme == "oss"
+        and parsed_url.netloc
+        and parsed_url.path
+        and "/" in parsed_url.path
+    )
+
+
+is_valid_cache_path = Validator(_is_valid_cache_path)
 
 
 _invalid_char_in_yaml_str = {'"', "'", "\n", "\\"}

maxframe/conftest.py

@@ -14,10 +14,13 @@
 
 import faulthandler
 import os
-from configparser import ConfigParser, NoOptionError
+from configparser import ConfigParser, NoOptionError, NoSectionError
 
 import pytest
 from odps import ODPS
+from odps.accounts import BearerTokenAccount
+
+from .config import options
 
 faulthandler.enable(all_threads=True)
 _test_conf_file_name = os.path.join(
@@ -32,12 +35,23 @@ def test_config():
     return config
 
 
-@pytest.fixture(scope="session", autouse=True)
-def odps_envs(test_config):
-    access_id = test_config.get("odps", "access_id")
-    secret_access_key = test_config.get("odps", "secret_access_key")
-    project = test_config.get("odps", "project")
-    endpoint = test_config.get("odps", "endpoint")
+def _get_odps_env(test_config: ConfigParser, section_name: str) -> ODPS:
+    try:
+        access_id = test_config.get(section_name, "access_id")
+    except NoOptionError:
+        access_id = test_config.get("odps", "access_id")
+    try:
+        secret_access_key = test_config.get(section_name, "secret_access_key")
+    except NoOptionError:
+        secret_access_key = test_config.get("odps", "secret_access_key")
+    try:
+        project = test_config.get(section_name, "project")
+    except NoOptionError:
+        project = test_config.get("odps", "project")
+    try:
+        endpoint = test_config.get(section_name, "endpoint")
+    except NoOptionError:
+        endpoint = test_config.get("odps", "endpoint")
     try:
         tunnel_endpoint = test_config.get("odps", "tunnel_endpoint")
     except NoOptionError:
@@ -53,12 +67,31 @@ def odps_envs(test_config):
         ],
     }
     token = entry.get_project().generate_auth_token(policy, "bearer", 5)
+    return ODPS(
+        account=BearerTokenAccount(token, 5),
+        project=project,
+        endpoint=endpoint,
+        tunnel_endpoint=tunnel_endpoint,
+    )
 
-    os.environ["ODPS_BEARER_TOKEN"] = token
-    os.environ["ODPS_PROJECT_NAME"] = project
-    os.environ["ODPS_ENDPOINT"] = endpoint
-    if tunnel_endpoint:
-        os.environ["ODPS_TUNNEL_ENDPOINT"] = tunnel_endpoint
+
+@pytest.fixture(scope="session")
+def odps_with_schema(test_config):
+    try:
+        return _get_odps_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")
+
+    os.environ["ODPS_BEARER_TOKEN"] = entry.account.token
+    os.environ["ODPS_PROJECT_NAME"] = entry.project
+    os.environ["ODPS_ENDPOINT"] = entry.endpoint
+    if entry.tunnel_endpoint:
+        os.environ["ODPS_TUNNEL_ENDPOINT"] = entry.tunnel_endpoint
 
     try:
         yield
@@ -77,16 +110,23 @@ def odps_envs(test_config):
                 pass
 
 
-@pytest.fixture
+@pytest.fixture(scope="session")
 def oss_config():
     config = ConfigParser()
     config.read(_test_conf_file_name)
 
+    old_role_arn = options.service_role_arn
+    old_cache_url = options.object_cache_url
+
     try:
         oss_access_id = config.get("oss", "access_id")
         oss_secret_access_key = config.get("oss", "secret_access_key")
         oss_bucket_name = config.get("oss", "bucket_name")
         oss_endpoint = config.get("oss", "endpoint")
+        oss_rolearn = config.get("oss", "rolearn")
+
+        options.service_role_arn = oss_rolearn
+        options.object_cache_url = f"oss://{oss_endpoint}/{oss_bucket_name}"
 
         config.oss_config = (
             oss_access_id,
@@ -99,9 +139,13 @@ def oss_config():
 
         auth = oss2.Auth(oss_access_id, oss_secret_access_key)
         config.oss_bucket = oss2.Bucket(auth, oss_endpoint, oss_bucket_name)
-        return config
+        config.oss_rolearn = oss_rolearn
+        yield config
     except (ConfigParser.NoSectionError, ConfigParser.NoOptionError, ImportError):
         return None
+    finally:
+        options.service_role_arn = old_role_arn
+        options.object_cache_url = old_cache_url
 
 
 @pytest.fixture(autouse=True)

maxframe/core/__init__.py

@@ -14,27 +14,18 @@
 
 # noinspection PyUnresolvedReferences
 from ..typing_ import ChunkType, EntityType, OperatorType, TileableType
-from .base import ExecutionError
+from .base import Base, ExecutionError
 from .entity import (
-    CHUNK_TYPE,
     ENTITY_TYPE,
-    FUSE_CHUNK_TYPE,
-    OBJECT_CHUNK_TYPE,
     OBJECT_TYPE,
     TILEABLE_TYPE,
-    Chunk,
-    ChunkData,
     Entity,
     EntityData,
     ExecutableTuple,
-    FuseChunk,
-    FuseChunkData,
     HasShapeTileable,
     HasShapeTileableData,
     NotSupportTile,
     Object,
-    ObjectChunk,
-    ObjectChunkData,
     ObjectData,
     OutputType,
     Tileable,
@@ -43,23 +34,18 @@ from .entity import (
     get_fetch_class,
     get_output_types,
     get_tileable_types,
-    register,
     register_fetch_class,
     register_output_types,
-    unregister,
 )
 
 # noinspection PyUnresolvedReferences
 from .graph import (
     DAG,
-    ChunkGraph,
-    ChunkGraphBuilder,
     DirectedGraph,
     GraphContainsCycleError,
     GraphSerializer,
     TileableGraph,
     TileableGraphBuilder,
-    TileContext,
-    TileStatus,
 )
 from .mode import enter_mode, is_build_mode, is_eager_mode, is_kernel_mode
+from .operator import build_fetch

maxframe/core/entity/__init__.py

@@ -12,18 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .chunks import CHUNK_TYPE, Chunk, ChunkData
 from .core import ENTITY_TYPE, Entity, EntityData
 from .executable import ExecutableTuple, _ExecuteAndFetchMixin
-from .fuse import FUSE_CHUNK_TYPE, FuseChunk, FuseChunkData
-from .objects import (
-    OBJECT_CHUNK_TYPE,
-    OBJECT_TYPE,
-    Object,
-    ObjectChunk,
-    ObjectChunkData,
-    ObjectData,
-)
+from .objects import OBJECT_TYPE, Object, ObjectData
 from .output_types import (
     OutputType,
     get_fetch_class,
@@ -39,6 +30,4 @@ from .tileables import (
     NotSupportTile,
     Tileable,
     TileableData,
-    register,
-    unregister,
 )

maxframe/core/entity/executable.py

@@ -46,7 +46,7 @@ class DecrefRunner:
                 break
 
             session = session_ref()
-            if session is None:
+            if session is None or session.closed:
                 fut.set_result(None)
                 continue
             try: