mage-ai 0.9.70__py3-none-any.whl → 0.9.71__py3-none-any.whl

mage_ai/shared/models.py

@@ -3,10 +3,11 @@ import typing
 from dataclasses import dataclass, make_dataclass
 from enum import Enum
 from functools import reduce
-from typing import Dict, List
+from typing import Any, Dict, List, Optional, Union
 
 import inflection
 
+from mage_ai.shared.environments import is_debug
 from mage_ai.shared.hash import merge_dict
 from mage_ai.shared.parsers import encode_complex
 
@@ -34,33 +35,41 @@ class BaseClass:
         return annotations
 
     @classmethod
-    def load(self, **kwargs):
-        annotations = self.all_annotations()
+    def load(cls, *args, **kwargs):
+        annotations = cls.all_annotations()
 
-        props_init = self.load_to_dict(**kwargs)
+        if args is not None and len(args) >= 1:
+            kwargs = kwargs or {}
+            if isinstance(args[0], dict):
+                kwargs.update(args[0])
+            elif isinstance(args[0], cls):
+                kwargs.update(args[0].to_dict())
+
+        props_init = cls.load_to_dict(**kwargs)
 
         props = {}
         props_not_set = {}
         if props_init:
             for key, value in props_init.items():
-                if self.attribute_aliases and key in self.attribute_aliases:
-                    key = self.attribute_aliases[key]
+                if cls.attribute_aliases and key in cls.attribute_aliases:
+                    key = cls.attribute_aliases[key]
 
                 annotation = annotations.get(key)
                 if annotation:
-                    props[key] = self.convert_value(value, annotation)
+                    props[key] = cls.convert_value(value, annotation)
                 else:
                     props_not_set[key] = value
 
-        model = self(**props)
+        model = cls(**props)
 
         for key, value in props_not_set.items():
             try:
                 if not callable(getattr(model, key)):
                     model.set_value(key, value)
             except AttributeError as err:
-                print(f'[WARNING] {self.__name__}.load: {err}')
-                raise err
+                print(f'[WARNING] {cls.__name__}.load: {err}')
+                if is_debug():
+                    raise err
 
         return model
 
@@ -85,11 +94,14 @@ class BaseClass:
             if ignore_empty and len(value) == 0:
                 return None
 
-            return [self.convert_value(
-                v,
-                convert_enum=convert_enum,
-                ignore_empty=ignore_empty,
-            ) for v in value]
+            return [
+                self.convert_value(
+                    v,
+                    convert_enum=convert_enum,
+                    ignore_empty=ignore_empty,
+                )
+                for v in value
+            ]
 
         if not annotation:
             annotation = type(value)
@@ -133,9 +145,13 @@ class BaseClass:
                     ignore_empty=ignore_empty,
                 )
 
-        is_enum_class = value is not None and inspect.isclass(
-            annotation,
-        ) and issubclass(annotation, Enum)
+        is_enum_class = (
+            value is not None
+            and inspect.isclass(
+                annotation,
+            )
+            and issubclass(annotation, Enum)
+        )
         is_enum = isinstance(value, Enum)
         if is_enum_class and not is_enum:
             try:
@@ -183,7 +199,13 @@ class BaseClass:
         try:
             value = getattr(self, attribute_name)
             if value and isinstance(value, str):
-                setattr(self, attribute_name, enum_class(value))
+                setattr(
+                    self,
+                    attribute_name,
+                    enum_class.from_value(value)
+                    if issubclass(enum_class, BaseEnum)
+                    else enum_class(value),
+                )
         except AttributeError as err:
             print(f'[WARNING] {self.__class__.__name__}.serialize_attribute_enum: {err}')
 
@@ -194,7 +216,11 @@ class BaseClass:
                 arr = []
                 for value in values:
                     if isinstance(value, str):
-                        arr.append(enum_class(value))
+                        arr.append(
+                            enum_class.from_value(value)
+                            if issubclass(enum_class, BaseEnum)
+                            else enum_class(value)
+                        )
                     else:
                         arr.append(value)
                 setattr(self, attribute_name, arr)
@@ -230,10 +256,11 @@ class BaseClass:
                 if self.attribute_aliases and key in self.attribute_aliases:
                     key = self.attribute_aliases[key]
 
-                if ignore_empty and \
-                        (isinstance(value, list) or isinstance(value, dict)) and \
-                        len(value) == 0:
-
+                if (
+                    ignore_empty
+                    and (isinstance(value, list) or isinstance(value, dict))
+                    and len(value) == 0
+                ):
                     continue
 
                 data[key] = encode_complex(value)
