maxframe 0.1.0b5__cp311-cp311-macosx_11_0_arm64.whl

maxframe/conftest.py

@@ -0,0 +1,139 @@
+# 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 faulthandler
+import os
+from configparser import ConfigParser, NoOptionError
+
+import pytest
+from odps import ODPS
+
+faulthandler.enable(all_threads=True)
+_test_conf_file_name = os.path.join(
+    os.path.dirname(os.path.abspath(__file__)), "tests", "test.conf"
+)
+
+
+@pytest.fixture(scope="session")
+def test_config():
+    config = ConfigParser()
+    config.read(_test_conf_file_name)
+    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")
+    try:
+        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
+    )
+    policy = {
+        "Version": "1",
+        "Statement": [
+            {"Action": ["odps:*"], "Resource": "acs:odps:*:*", "Effect": "Allow"}
+        ],
+    }
+    token = entry.get_project().generate_auth_token(policy, "bearer", 5)
+
+    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
+
+    try:
+        yield
+    finally:
+        os.environ.pop("ODPS_BEARER_TOKEN", None)
+        os.environ.pop("ODPS_PROJECT_NAME", None)
+        os.environ.pop("ODPS_ENDPOINT", None)
+        os.environ.pop("ODPS_TUNNEL_ENDPOINT", None)
+
+        from .tests.utils import _test_tables_to_drop
+
+        for table_name in _test_tables_to_drop:
+            try:
+                entry.delete_table(table_name, wait=False)
+            except:
+                pass
+
+
+@pytest.fixture
+def oss_config():
+    config = ConfigParser()
+    config.read(_test_conf_file_name)
+
+    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")
+
+        config.oss_config = (
+            oss_access_id,
+            oss_secret_access_key,
+            oss_bucket_name,
+            oss_endpoint,
+        )
+
+        import oss2
+
+        auth = oss2.Auth(oss_access_id, oss_secret_access_key)
+        config.oss_bucket = oss2.Bucket(auth, oss_endpoint, oss_bucket_name)
+        return config
+    except (ConfigParser.NoSectionError, ConfigParser.NoOptionError, ImportError):
+        return None
+
+
+@pytest.fixture(autouse=True)
+def apply_engine_selection(request):
+    try:
+        from maxframe_framedriver.services.analyzer import DagAnalyzer
+    except ImportError:
+        DagAnalyzer = None
+    try:
+        marks = list(request.node.iter_markers())
+        if request.node.parent:
+            marks.extend(request.node.parent.iter_markers())
+        marks = [m for m in marks if m.name == "maxframe_engine"]
+
+        if DagAnalyzer:
+            engines = set()
+            for mark in marks:
+                engines.update(mark.args[0])
+            DagAnalyzer._enabled_engines = set(engines)
+        yield
+    finally:
+        if DagAnalyzer:
+            DagAnalyzer._enabled_engines = set()
+
+
+@pytest.fixture
+def local_test_envs():
+    spe_launcher_env = "MAXFRAME_SPE_LAUNCHER"
+    old_value = os.getenv(spe_launcher_env)
+    os.environ[spe_launcher_env] = "local"
+    yield
+    if old_value is not None:
+        os.environ[spe_launcher_env] = old_value
+    else:
+        del os.environ[spe_launcher_env]

maxframe/core/__init__.py

@@ -0,0 +1,65 @@
+# 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.
+
+# noinspection PyUnresolvedReferences
+from ..typing_ import ChunkType, EntityType, OperatorType, TileableType
+from .base import 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,
+    TileableData,
+    _ExecuteAndFetchMixin,
+    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

maxframe/core/base.py

