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

mage_ai/ai/utils/xgboost.py

@@ -0,0 +1,222 @@
+import base64
+import io
+import json
+import os
+from typing import Any, Optional, Tuple
+
+from mage_ai.data_preparation.models.variables.constants import (
+    CONFIG_JSON_FILE,
+    MEDIA_IMAGE_VISUALIZATION_FILE,
+    UBJSON_MODEL_FILENAME,
+)
+from mage_ai.settings.server import MAX_OUTPUT_IMAGE_PREVIEW_SIZE
+from mage_ai.shared.environments import is_debug
+
+
+# Function to check if the booster is trained
+def is_booster_trained(booster: Any, raise_exception: bool = True) -> bool:
+    # Check if the booster has been trained by checking for the existence of an attribute
+    # like 'feature_names'; this attribute is set upon training.
+    # If the model is not trained, accessing this attribute should raise an AttributeError.
+    try:
+        import xgboost as xgb
+
+        # If feature_names exists, it means the model has been trained.
+        return booster.num_boosted_rounds() >= 1
+    except xgb.core.XGBoostError as err:
+        message = f'XGBoost model is not trained. {err}'
+
+        if raise_exception:
+            raise Exception(message)
+
+        print(message)
+
+        # If the model is not trained, an AttributeError is raised.
+        return False
+
+
+def load_model(
+    model_dir: str,
+    model_filename: str = UBJSON_MODEL_FILENAME,
+    config_filename: str = CONFIG_JSON_FILE,
+    raise_exception: bool = True,
+) -> Optional[Any]:
+    try:
+        import xgboost as xgb
+
+        model_path = os.path.join(model_dir, model_filename)
+        model = xgb.Booster()
+        model.load_model(model_path)
+
+        config_path = os.path.join(model_dir, config_filename)
+        with open(config_path, 'r') as file:
+            model_config = json.load(file)
+
+        model_config_str = json.dumps(model_config)
+        # Apply the saved configuration to the model
+        model.load_config(model_config_str)
+
+        return model
+    except Exception as err:
+        if raise_exception or is_debug():
+            raise err
+        print(f'[ERROR] XGBoost.load_model: {err}')
+
+    return None
+
+
+def save_model(
+    booster: Any,
+    model_dir: str,
+    model_filename: str = UBJSON_MODEL_FILENAME,
+    config_filename: str = CONFIG_JSON_FILE,
+    image_filename: str = MEDIA_IMAGE_VISUALIZATION_FILE,
+    raise_exception: bool = True,
+) -> bool:
+    if not is_booster_trained(booster, raise_exception=raise_exception):
+        return False
+
+    os.makedirs(model_dir, exist_ok=True)
+
+    # Save detailed configuration of the model that includes all the hyperparameters
+    # and settings
+    model_path = os.path.join(model_dir, model_filename)
+    booster.save_model(model_path)
+
+    # Save the structure of the trees (for tree-based models like gradient boosting)
+    # along with some basic configurations necessary to understand the model structure itself
+    config_path = os.path.join(model_dir, config_filename)
+    with open(config_path, 'w') as f:
+        f.write(booster.save_config())
+
+    if image_filename:
+        image_path = os.path.join(model_dir, image_filename)
+        try:
+            create_tree_visualization(booster, image_path=image_path)
+        except Exception as err:
+            print(f'[ERROR] XGBoost.load_model: {err}')
+
+    return True
+
+
+def create_tree_visualization(
+    model: Any,
+    image_path: Optional[str] = None,
+    max_render_size: int = MAX_OUTPUT_IMAGE_PREVIEW_SIZE,
+    max_trees: int = 12,
+    num_trees: int = 0,
+) -> Tuple[Optional[str], bool]:
+    try:
+        import xgboost as xgb
+        from PIL import Image
+
+        # Increase the maximum allowed image size to, say, 500 million pixels
+        Image.MAX_IMAGE_PIXELS = 1024 * 1024 * 500
+
+        trees_to_render = 0
+        n_trees = model.num_boosted_rounds()
+        if n_trees > max_trees or n_trees == 0:
+            trees_to_render = max_trees
+        elif n_trees >= 1 and num_trees < n_trees:
+            trees_to_render = num_trees
+        elif n_trees == max_trees:
+            trees_to_render = 0
+
+        graph = xgb.to_graphviz(
+            model, num_trees=trees_to_render, rankdir='TB', format='png'
+        )
+
+        if image_path:
+            # Remove the '.png' extension when specifying the filename
+            base_image_path = (
+                image_path.rsplit('.', 1)[0] if '.' in image_path else image_path
+            )
+
+            # Pass the adjusted filename without the extension
+            graph.render(filename=base_image_path, cleanup=True, format='png')
+            # Since the 'format' is 'png', Graphviz will output 'visualization.png'
+            return image_path, True
+
+        # Save the graph to a temporary PNG file (or use BytesIO directly with some adjustments)
+        png_bytes = graph.pipe(format='png')
+
+        # Convert PNG bytes to an PIL Image object
+        image = Image.open(io.BytesIO(png_bytes))
+
+        # Save or Display your image directly if needed
+        # image.show() # Uncomment to display the image directly
+
+        # Convert the PIL Image to a base64 string
+        buffered = io.BytesIO()
+        image.save(buffered, format='PNG')
+        img_str = base64.b64encode(buffered.getvalue()).decode()
+
+        buffered.seek(0)  # Reset pointer to the beginning of the buffer
+        file_size_bytes = buffered.getbuffer().nbytes
+        if file_size_bytes > max_render_size:
+            message = (
+                'XGBoost tree visualization created an image that exceeds '
+                f'{max_render_size} bytes (actual size is {file_size_bytes} bytes). '
+                'No preview will be shown in the browser. '
+                'To increase the preview limit, set the environment variable '
+                'MAX_OUTPUT_IMAGE_PREVIEW_SIZE to a larger byte size value.'
+            )
+            return message, False
+
+        return img_str, True
+    except Exception as err:
+        print(f'[ERROR] XGBoost.create_tree_visualization: {err}')
+        return str(err), False
+
+
+def render_tree_visualization(
+    image_dir: str,
+    image_filename: Optional[str] = None,
+    max_render_size: Optional[int] = None,
+) -> Tuple[Optional[str], bool]:
+    # Load the model’s tree from a PNG file into base64 format
+    try:
+        image_filename = image_filename or MEDIA_IMAGE_VISUALIZATION_FILE
+        max_render_size = max_render_size or MAX_OUTPUT_IMAGE_PREVIEW_SIZE
+        image_path = os.path.join(image_dir, image_filename)
+
+        # Check for file size before opening
+        file_size_bytes = os.path.getsize(image_path)
+
+        if file_size_bytes > max_render_size:
+            message = (
+                'XGBoost tree visualization created an image that exceeds '
+                f'{max_render_size} bytes (actual size is {file_size_bytes} bytes). '
+                'No preview will be shown in the browser. '
+                'To increase the preview limit, set the environment variable '
+                'MAX_OUTPUT_IMAGE_PREVIEW_SIZE to a larger byte size value.'
+            )
+            return message, False
+
+        with open(image_path, 'rb') as image_file:
+            encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
+        return encoded_string, True
+    except Exception as err:
+        print(f'[ERROR] XGBoost.render_tree_visualization: {err}')
+        return str(err), False
+
+
+def create_tree_plot(model: Any, image_path: str, num_trees: int = 0) -> str:
+    try:
+        import matplotlib.pyplot as plt
+        import xgboost as xgb
+
+        plt.close('all')
+
+        plt.figure(dpi=300)
+        xgb.plot_tree(model, num_trees=5)
+        plt.tight_layout()
+
+        plt.savefig(image_path, dpi=300, bbox_inches='tight')
+
+        plt.close()
+
+        return image_path
+    except Exception as err:
+        print(f'[ERROR] XGBoost.create_tree_plot: {err}')
+        return str(err)