@@ -264,3 +291,34 @@ class BaseDataClass(BaseClass):
         BaseClass.__init__(obj, *args, **kwargs)
 
         return obj
+
+
+class DelegatorTarget:
+    def __init__(self, target):
+        self._target = target
+
+    def __getattr__(self, item):
+        return getattr(self._target, item)
+
+
+class Delegator:
+    def __init__(self, target: Any):
+        self.target = target
+        self.delegate = DelegatorTarget(self.target)
+
+
+class BaseEnum(str, Enum):
+    @classmethod
+    def has_value(cls, value: Union[Any, str]) -> bool:
+        if isinstance(value, cls):
+            return True
+        return isinstance(value, str) and value.upper() in (name for name in cls.__members__)
+
+    @classmethod
+    def from_value(cls, value: Union[Any, str]) -> Optional[Any]:
+        if not cls.has_value(value):
+            return None
+        if isinstance(value, cls):
+            return value
+        if isinstance(value, str):
+            return cls[value.upper()]

mage_ai/shared/outputs.py

@@ -0,0 +1,87 @@
+import os
+from typing import Any, Dict, Optional, Tuple
+
+import joblib
+
+from mage_ai.ai.utils.xgboost import load_model as load_model_xgboost
+from mage_ai.ai.utils.xgboost import save_model as save_model_xgboost
+
+# MATRIX_NPZ_FILE, MATRIX_SAMPLE_NPZ_FILE
+from mage_ai.data_preparation.models.variables.constants import (
+    CONFIG_JSON_FILE,
+    JOBLIB_FILE,
+    JOBLIB_OBJECT_FILE,
+    MEDIA_IMAGE_VISUALIZATION_FILE,
+    UBJSON_MODEL_FILENAME,
+    VariableType,
+)
+from mage_ai.shared.array import is_iterable
+from mage_ai.shared.parsers import object_to_dict
+
+# from scipy.sparse import load_npz, save_npz
+
+
+def save_custom_object(
+    data: Any,
+    variable_path: str,
+    variable_type: Optional[VariableType] = None,
+) -> Tuple[Dict, Optional[str]]:
+    is_object = False
+    full_path = None
+
+    if VariableType.MODEL_SKLEARN == variable_type:
+        is_object = True
+        os.makedirs(variable_path, exist_ok=True)
+        full_path = os.path.join(variable_path, JOBLIB_FILE)
+        joblib.dump(data, full_path)
+    elif VariableType.MODEL_XGBOOST == variable_type:
+        is_object = True
+
+        save_model_xgboost(
+            data,
+            model_dir=variable_path,
+            model_filename=UBJSON_MODEL_FILENAME,
+            config_filename=CONFIG_JSON_FILE,
+            image_filename=MEDIA_IMAGE_VISUALIZATION_FILE,
+        )
+    # elif VariableType.MATRIX_SPARSE == variable_type:
+    #     save_npz(os.path.join(variable_path, MATRIX_NPZ_FILE), data)
+    elif VariableType.CUSTOM_OBJECT == variable_type:
+        is_object = True
+        os.makedirs(variable_path, exist_ok=True)
+
+        if not is_iterable(data):
+            full_path = os.path.join(variable_path, JOBLIB_OBJECT_FILE)
+            joblib.dump(data, full_path)
+    elif VariableType.DICTIONARY_COMPLEX == variable_type:
+        is_object = True
+        pass
+
+    if is_object:
+        data = object_to_dict(data, variable_type=variable_type)
+
+    return data, full_path
+
+
+def load_custom_object(
+    variable_path: str,
+    variable_type: Optional[VariableType] = None,
+) -> Optional[Any]:
+    try:
+        if VariableType.MODEL_SKLEARN == variable_type:
+            return joblib.load(os.path.join(variable_path, JOBLIB_FILE))
+        elif VariableType.MODEL_XGBOOST == variable_type:
+            return load_model_xgboost(
+                model_dir=variable_path,
+                model_filename=UBJSON_MODEL_FILENAME,
+                config_filename=CONFIG_JSON_FILE,
+                raise_exception=False,
+            )
+        # elif VariableType.MATRIX_SPARSE == variable_type:
+        #     return load_npz(os.path.join(variable_path, MATRIX_NPZ_FILE))
+        elif VariableType.CUSTOM_OBJECT == variable_type:
+            return joblib.load(os.path.join(variable_path, JOBLIB_OBJECT_FILE))
+        elif VariableType.DICTIONARY_COMPLEX == variable_type:
+            pass
+    except Exception as e:
+        print(f'Error loading custom object: {e}')