@@ -0,0 +1,156 @@
+# 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 typing import Dict, Optional, Tuple, Type
+
+from odps.errors import ODPSError
+
+from ..serialization.serializables import Serializable, StringField
+from ..serialization.serializables.core import SerializableSerializer
+from ..utils import tokenize
+
+
+class Base(Serializable):
+    _no_copy_attrs_ = {"_id"}
+    _init_update_key_ = True
+
+    _key = StringField("key", default=None)
+    _id = StringField("id")
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if self._init_update_key_ and (not hasattr(self, "_key") or not self._key):
+            self._update_key()
+        if not hasattr(self, "_id") or not self._id:
+            self._id = str(id(self))
+
+    @property
+    def _keys_(self):
+        cls = type(self)
+        member = "__keys_" + cls.__name__
+        try:
+            return getattr(cls, member)
+        except AttributeError:
+            slots = sorted(self._FIELDS)
+            setattr(cls, member, slots)
+            return slots
+
+    @property
+    def _copy_tags_(self):
+        cls = type(self)
+        member = f"__copy_tags_{cls.__name__}"
+        try:
+            return getattr(cls, member)
+        except AttributeError:
+            slots = sorted(
+                f.name for k, f in self._FIELDS.items() if k not in self._no_copy_attrs_
+            )
+            setattr(cls, member, slots)
+            return slots
+
+    @property
+    def _values_(self):
+        values = []
+        fields = self._FIELDS
+        for k in self._copy_tags_:
+            try:
+                values.append(fields[k].get(self))
+            except AttributeError:
+                values.append(None)
+        return values
+
+    def __maxframe_tokenize__(self):
+        try:
+            return self._key
+        except AttributeError:  # pragma: no cover
+            self._update_key()
+            return self._key
+
+    def _obj_set(self, k, v):
+        object.__setattr__(self, k, v)
+
+    def _update_key(self):
+        self._obj_set("_key", tokenize(type(self).__name__, *self._values_))
+        return self
+
+    def reset_key(self):
+        self._obj_set("_key", None)
+        return self
+
+    def __copy__(self):
+        return self.copy()
+
+    def copy(self):
+        return self.copy_to(type(self)(_key=self.key))
+
+    def copy_to(self, target: "Base"):
+        target_fields = target._FIELDS
+        no_copy_attrs = self._no_copy_attrs_
+        for k, field in self._FIELDS.items():
+            if k in no_copy_attrs:
+                continue
+            try:
+                # Slightly faster than getattr.
+                value = field.__get__(self, k)
+                try:
+                    target_fields[k].set(target, value)
+                except KeyError:
+                    target_fields["_" + k].set(target, value)
+            except AttributeError:
+                continue
+
+        return target
+
+    def copy_from(self, obj):
+        obj.copy_to(self)
+
+    @property
+    def key(self):
+        return self._key
+
+    @property
+    def id(self):
+        return self._id
+
+    def to_kv(
+        self,
+        exclude_fields: Tuple[str],
+        accept_value_types: Optional[Tuple[Type]] = None,
+    ):
+        fields = self._FIELDS
+        kv = {}
+        no_value = object()
+        for name, field in fields.items():
+            if name not in exclude_fields:
+                value = getattr(self, name, no_value)
+                if value is not no_value and (
+                    accept_value_types is None or isinstance(value, accept_value_types)
+                ):
+                    kv[field.tag] = value
+        return kv
+
+
+class BaseSerializer(SerializableSerializer):
+    def serial(self, obj: Base, context: Dict):
+        return super().serial(obj, context)
+
+
+BaseSerializer.register(Base)
+
+
+class ExecutionError(ODPSError):
+    def __init__(self, nested_error: BaseException):
+        super().__init__(nested_error)
+        self.nested_error = nested_error

maxframe/core/entity/__init__.py

@@ -0,0 +1,44 @@
+# 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 .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 .output_types import (
+    OutputType,
+    get_fetch_class,
+    get_output_types,
+    get_tileable_types,
+    register_fetch_class,
+    register_output_types,
+)
+from .tileables import (
+    TILEABLE_TYPE,
+    HasShapeTileable,
+    HasShapeTileableData,
+    NotSupportTile,
+    Tileable,
+    TileableData,
+    register,
+    unregister,
+)

maxframe/core/entity/chunks.py