mage_ai/api/errors.py

@@ -3,46 +3,58 @@ from typing import Dict
 
 class ApiError(Exception):
     UNAUTHORIZED_ACCESS = {
-        'code': 403,
-        'message': 'Unauthorized access.',
-        'type': 'unauthorized_access',
+        "code": 403,
+        "message": "Unauthorized access.",
+        "type": "unauthorized_access",
     }
 
     RESOURCE_ERROR = {
-        'code': 500,
-        'message': 'API resource error.',
-        'type': 'api_resource_error',
+        "code": 500,
+        "message": "API resource error.",
+        "type": "api_resource_error",
     }
     RESOURCE_INVALID = {
-        'code': 402,
-        'message': 'Record is invalid.',
-        'type': 'record_invalid',
+        "code": 402,
+        "message": "Record is invalid.",
+        "type": "record_invalid",
     }
     RESOURCE_NOT_FOUND = {
-        'code': 404,
-        'message': 'Record not found.',
-        'type': 'record_not_found',
+        "code": 404,
+        "message": "Record not found.",
+        "type": "record_not_found",
     }
 
     INVALID_API_KEY = {
-        'code': 403,
-        'message': 'Invalid API key.',
-        'type': 'invalid_api_key',
+        "code": 403,
+        "message": "Invalid API key.",
+        "type": "invalid_api_key",
     }
     EXPIRED_OAUTH_TOKEN = {
-        'code': 401,
-        'message': 'Expired OAuth token.',
-        'type': 'expired_oauth_token',
+        "code": 401,
+        "message": "Expired OAuth token.",
+        "type": "expired_oauth_token",
     }
     INVALID_OAUTH_TOKEN = {
-        'code': 401,
-        'message': 'Invalid OAuth token.',
-        'type': 'invalid_oauth_token',
+        "code": 401,
+        "message": "Invalid OAuth token.",
+        "type": "invalid_oauth_token",
     }
 
     def __init__(self, opts: Dict = None):
         if opts:
-            self.code = opts.get('code')
-            self.errors = opts.get('errors')
-            self.message = opts.get('message')
-            self.type = opts.get('type')
+            self.code = opts.get("code")
+            self.errors = opts.get("errors")
+            self.message = opts.get("message")
+            self.type = opts.get("type")
+
+    def to_dict(self, **kwargs) -> Dict:
+        errors = self.errors
+        if not isinstance(errors, list):
+            errors = str(errors).split("\n")
+
+        return dict(
+            code=self.code,
+            errors=errors,
+            message=self.message,
+            type=self.type,
+        )

mage_ai/api/operations/base.py

@@ -72,6 +72,7 @@ class BaseOperation():
         self.__combined_options_attr = None
         self.__presentation_format_attr = None
         self.__updated_options_attr = None
+        self.__flags = dict()
 
     async def execute(self):
         db_connection.start_cache()
@@ -295,7 +296,15 @@ class BaseOperation():
         resource_parent: Any = None,
         resources: List[Dict] = None,
     ) -> List[Hook]:
-        if not Project.is_feature_enabled_in_root_or_active_project(FeatureUUID.GLOBAL_HOOKS):
+        if FeatureUUID.GLOBAL_HOOKS not in self.__flags:
+            self.__flags[FeatureUUID.GLOBAL_HOOKS] = \
+                Project.is_feature_enabled_in_root_or_active_project(
+                    FeatureUUID.GLOBAL_HOOKS,
+                    context_data=self.context.data,
+                    user=self.user,
+            )
+
+        if not self.__flags[FeatureUUID.GLOBAL_HOOKS]:
             return None
 
         operation_types = [operation_type]
@@ -744,9 +753,11 @@ class BaseOperation():
                 try:
                     model = parent_resource_class.get_model(
                         self.resource_parent_id,
+                        context_data=self.context.data,
                         query=self.query,
                         resource_class=self.__resource_class(),
                         resource_id=self.pk,
+                        user=self.user,
                     )
                     if inspect.isawaitable(model):
                         model = await model