mage_ai/shared/parsers.py

@@ -1,11 +1,21 @@
+import base64
+import inspect
+import io
+import traceback
+from collections.abc import Generator
 from datetime import datetime
 from enum import Enum
 from json import JSONDecoder
+from typing import Any, Dict, List, Optional, Union
 
 import numpy as np
 import pandas as pd
+import polars as pl
+import scipy
 
+from mage_ai.data_preparation.models.variables.constants import VariableType
 from mage_ai.orchestration.db.models.base import BaseModel
+from mage_ai.shared.complex import is_model_sklearn, is_model_xgboost
 
 INTS = (
     np.int16,
@@ -14,6 +24,7 @@ INTS = (
     np.int8,
     np.int_,
     np.intc,
+    np.integer,
     np.intp,
     np.uint16,
     np.uint32,
@@ -24,6 +35,10 @@ INTS = (
 MAX_ITEMS_IN_SAMPLE_OUTPUT = 20
 
 
+def has_to_dict(obj) -> bool:
+    return hasattr(obj, 'to_dict')
+
+
 def encode_complex(obj):
     from mage_ai.shared.models import BaseDataClass
 
@@ -35,35 +50,47 @@ def encode_complex(obj):
         return obj.to_dict()
     elif isinstance(obj, Enum):
         return obj.value
-    elif hasattr(obj, 'isoformat') and 'method' in type(obj.isoformat).__name__:
-        return obj.isoformat()
-    elif isinstance(obj, np.integer):
-        return int(obj)
-    elif isinstance(obj, np.floating):
-        return float(obj)
-    elif isinstance(obj, np.ndarray):
-        return obj.tolist()
-    elif isinstance(obj, datetime):
+    elif (
+        isinstance(obj, datetime)
+        or hasattr(obj, 'isoformat')
+        and 'method' in type(obj.isoformat).__name__
+    ):
         return obj.isoformat()
     elif isinstance(obj, INTS):
         return int(obj)
-    elif isinstance(obj, (np.float_, np.float16, np.float32, np.float64)):
+    elif isinstance(obj, (np.float_, np.float16, np.float32, np.float64, np.floating)):
         return float(obj)
     elif isinstance(obj, (np.complex_, np.complex64, np.complex128)):
         return {'real': obj.real, 'imag': obj.imag}
     elif isinstance(obj, (np.ndarray,)):
+        # np.array is a function
         return obj.tolist()
     elif isinstance(obj, (np.bool_)):
         return bool(obj)
-    elif isinstance(obj, (np.void)):
+    elif isinstance(obj, (np.void, pd._libs.missing.NAType)):
+        # Convert pandas._libs.missing.NAType to None
         return None
     elif isinstance(obj, pd.DataFrame):
         return obj.to_dict(orient='records')
-    elif isinstance(obj, pd.Series):
+    elif isinstance(obj, pl.DataFrame):
+        return obj.to_dicts()
+    elif isinstance(obj, (pd.Index, pd.Series, pl.Series)):
         return obj.to_list()
-    # Convert pandas._libs.missing.NAType to None
-    elif isinstance(obj, pd._libs.missing.NAType):
-        return None
+    elif isinstance(obj, scipy.sparse.csr_matrix):
+        return serialize_matrix(obj)
+    elif is_model_sklearn(obj) or is_model_xgboost(obj) or inspect.isclass(obj):
+        return object_to_uuid(obj)
+    elif has_to_dict(obj):
+        return obj.to_dict()
+    elif isinstance(obj, Generator):
+        return object_to_uuid(obj)
+    elif isinstance(obj, Exception):
+        # Serialize the exception
+        return {
+            'type': type(obj).__name__,
+            'message': str(obj),
+            'traceback': traceback.format_tb(obj.__traceback__),
+        }
 
     return obj
 
@@ -110,3 +137,102 @@ def sample_output(obj):
             output[k] = v
         return output, sampled
     return obj, False
+
+
+def serialize_matrix(csr_matrix: scipy.sparse._csr.csr_matrix) -> Dict:
+    with io.BytesIO() as buffer:
+        scipy.sparse.save_npz(buffer, csr_matrix)
+        buffer.seek(0)
+        data = base64.b64encode(buffer.read()).decode('ascii')
+
+    return {'__type__': 'scipy.sparse.csr_matrix', '__data__': data}
+
+
+def deserialize_matrix(json_dict: Dict) -> scipy.sparse._csr.csr_matrix:
+    data = json_dict['__data__']
+    data = base64.b64decode(data.encode('ascii'))
+
+    with io.BytesIO(data) as buffer:
+        buffer.seek(0)
+        csr_matrix = scipy.sparse.load_npz(buffer)
+
+    return csr_matrix
+
+
+def convert_matrix_to_dataframe(csr_matrix: scipy.sparse.csr_matrix) -> pd.DataFrame:
+    if isinstance(csr_matrix, scipy.sparse.csr_matrix):
+        n_columns = csr_matrix.shape[1]
+        return pd.DataFrame(csr_matrix.toarray(), columns=[str(i) for i in range(n_columns)])
+    return csr_matrix
+
+
+def polars_to_dict_split(df: pl.DataFrame) -> Dict:
+    # Extract column names
+    columns = df.columns
+
+    # Extract rows as list of lists
+    data = df.to_numpy().tolist()
+
+    # Construct and return the dictionary
+    return dict(columns=columns, data=data)
+
+
+def object_to_hash(obj: Any) -> str:
+    return hex(id(obj))
+
+
+def object_uuid_parts(obj: Any) -> Dict[str, Union[bool, str, List[str]]]:
+    if inspect.isclass(obj):
+        obj_cls = obj
+    else:
+        obj_cls = obj.__class__
+
+    return dict(
+        module=obj_cls.__module__,
+        name=obj_cls.__name__,
+    )
+
+
+def object_to_uuid(obj: Any, include_hash: bool = False) -> str:
+    hash_uuid = None
+    if include_hash:
+        hash_uuid = object_to_hash(obj)
+
+    parts = object_uuid_parts(obj)
+    uuid = '.'.join([
+        t
+        for t in [
+            str(parts['module']),
+            str(parts['name']),
+        ]
+        if t
+    ])
+
+    if hash_uuid:
+        return f'{uuid} {hash_uuid}'
+
+    return uuid
+
+
+def object_to_dict(
+    obj: Any,
+    include_hash: bool = True,
+    include_uuid: bool = True,
+    variable_type: Optional[VariableType] = None,
+) -> Dict[str, Union[bool, str, List[str]]]:
+    is_class = inspect.isclass(obj)
+
+    data_dict = object_uuid_parts(obj)
+    data_dict['type'] = 'class' if is_class else 'instance'
+    data_dict['uuid'] = object_to_uuid(obj, include_hash=False)
+
+    if include_uuid:
+        data_dict['uuid'] = object_to_uuid(obj, include_hash=False)
+
+    if include_hash and not is_class:
+        data_dict['hash'] = object_to_hash(obj)
+
+    if variable_type:
+        data_dict['variable_type'] = variable_type.value
+
+    return data_dict

mage_ai/shared/path_fixer.py

@@ -33,7 +33,9 @@ def add_absolute_path(file_path: str, add_base_repo_path: bool = True) -> str:
                     # DBT v1 paths
                     full_path = os.path.join(repo_path, 'dbt', file_path)
                 else:
-                    full_path = find_directory(repo_path, lambda x: x.endswith(file_path))
+                    full_path = find_directory(
+                        repo_path, lambda x: x.endswith(file_path)
+                    )
 
     if full_path:
         # Remove the repo path
@@ -93,7 +95,7 @@ def remove_base_repo_name(file_path: str) -> str:
 
 def remove_base_repo_path_or_name(file_path: str) -> str:
     if not file_path:
-        return
+        return file_path
 
     if os.path.isabs(file_path):
         value = str(remove_base_repo_path(file_path))
@@ -121,11 +123,13 @@ def remove_base_repo_directory_name(file_path: str) -> str:
 
 def remove_repo_names(file_path: str) -> str:
     try:
-        path = Path(file_path).relative_to(get_repo_path(
-            absolute_path=False,
-            file_path=file_path,
-            root_project=False,
-        ))
+        path = Path(file_path).relative_to(
+            get_repo_path(
+                absolute_path=False,
+                file_path=file_path,
+                root_project=False,
+            )
+        )
         return str(path)
     except ValueError:
         return file_path

mage_ai/shared/singletons/base.py

@@ -0,0 +1,47 @@
+import asyncio
+import threading
+from functools import wraps
+
+
+class SingletonBase:
+    def __init__(self):
+        self.value = None
+        self._lock = asyncio.Lock()
+
+
+class SingletonMeta(type):
+    _instances = {}
+    _lock: threading.Lock = threading.Lock()
+
+    def __call__(cls, *args, **kwargs):
+        with cls._lock:
+            if cls not in cls._instances:
+                instance = super().__call__(*args, **kwargs)
+                cls._instances[cls] = instance
+        return cls._instances[cls]
+
+
+def singleton(cls):
+    class Wrapper(cls, metaclass=SingletonMeta):
+        __doc__ = cls.__doc__
+
+    Wrapper.__name__ = cls.__name__
+    return Wrapper
+
+
+def synchronized_method(method):
+    @wraps(method)
+    def wrapper(self, *args, **kwargs):
+        with self._lock:
+            return method(self, *args, **kwargs)
+
+    return wrapper
+
+
+def asynchronized_method(method):
+    @wraps(method)
+    async def wrapper(self, *args, **kwargs):
+        async with self._lock:
+            return method(self, *args, **kwargs)
+
+    return wrapper

mage_ai/shared/singletons/memory.py

@@ -0,0 +1,38 @@
+from typing import Any, Optional
+
+from mage_ai.shared.singletons.base import SingletonBase, singleton
+
+
+@singleton
+class MemoryManagerController(SingletonBase):
+    def __init__(self):
+        super().__init__()  # Ensures self._lock is initialized
+        self.events = {}
+
+    def add_event_monitor(self, key: str, stop_event, monitor_thread):
+        self.stop_event(key)
+        self.events[key] = (stop_event, monitor_thread)
+
+    def stop_event(
+        self, key, stop_event: Optional[Any] = None, monitor_thread: Optional[Any] = None
+    ):
+        if self.events.get(key):
+            if not stop_event and not monitor_thread:
+                stop_event, monitor_thread = self.events[key]
+
+            if stop_event:
+                stop_event.set()
+                stop_event = None
+            if monitor_thread:
+                monitor_thread.join()
+                monitor_thread = None
+            del self.events[key]
+
+    def stop_all_events(self):
+        for key, pair in self.events.items():
+            self.stop_event(key, *pair)
+            print(f'Event stopped: {key}')
+
+
+def get_memory_manager_controller():
+    return MemoryManagerController()

mage_ai/shared/strings.py

@@ -1,9 +1,20 @@
+import json
 import re
 from typing import List
 
 import inflection
 
 
+def is_json(myjson):
+    try:
+        json_object = json.loads(str(myjson))
+        if isinstance(json_object, dict):
+            return json_object
+        return False
+    except Exception:
+        return False
+
+
 def camel_to_snake_case(name):
     name = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', name)
     name = re.sub('__([A-Z])', r'_\1', name)
@@ -30,6 +41,28 @@ def is_number(s) -> bool:
         return False
 
 
+def string_to_number(value: str):
+    """
+    Attempts to convert a string to an integer or float.
+
+    Parameters:
+    - value: The string to convert.
+
+    Returns:
+    - int, float, or original str depending on its content.
+    """
+    try:
+        # First, try to convert it to an integer.
+        return int(value)
+    except ValueError:
+        # If it fails, try to convert it to a float.
+        try:
+            return float(value)
+        except ValueError:
+            # If it also fails, return the original string.
+            return value
+
+
 def replacer(s, newstring, index, nofail=False):
     # raise an error if index is outside of the string
     if not nofail and index not in range(len(s)):
@@ -42,7 +75,7 @@ def replacer(s, newstring, index, nofail=False):
         return s + newstring
 
     # insert the new string between 'slices' of the original
-    return s[:index] + newstring + s[index + 1:]
+    return s[:index] + newstring + s[index + 1 :]
 
 
 def remove_extension_from_filename(filename: str) -> str:

mage_ai/shared/yaml.py

@@ -0,0 +1,24 @@
+from typing import Dict, List, Union
+
+import yaml
+
+try:
+    # Speed up yaml load performance with CLoader
+    from yaml import CLoader as Loader
+except Exception:
+    from yaml import FullLoader as Loader
+
+
+def trim_strings(data: Union[Dict, List, str]) -> Union[Dict, List, str]:
+    if isinstance(data, dict):
+        return {k: trim_strings(v) for k, v in data.items()}
+    elif isinstance(data, list):
+        return [trim_strings(item) for item in data]
+    elif isinstance(data, str):
+        return data.rstrip()  # Remove trailing whitespace and newlines
+    else:
+        return data
+
+
+def load_yaml(data):
+    return yaml.load(data, Loader=Loader)