@@ -0,0 +1,68 @@
+# 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 ...serialization.serializables import BoolField, FieldTypes, TupleField
+from ...utils import tokenize
+from .core import Entity, EntityData
+
+
+class ChunkData(EntityData):
+    __slots__ = ()
+
+    is_broadcaster = BoolField("is_broadcaster", default=False)
+    # If the operator is a shuffle mapper, this flag indicates whether the current chunk is mapper chunk when
+    # the operator produce multiple chunks such as TensorUnique.
+    is_mapper = BoolField("is_mapper", default=None)
+    # optional fields
+    _index = TupleField("index", FieldTypes.uint32)
+
+    def __repr__(self):
+        if self.op.stage is None:
+            return (
+                f"{type(self).__name__} <op={type(self.op).__name__}, "
+                f"key={self.key}>"
+            )
+        else:
+            return (
+                f"{type(self).__name__} <op={type(self.op).__name__}, "
+                f"stage={self.op.stage.name}, key={self.key}>"
+            )
+
+    @property
+    def index(self):
+        return getattr(self, "_index", None)
+
+    @property
+    def device(self):
+        return self.op.device
+
+    def _update_key(self):
+        object.__setattr__(
+            self,
+            "_key",
+            tokenize(
+                type(self).__name__,
+                *(getattr(self, k, None) for k in self._keys_ if k != "_index"),
+            ),
+        )
+
+
+class Chunk(Entity):
+    _allow_data_type_ = (ChunkData,)
+
+    def __repr__(self):
+        return f"{type(self).__name__}({self._data.__repr__()})"
+
+
+CHUNK_TYPE = (Chunk, ChunkData)

maxframe/core/entity/core.py

@@ -0,0 +1,152 @@
+# 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 ...serialization.serializables import (
+    DictField,
+    FieldTypes,
+    ReferenceField,
+    Serializable,
+)
+from ...utils import AttributeDict
+from ..base import Base
+
+
+class EntityData(Base):
+    __slots__ = ("_siblings",)
+    type_name = None
+
+    # required fields
+    _op = ReferenceField("op", "maxframe.core.operator.base.Operator")
+    # optional fields
+    _extra_params = DictField("extra_params", key_type=FieldTypes.string)
+
+    def __init__(self, *args, **kwargs):
+        extras = AttributeDict(
+            (k, kwargs.pop(k)) for k in set(kwargs) - set(self._FIELDS)
+        )
+        kwargs["_extra_params"] = kwargs.pop("_extra_params", extras)
+        super().__init__(*args, **kwargs)
+
+    @property
+    def op(self):
+        return self._op
+
+    @property
+    def inputs(self):
+        return self.op.inputs
+
+    @inputs.setter
+    def inputs(self, new_inputs):
+        self.op.inputs = new_inputs
+
+    def is_sparse(self):
+        return self.op.is_sparse()
+
+    issparse = is_sparse
+
+    @property
+    def extra_params(self):
+        return self._extra_params
+
+    def build_graph(self, **kw):
+        from ..graph.builder.utils import build_graph
+
+        return build_graph([self], **kw)
+
+    def visualize(self, graph_attrs=None, node_attrs=None, **kw):
+        from graphviz import Source
+
+        g = self.build_graph(**kw)
+        dot = g.to_dot(
+            graph_attrs=graph_attrs,
+            node_attrs=node_attrs,
+            result_chunk_keys={c.key for c in self.chunks},
+        )
+
+        return Source(dot)
+
+    def _need_execution(self):  # pylint: disable=no-self-use
+        # some tileable may generate unknown meta,
+        # they need to be executed first
+        return False
+
+
+class Entity(Serializable):
+    _allow_data_type_ = ()
+    type_name = None
+
+    _data = ReferenceField("data", EntityData)
+
+    def __init__(self, data=None, **kw):
+        super().__init__(_data=data, **kw)
+
+    def __dir__(self):
+        obj_dir = object.__dir__(self)
+        if self._data is not None:
+            obj_dir = sorted(set(dir(self._data) + obj_dir))
+        return obj_dir
+
+    def __str__(self):
+        return self._data.__str__()
+
+    def __repr__(self):
+        return self._data.__repr__()
+
+    def _check_data(self, data):
+        if data is not None and not isinstance(data, self._allow_data_type_):
+            raise TypeError(f"Expect {self._allow_data_type_}, got {type(data)}")
+
+    @property
+    def data(self):
+        return self._data
+
+    @data.setter
+    def data(self, new_data):
+        self._check_data(new_data)
+        self._data = new_data
+
+    def __copy__(self):
+        return self.copy()
+
+    def copy(self):
+        return self.copy_to(type(self)(None))
+
+    def copy_to(self, target):
+        target.data = self._data
+        return target
+
+    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)
+
+    def __setattr__(self, key, value):
+        try:
+            object.__setattr__(self, key, value)
+        except AttributeError:
+            return setattr(self._data, key, value)
+
+    def _need_execution(self):
+        return self._data._need_execution()
+
+
+ENTITY_TYPE = (Entity, EntityData)