@@ -758,6 +769,7 @@ class BaseOperation():
         if not self.__combined_options_attr:
             self.__combined_options_attr = {
                 'context': self.context,
+                'context_data': self.context.data,
                 'meta': self.meta,
                 'options': self.options,
                 'payload': self.payload,

mage_ai/api/parsers/PipelineScheduleParser.py

@@ -40,7 +40,7 @@ async def parse_write(
     if not pipeline:
         pipeline = parser.policy.resource.pipeline
 
-    if Project(pipeline.repo_config if pipeline else None).is_feature_enabled(
+    if Project(repo_config=pipeline.repo_config if pipeline else None).is_feature_enabled(
         FeatureUUID.INTERACTIONS,
     ):
         pipeline_interaction = (parser.policy.result_set().context.data or {}).get(

mage_ai/api/policies/BackfillPolicy.py

@@ -32,6 +32,7 @@ BackfillPolicy.allow_actions([
 
 BackfillPolicy.allow_read([
     'pipeline_run_dates',
+    'run_status_counts',
     'total_run_count',
 ] + BackfillPresenter.default_attributes, scopes=[
     OauthScope.CLIENT_PRIVATE,

mage_ai/api/policies/BlockOutputPolicy.py

@@ -8,22 +8,45 @@ class BlockOutputPolicy(BasePolicy):
     pass
 
 
-BlockOutputPolicy.allow_actions([
-    constants.DETAIL,
-], scopes=[
-    OauthScope.CLIENT_PRIVATE,
-], condition=lambda policy: policy.has_at_least_viewer_role())
+BlockOutputPolicy.allow_actions(
+    [
+        constants.DETAIL,
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    condition=lambda policy: policy.has_at_least_viewer_role(),
+)
 
-BlockOutputPolicy.allow_read(BlockOutputPresenter.default_attributes, scopes=[
-    OauthScope.CLIENT_PRIVATE,
-], on_action=[
-    constants.DETAIL,
-], condition=lambda policy: policy.has_at_least_viewer_role())
+BlockOutputPolicy.allow_read(
+    BlockOutputPresenter.default_attributes,
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        constants.DETAIL,
+    ],
+    condition=lambda policy: policy.has_at_least_viewer_role(),
+)
 
-BlockOutputPolicy.allow_query([
-    'pipeline_uuid',
-], scopes=[
-    OauthScope.CLIENT_PRIVATE,
-], on_action=[
-    constants.DETAIL,
-], condition=lambda policy: policy.has_at_least_viewer_role())
+BlockOutputPolicy.allow_query(
+    [
+        'block_uuid',
+        'csv_lines_only',
+        'dynamic_block_index',
+        'exclude_blank_variable_uuids',
+        'execution_partition',
+        'include_print_outputs',
+        'pipeline_uuid',
+        'sample',
+        'sample_count',
+        'variable_type',
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        constants.DETAIL,
+    ],
+    condition=lambda policy: policy.has_at_least_viewer_role(),
+)

mage_ai/api/policies/GlobalDataProductPolicy.py

@@ -8,44 +8,94 @@ class GlobalDataProductPolicy(BasePolicy):
     pass
 
 
-GlobalDataProductPolicy.allow_actions([
-    constants.DETAIL,
-    constants.LIST,
-], scopes=[
-    OauthScope.CLIENT_PRIVATE,
-], condition=lambda policy: policy.has_at_least_viewer_role())
-
-
-GlobalDataProductPolicy.allow_actions([
-    constants.CREATE,
-    constants.DELETE,
-    constants.UPDATE,
-], scopes=[
-    OauthScope.CLIENT_PRIVATE,
-], condition=lambda policy: policy.has_at_least_editor_role())
-
-
-GlobalDataProductPolicy.allow_read(GlobalDataProductPresenter.default_attributes + [], scopes=[
-    OauthScope.CLIENT_PRIVATE,
-], on_action=[
-    constants.DETAIL,
-    constants.LIST,
-], condition=lambda policy: policy.has_at_least_viewer_role())
-
-
-GlobalDataProductPolicy.allow_read(GlobalDataProductPresenter.default_attributes + [], scopes=[
-    OauthScope.CLIENT_PRIVATE,
-], on_action=[
-    constants.CREATE,
-    constants.DELETE,
-    constants.UPDATE,
-], condition=lambda policy: policy.has_at_least_editor_role())
-
-
-GlobalDataProductPolicy.allow_write(GlobalDataProductPresenter.default_attributes + [], scopes=[
-    OauthScope.CLIENT_PRIVATE,
-], on_action=[
-    constants.CREATE,
-    constants.DELETE,
-    constants.UPDATE,
-], condition=lambda policy: policy.has_at_least_editor_role())
+GlobalDataProductPolicy.allow_actions(
+    [
+        constants.DETAIL,
+        constants.LIST,
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    condition=lambda policy: policy.has_at_least_viewer_role(),
+)
+
+
+GlobalDataProductPolicy.allow_actions(
+    [
+        constants.CREATE,
+        constants.DELETE,
+        constants.UPDATE,
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    condition=lambda policy: policy.has_at_least_editor_role(),
+)
+
+
+GlobalDataProductPolicy.allow_read(
+    GlobalDataProductPresenter.default_attributes + [],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        constants.DETAIL,
+        constants.LIST,
+    ],
+    condition=lambda policy: policy.has_at_least_viewer_role(),
+)
+
+
+GlobalDataProductPolicy.allow_read(
+    GlobalDataProductPresenter.default_attributes + [],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        constants.CREATE,
+        constants.DELETE,
+        constants.UPDATE,
+    ],
+    condition=lambda policy: policy.has_at_least_editor_role(),
+)
+
+
+GlobalDataProductPolicy.allow_write(
+    GlobalDataProductPresenter.default_attributes + [],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        constants.CREATE,
+        constants.DELETE,
+        constants.UPDATE,
+    ],
+    condition=lambda policy: policy.has_at_least_editor_role(),
+)
+
+GlobalDataProductPolicy.allow_query(
+    [
+        'current_project',
+        'repo_path',
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        constants.LIST,
+    ],
+    condition=lambda policy: policy.has_at_least_viewer_role(),
+)
+
+GlobalDataProductPolicy.allow_query(
+    [
+        'project',
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        constants.DETAIL,
+    ],
+    condition=lambda policy: policy.has_at_least_viewer_role(),
+)

mage_ai/api/policies/KernelPolicy.py

@@ -11,35 +11,58 @@ class KernelPolicy(BasePolicy):
         return Entity.ANY, None
 
 
-KernelPolicy.allow_actions([
-    constants.LIST,
-], scopes=[
-    OauthScope.CLIENT_PRIVATE,
-], condition=lambda policy: policy.has_at_least_viewer_role())
-
-KernelPolicy.allow_actions([
-    constants.UPDATE,
-], scopes=[
-    OauthScope.CLIENT_PRIVATE,
-], condition=lambda policy: policy.has_at_least_editor_role_and_notebook_edit_access())
-
-KernelPolicy.allow_read(KernelPresenter.default_attributes + [], scopes=[
-    OauthScope.CLIENT_PRIVATE,
-], on_action=[
-    constants.LIST,
-], condition=lambda policy: policy.has_at_least_viewer_role())
-
-
-KernelPolicy.allow_read(KernelPresenter.default_attributes + [], scopes=[
-    OauthScope.CLIENT_PRIVATE,
-], on_action=[
-    constants.UPDATE,
-], condition=lambda policy: policy.has_at_least_editor_role_and_notebook_edit_access())
-
-KernelPolicy.allow_write([
-  'action_type',
-], scopes=[
-    OauthScope.CLIENT_PRIVATE,
-], on_action=[
-    constants.UPDATE,
-], condition=lambda policy: policy.has_at_least_editor_role_and_notebook_edit_access())
+KernelPolicy.allow_actions(
+    [
+        constants.LIST,
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    condition=lambda policy: policy.has_at_least_viewer_role(),
+)
+
+KernelPolicy.allow_actions(
+    [
+        constants.UPDATE,
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    condition=lambda policy: policy.has_at_least_editor_role_and_notebook_edit_access(),
+)
+
+KernelPolicy.allow_read(
+    KernelPresenter.default_attributes + [],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        constants.LIST,
+    ],
+    condition=lambda policy: policy.has_at_least_viewer_role(),
+)
+
+
+KernelPolicy.allow_read(
+    KernelPresenter.default_attributes + [],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        constants.UPDATE,
+    ],
+    condition=lambda policy: policy.has_at_least_editor_role_and_notebook_edit_access(),
+)
+
+KernelPolicy.allow_write(
+    [
+        'action_type',
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        constants.UPDATE,
+    ],
+    condition=lambda policy: policy.has_at_least_editor_role_and_notebook_edit_access(),
+)