maxframe 1.3.0__cp38-cp38-win_amd64.whl → 2.0.0__cp38-cp38-win_amd64.whl

maxframe/tests/utils.py

@@ -13,8 +13,10 @@
 # limitations under the License.
 
 import asyncio
+import contextlib
 import functools
 import hashlib
+import logging
 import os
 import queue
 import socket
@@ -158,6 +160,13 @@ def require_cudf(func):
     return func
 
 
+def require_ray(func):
+    if pytest:
+        func = pytest.mark.ray(func)
+    func = pytest.mark.skipif(ray is None, reason="ray not installed")(func)
+    return func
+
+
 def require_hadoop(func):
     if pytest:
         func = pytest.mark.hadoop(func)
@@ -168,7 +177,7 @@ def require_hadoop(func):
 
 
 def get_test_unique_name(size=None):
-    test_name = os.getenv("PYTEST_CURRENT_TEST", "pyodps_test")
+    test_name = os.getenv("PYTEST_CURRENT_TEST", "maxframe_test")
     digest = hashlib.md5(to_binary(test_name)).hexdigest()
     if size:
         digest = digest[:size]
@@ -181,6 +190,55 @@ def assert_mf_index_dtype(idx_obj, dtype):
     assert isinstance(idx_obj, IndexValue.IndexBase) and idx_obj.dtype == dtype
 
 
+@contextlib.contextmanager
+def create_test_volume(vol_name, oss_config):
+    odps_entry = ODPS.from_environments()
+
+    oss_test_dir_name = "test_dir_" + vol_name
+    if oss_config is None:
+        pytest.skip("Need oss and its config to run this test")
+    (
+        oss_access_id,
+        oss_secret_access_key,
+        oss_bucket_name,
+        oss_endpoint,
+    ) = oss_config.oss_config
+
+    if "test" in oss_endpoint:
+        # offline config
+        test_location = "oss://%s:%s@%s/%s/%s" % (
+            oss_access_id,
+            oss_secret_access_key,
+            oss_endpoint,
+            oss_bucket_name,
+            oss_test_dir_name,
+        )
+        rolearn = None
+    else:
+        # online config
+        endpoint_parts = oss_endpoint.split(".", 1)
+        if "-internal" not in endpoint_parts[0]:
+            endpoint_parts[0] += "-internal"
+        test_location = "oss://%s/%s/%s" % (
+            ".".join(endpoint_parts),
+            oss_bucket_name,
+            oss_test_dir_name,
+        )
+        rolearn = oss_config.oss_rolearn
+
+    oss_config.oss_bucket.put_object(oss_test_dir_name + "/", b"")
+    odps_entry.create_external_volume(vol_name, location=test_location, rolearn=rolearn)
+
+    try:
+        yield vol_name
+    finally:
+        try:
+            logging.warning("Deleting test volume %s", vol_name)
+            odps_entry.delete_volume(vol_name, auto_remove_dir=True, recursive=True)
+        except:
+            pass
+
+
 def ensure_table_deleted(odps_entry: ODPS, table_name: str) -> None:
     retry_times = 20
     while odps_entry.exist_table(table_name) and retry_times > 0:

maxframe/udf.py

@@ -17,6 +17,7 @@ from typing import Callable, List, Optional, Union
 
 from odps.models import Resource
 
+from .serialization import load_member
 from .serialization.serializables import (
     BoolField,
     DictField,
@@ -26,7 +27,7 @@ from .serialization.serializables import (
     Serializable,
     StringField,
 )
-from .utils import tokenize
+from .utils import extract_class_name, tokenize
 
 
 class PythonPackOptions(Serializable):
@@ -51,6 +52,52 @@ class PythonPackOptions(Serializable):
         return f"<PythonPackOptions {self.requirements} {args_str}>"
 
 
+class BuiltinFunction(Serializable):
+    """
+    Record reference for builtin functions. The function body
+    will NOT be serialized when submitting jobs.
+    """
+
+    __slots__ = ("_func",)
+
+    func_name = StringField("func_name")
+
+    def __init__(self, func: Optional[Callable] = None, **kw):
+        self._func = func
+        if func is not None:
+            func_name = extract_class_name(func)
+            if "<" in func_name:
+                raise ValueError("Cannot be a local or lambda function")
+            kw["func_name"] = kw.get("func_name") or func_name
+        super().__init__(**kw)
+
+    @property
+    def func(self):
+        if getattr(self, "_func", None) is None:
+            assert isinstance(self.func_name, str)
+            self._func = load_member(self.func_name, type(self)).func
+        return self._func
+
+    @property
+    def module(self):
+        # use func_name instead of func itself to avoid it
+        # been imported (might cause infinite recursion!)
+        assert isinstance(self.func_name, str)
+        return self.func_name.split("#")[0]
+
+    def __getattr__(self, item):
+        if item.startswith("_") and not item.endswith("_"):
+            return super().__getattribute__(item)
+        return getattr(self.func, item)
+
+    def __call__(self, *args, **kw):
+        return self.func(*args, **kw)
+
+
+def builtin_function(func: Callable) -> BuiltinFunction:
+    return BuiltinFunction(func=func)
+
+
 class MarkedFunction(Serializable):
     func = FunctionField("func")
     resources = ListField("resources", FieldTypes.string, default_factory=list)
@@ -129,16 +176,39 @@ def with_running_options(
     *,
     engine: Optional[str] = None,
     cpu: Optional[int] = None,
-    memory: Optional[int] = None,
+    memory: Optional[Union[str, int]] = None,
     **kwargs,
 ):
+    """
+    Set running options for the UDF.
+
+    Parameters
+    ----------
+    engine: Optional[str]
+        The engine to run the UDF.
+    cpu: Optional[int]
+        The CPU to run the UDF.
+    memory: Optional[Union[str, int]]
+        The memory to run the UDF. If it is an int, it is in GB.
+        If it is a str, it is in the format of "10GiB", "30MiB", etc.
+    kwargs
+        Other running options.
+    """
     engine = engine.upper() if engine else None
-    resources = {"cpu": cpu, "memory": memory, **kwargs}
+    resources = kwargs.copy()
 
-    if cpu is not None and cpu <= 0:
+    if cpu is not None and isinstance(cpu, int) and cpu <= 0:
         raise ValueError("cpu must be greater than 0")
-    if memory is not None and memory <= 0:
-        raise ValueError("memory must be greater than 0")
+    if memory is not None:
+        if not isinstance(memory, (int, str)):
+            raise TypeError("memory must be an int or str")
+        if isinstance(memory, int) and memory <= 0:
+            raise ValueError("memory must be greater than 0")
+
+    if cpu:
+        resources["cpu"] = cpu
+    if memory:
+        resources["memory"] = memory
 
     def func_wrapper(func):
         if all(v is None for v in (engine, cpu, memory)):