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

mage_ai/streaming/sinks/oracledb.py

@@ -0,0 +1,57 @@
+import json
+import time
+from dataclasses import dataclass
+from typing import Dict, List
+
+import pandas as pd
+
+from mage_ai.io.base import ExportWritePolicy
+from mage_ai.io.oracledb import OracleDB
+from mage_ai.shared.config import BaseConfig
+from mage_ai.streaming.sinks.base import BaseSink
+
+
+@dataclass
+class OracleDbConfig(BaseConfig):
+    user: str
+    password: int
+    host: str
+    port: int
+    table: str
+    service_name: str
+    verbose: bool = False
+    mode: str = 'thin'
+
+
+class OracleDbSink(BaseSink):
+    config_class = OracleDbConfig
+
+    def init_client(self):
+        self.oracledb_client = OracleDB(
+            user=self.config.user,
+            password=self.config.password,
+            host=self.config.host,
+            port=self.config.port,
+            service_name=self.config.service_name,
+            verbose=self.config.verbose,
+            mode=self.config.mode,
+        )
+        self.oracledb_client.open()
+
+    def write(self, message: Dict):
+        self.batch_write([message])
+
+    def batch_write(self, messages: List[Dict]):
+        if not messages:
+            return
+        self._print(
+            f'Batch ingest {len(messages)} records, time={time.time()}. Sample: {messages[0]}')
+        # Convert string formated dictionary array to array of dictionaries.
+        # To be used for dataframe creation.
+        dictionary_messages = [json.loads(doc_string) for doc_string in messages]
+        df = pd.DataFrame(dictionary_messages)
+        self.oracledb_client.export(
+            df,
+            self.config.table,
+            if_exists=ExportWritePolicy.APPEND,
+        )

mage_ai/streaming/sinks/sink_factory.py

@@ -59,6 +59,10 @@ class SinkFactory:
             from mage_ai.streaming.sinks.opensearch import OpenSearchSink
 
             return OpenSearchSink(config, **kwargs)
+        elif connector_type == SinkType.ORACLEDB:
+            from mage_ai.streaming.sinks.oracledb import OracleDbSink
+
+            return OracleDbSink(config, **kwargs)
         elif connector_type == SinkType.POSTGRES:
             from mage_ai.streaming.sinks.postgres import PostgresSink
 

mage_ai/system/constants.py

@@ -0,0 +1,14 @@
+from enum import Enum
+
+
+class LogType(str, Enum):
+    END = 'END'
+    GENERIC = 'GENERIC'
+    MEMORY = 'MEMORY'
+    METRICS = 'METRICS'
+    PROCESS = 'PROCESS'
+    START = 'START'
+
+
+METRICS_DIRECTORY = 'metrics'
+SYSTEM_DIRECTORY = 'system'

mage_ai/system/memory/constants.py

@@ -0,0 +1 @@
+MEMORY_DIRECTORY = 'memory'

mage_ai/system/memory/manager.py

@@ -0,0 +1,174 @@
+import os
+from datetime import datetime
+from typing import Dict, Optional
+
+import aiofiles
+
+from mage_ai.data.constants import SUPPORTED_VARIABLE_TYPES
+from mage_ai.settings.server import (
+    MEMORY_MANAGER_PANDAS_VERSION,
+    MEMORY_MANAGER_POLARS_VERSION,
+    MEMORY_MANAGER_V2,
+    MEMORY_MANAGER_VERSION,
+    SYSTEM_LOGS_PARTITIONS,
+    SYSTEM_LOGS_POLL_INTERVAL,
+)
+from mage_ai.shared.files import makedirs_async, makedirs_sync
+from mage_ai.shared.singletons.memory import get_memory_manager_controller
+from mage_ai.system.constants import LogType
+from mage_ai.system.memory.constants import MEMORY_DIRECTORY
+from mage_ai.system.memory.utils import (
+    current_memory_usage,
+    format_log_message,
+    get_log_directory,
+    monitor_memory_usage,
+    monitor_memory_usage_async,
+)
+
+
+class MemoryManager:
+    def __init__(
+        self,
+        scope_uuid: str,
+        process_uuid: str,
+        poll_interval: Optional[float] = None,
+        repo_path: Optional[str] = None,
+        metadata: Optional[Dict] = None,
+    ):
+        """
+        Usage example with EnhancedMemoryUsage class for enhanced memory usage tracking:
+
+        def run_operations():
+            # Simulating memory usage scenario
+            a = [1] * (10**6)
+            time.sleep(1)  # Simulate time delay
+            b = [2] * (10**6)
+            del a  # Manual deletion to potentially influence memory usage metrics
+
+        with EnhancedMemoryUsage() as emu:
+            result = emu.profile_memory(run_operations)
+        print(emu.report())
+        """
+        self.monitor = None
+        self.poll_interval = poll_interval or SYSTEM_LOGS_POLL_INTERVAL
+        self.process_uuid = process_uuid  # [process_uuid]
+        self.repo_path = repo_path
+        self.scope_uuid = scope_uuid  # [pipeline_uuid]/[block_uuid]
+        self._log_path = None
+        self._metadata = metadata or {}
+
+    @property
+    def log_path(self) -> str:
+        """
+        /root/.mage_data/[project]
+            /system/metrics
+                /pipelines/[pipeline_uuid]/[block_uuid]
+                    /[date]/[hour]
+                        /[metric]
+                            /[process_uuid].log
+        """
+        if not self._log_path:
+            now = datetime.utcnow()
+
+            datetime_partitions = []
+            for partition_name in SYSTEM_LOGS_PARTITIONS:
+                if 'ds' == partition_name:
+                    datetime_partitions.append(now.strftime('%Y-%m-%d'))
+                elif 'hr' == partition_name:
+                    datetime_partitions.append(now.strftime('%H'))
+
+            self._log_path = os.path.join(
+                get_log_directory(self.scope_uuid, repo_path=self.repo_path),
+                *datetime_partitions,
+                MEMORY_DIRECTORY,
+                f'{self.process_uuid}.log',
+            )
+        return self._log_path
+
+    @property
+    def metadata(self) -> Dict:
+        supported_variable_types = [v.value for v in SUPPORTED_VARIABLE_TYPES]
+
+        self._metadata.update(
+            dict(
+                memory_manager_version=MEMORY_MANAGER_VERSION,
+                memory_manager_pandas_version=MEMORY_MANAGER_PANDAS_VERSION,
+                memory_manager_polars_version=MEMORY_MANAGER_POLARS_VERSION,
+                supported_variable_types=supported_variable_types,
+            ),
+        )
+        return self._metadata
+
+    def __enter__(self):
+        if not MEMORY_MANAGER_V2:
+            return self
+
+        makedirs_sync(os.path.dirname(self.log_path))
+
+        self.stop()
+        self.start(
+            *monitor_memory_usage(
+                callback=self.__write_sync,
+                interval_seconds=self.poll_interval,
+            )
+        )
+
+        with open(self.log_path, 'a') as f:
+            f.write(format_log_message(log_type=LogType.START, metadata=self.metadata))
+            f.write(format_log_message(message=current_memory_usage()))
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if not MEMORY_MANAGER_V2:
+            return
+
+        try:
+            with open(self.log_path, 'a') as f:
+                f.write(format_log_message(message=current_memory_usage()))
+                f.write(format_log_message(log_type=LogType.END, metadata=self.metadata))
+        except Exception as err:
+            print(f'[MemoryManager] Error writing memory usage logs: {err}')
+        finally:
+            self.stop()
+
+    async def __aenter__(self):
+        await makedirs_async(os.path.dirname(self.log_path))
+
+        self.stop()
+        self.start(
+            *(
+                await monitor_memory_usage_async(
+                    callback=self.__write_async,
+                    interval_seconds=self.poll_interval,
+                )
+            )
+        )
+
+        async with aiofiles.open(self.log_path, mode='a', encoding='utf-8') as fp:
+            await fp.write(format_log_message(log_type=LogType.START, metadata=self.metadata))
+            await fp.write(format_log_message(message=current_memory_usage()))
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        try:
+            async with aiofiles.open(self.log_path, mode='a', encoding='utf-8') as fp:
+                await fp.write(format_log_message(message=current_memory_usage()))
+                await fp.write(format_log_message(log_type=LogType.END, metadata=self.metadata))
+        finally:
+            self.stop()
+
+    def __write_sync(self, memory: float) -> None:
+        with open(self.log_path, 'a') as f:
+            f.write(format_log_message(message=memory))
+
+    async def __write_async(self, memory: float) -> None:
+        async with aiofiles.open(self.log_path, mode='a', encoding='utf-8') as fp:
+            await fp.write(format_log_message(message=memory))
+
+    def start(self, stop_event, monitor_thread) -> None:
+        get_memory_manager_controller().add_event_monitor(
+            self.log_path, stop_event, monitor_thread
+        )
+
+    def stop(self) -> None:
+        get_memory_manager_controller().stop_event(self.log_path)

mage_ai/system/memory/presenters.py

@@ -0,0 +1,158 @@
+import os
+from pathlib import Path
+from typing import Callable, Dict, List, Optional
+
+import polars as pl
+
+from mage_ai.settings.server import SYSTEM_LOGS_PARTITIONS
+from mage_ai.system.constants import LogType
+from mage_ai.system.memory.constants import MEMORY_DIRECTORY
+from mage_ai.system.memory.utils import get_log_directory
+
+DEFAULT_SCHEMA = {
+    'block_uuid': pl.Utf8,
+    'log_type': pl.Utf8,
+    'message': pl.Utf8,
+    'metadata': pl.Utf8,
+    'metric': pl.Utf8,
+    'pipeline_uuid': pl.Utf8,
+    'process_uuid': pl.Utf8,
+    'scope_uuid': pl.Utf8,
+    'timestamp': pl.Int64,
+    'value': pl.Float64,
+}
+
+
+def parse_line(line: str) -> dict:
+    # Initial parsing by removing brackets and splitting
+    parts = line.strip('[]').split(']', 1)
+    timestamp = parts[0]
+    rest = parts[1].strip() if len(parts) > 1 else ''
+
+    data = {}
+    for k in DEFAULT_SCHEMA.keys():
+        if k not in ['timestamp', 'value']:
+            data[k] = None
+    data['timestamp'] = int(timestamp)
+    data['value'] = float('nan')
+
+    # Extract the log type and the rest
+    if rest:
+        segments = rest.split(' ', 2)  # Expecting type, potential message, and metadata
+        log_type = segments[0].replace('[', '').replace(']', '')
+        data['log_type'] = log_type
+
+        # The rest could contain a message and/or metadata
+        rest_segments = ' '.join(segments[1:]).split('=', 1)
+
+        message = None
+        if len(rest_segments) > 1:  # There's metadata
+            # Attempt to identify potential message before metadata
+            pre_meta = rest_segments[0].rpartition(' ')
+
+            message = pre_meta[0].strip()
+            # Reconstruct metadata split and parse
+            metadata_content = pre_meta[2] + '=' + rest_segments[1]
+            data['metadata'] = metadata_content
+        else:
+            # No metadata, the rest could be a message
+            message = rest_segments[0].strip()
+
+        data['message'] = message
+        if LogType.MEMORY.value == log_type:
+            data['metric'] = LogType.MEMORY.value
+            if message.isdigit():
+                data['value'] = float(message)
+
+    return data
+
+
+def parse_log_to_polars(log_contents: str, process_uuid: str) -> pl.DataFrame:
+    data_records: List[dict] = []
+    lines = log_contents.strip().split('\n')
+
+    for line in lines:
+        log_data = parse_line(line)
+        log_data['process_uuid'] = process_uuid  # Add process_uuid to each log data record
+        data_records.append(log_data)
+
+    return pl.DataFrame(data_records)
+
+
+def parse_directory_to_df(
+    directory_path: str, extra_columns_parser: Optional[Callable[[str, str], dict]] = None
+) -> pl.DataFrame:
+    """
+    Parses a directory of log files into a Polars DataFrame.
+
+    Parameters:
+    - directory_path: Path to the directory containing log files.
+    - extra_columns_parser: [Optional] A callable function that receives the file path of
+        each log and returns a dictionary of extra columns and their values.
+
+    Returns:
+    - A Polars DataFrame containing the merged logs with potential extra columns.
+    """
+    all_dataframes = []
+    for root, _dirs, files in os.walk(directory_path):
+        for file_name in files:
+            if file_name.endswith('.log'):
+                process_uuid = file_name.replace('.log', '')
+                file_path = os.path.join(root, file_name)
+                with open(file_path, 'r') as file:
+                    log_contents = file.read()
+                    df = parse_log_to_polars(log_contents, process_uuid)
+
+                    extra_columns = pipeline_block_columns_parser(directory_path, file_path)
+
+                    if extra_columns_parser:
+                        extra_columns.update(extra_columns_parser(directory_path, file_path) or {})
+
+                    if extra_columns:
+                        for col, val in extra_columns.items():
+                            df = df.with_columns((pl.lit(val)).alias(col))
+
+                    all_dataframes.append(df)
+
+    if all_dataframes:
+        final_df = pl.concat(all_dataframes)
+        return final_df
+    else:
+        return pl.DataFrame(
+            [],
+            schema=DEFAULT_SCHEMA,
+        )
+
+
+def to_dataframe(
+    scope_uuid: str, extra_columns_parser: Optional[Callable[[str, str], dict]] = None
+) -> pl.DataFrame:
+    log_directory = get_log_directory(scope_uuid)
+    return parse_directory_to_df(
+        log_directory,
+        extra_columns_parser=extra_columns_parser,
+    )
+
+
+def pipeline_block_columns_parser(directory_path: str, file_path: str) -> Dict:
+    try:
+        parts = Path(file_path).relative_to(Path(directory_path)).parts
+        if MEMORY_DIRECTORY in parts:
+            parts = parts[: parts.index(MEMORY_DIRECTORY)]
+
+        parts = parts[: -1 * len(SYSTEM_LOGS_PARTITIONS)]
+
+        data = {}
+
+        # If the pipeline is in the directory path, there won’t be a pipeline UUID in the parts
+        if len(parts) >= 1:
+            data['scope_uuid'] = os.path.join(*parts)
+            if len(parts) >= 2:
+                data['pipeline_uuid'] = os.path.join(*parts[:1])
+                data['block_uuid'] = os.path.join(*parts[1:])
+            else:
+                data['block_uuid'] = os.path.join(*parts)
+
+        return data
+    except ValueError:
+        return {}

mage_ai/system/memory/process.py

@@ -0,0 +1,216 @@
+import functools
+import gc
+import os
+import sys
+import tracemalloc
+from multiprocessing import Process, Queue
+from typing import Any, Dict, List, Tuple
+
+import psutil
+
+
+def track_usage(queue):
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            process_id = os.getpid()
+            # Memory usage before
+            proc = psutil.Process(process_id)
+            mem_before = proc.memory_info().rss
+
+            # Track objects before the function call
+            objects_before = set(id(obj) for obj in gc.get_objects())
+
+            result = func(*args, **kwargs)  # Call the original function
+
+            objects_after = set(id(obj) for obj in gc.get_objects())
+            new_objects = [
+                [str(type(obj)), sys.getsizeof(obj)]
+                for obj in gc.get_objects()
+                if id(obj) in objects_after - objects_before
+            ]
+
+            # Memory usage after
+            mem_after = proc.memory_info().rss
+            print(
+                f'[{func.__name__}] Memory Usage: '
+                f'Before={mem_before} After={mem_after} Difference={mem_after - mem_before} bytes'
+            )
+
+            # Put tracking object/details into the queue
+            print(len(new_objects))
+            memory_usage_info = {
+                'function_name': func.__name__,
+                'memory_used': mem_after - mem_before,
+                'process_id': process_id,
+                'uncollected_objects': new_objects,
+            }
+            queue.put(memory_usage_info)
+            print('Added to queue')
+
+            return result
+
+        return wrapper
+
+    return decorator
+
+
+def wrapped_func(queue, code_string, func_name, *args, **kwargs):
+    # Prepare a namespace dict where the executed code will live
+    namespace = (
+        globals().copy()
+    )  # This copies the current global symbols to the namespace
+
+    # Execute the code_string within the namespace
+    exec(
+        code_string, namespace, namespace
+    )  # Now both global and local scopes are the same
+
+    # Now, retrieve the function from this namespace
+    actual_function = namespace.get(func_name)
+
+    decorated_func = track_usage(queue)(actual_function)
+
+    try:
+        return decorated_func(*args, **kwargs)
+    except Exception as e:
+        queue.put({'error': str(e)})
+    finally:
+        # Signal that this process is done
+        queue.put(None)
+
+
+def begin_working(jobs: List[Tuple[str, str, List[Any], Dict[str, Any]]]):
+    gc.collect()  # Force a collection to start from a clean slate
+
+    queue = Queue()
+
+    tracemalloc.start()
+
+    # Take a snapshot before the function call
+    snapshot_before = tracemalloc.take_snapshot()
+
+    num_processes = 0
+    for file_path, func_name, args, kwargs in jobs:
+        print(f'Processing {func_name} in {file_path}')
+        with open(file_path, 'r') as file:
+            code_string = file.read()
+
+        p = Process(
+            target=wrapped_func,
+            args=(
+                queue,
+                code_string,
+                func_name,
+                *args,
+            ),
+            kwargs=kwargs,
+        )
+        p.start()
+        num_processes += 1
+        p.join(timeout=10 * 60)  # Adjust the timeout as needed
+        if p.is_alive():
+            print(f'Process {p.name} did not terminate as expected.')
+            p.terminate()  # Forcefully terminate
+
+    # Retrieve and print memory usage information
+    uncollected_objects = []
+    func_names = []
+    process_ids = []
+    memory_uses = []
+
+    sentinel_count = 0
+    while sentinel_count < num_processes:
+        memory_usage_info = queue.get()
+        if memory_usage_info is None:
+            sentinel_count += 1
+            continue  # Skip the rest of the loop
+        # Existing code to process non-sentinel items...
+
+        if 'error' in memory_usage_info:
+            raise Exception(memory_usage_info['error'])
+
+        memory_uses.append(memory_usage_info.get('memory_used'))
+        func_names.append(memory_usage_info.get('function_name'))
+        process_ids.append(memory_usage_info.get('process_id'))
+        uncollected_objects.extend(memory_usage_info.get('uncollected_objects') or [])
+
+    print(f'Total memory used: {sum(memory_uses)} bytes across:')
+    for func_name, process_id in zip(func_names, process_ids):
+        print(f'\t{func_name} (PID={process_id}): {memory_uses.pop(0)} bytes')
+
+    print('\n')
+
+    sizes_by_type = {}
+    # Find new objects created during the function call
+    for obj_type, obj_size in uncollected_objects:
+        if obj_type not in sizes_by_type:
+            sizes_by_type[obj_type] = 0
+        sizes_by_type[obj_type] += obj_size
+
+    print(
+        f'Objects created that wasn’t uncollected: {len(uncollected_objects)} '
+        f'({sum(list(sizes_by_type.values()))})'
+    )
+    for k, v in sizes_by_type.items():
+        print(f'\t{k}: {v} bytes')
+
+    del uncollected_objects
+    del func_names
+    del process_ids
+    del memory_uses
+    del sizes_by_type
+
+    # Take another snapshot after
+    snapshot_after = tracemalloc.take_snapshot()
+
+    # Compare the two snapshots to see what has been allocated in between
+    """
+    This comparison highlights allocations, not necessarily what's "still in use" because it
+    doesn't track deallocations directly. If an object is allocated and then deallocated between
+    the snapshots, it won't appear. However, if there's a positive `size_diff`,
+    it suggests memory allocated during the snapshot window hasn't been released by the time of
+    the second snapshot.
+    """
+    stats = snapshot_after.compare_to(snapshot_before, 'lineno')
+
+    # A more focused approach to filtering the stats
+    for stat in stats:
+        trace = stat.traceback.format()
+        # Check if any line in the traceback belongs to your function's code file
+        if any('worker' in line for line in trace):
+            for line in trace:
+                print(line)
+            print(stat)
+            """
+            For understanding "memory still being used," you're mainly concerned with `size_diff`
+            and allocations (`size` and `count`) not countered by deallocations within the
+            window of your snapshots.
+            """
+            # print('\t', stat.traceback.format())
+            # print('\t', stat.count)
+            # print('\t', stat.count_diff)
+            # print('\t', stat.size)
+            print('\t', stat.size_diff)
+            print('\n')
+
+    gc.collect()  # Force a collection to start from a clean slate
+
+
+if __name__ == '__main__':
+    begin_working(
+        [
+            (
+                '/Users/dangerous/Code/materia/mage-ai/mage_ai/system/memory/samples.py',
+                'worker',
+                [1],
+                dict(another_arg=2),
+            ),
+            (
+                '/Users/dangerous/Code/materia/mage-ai/mage_ai/system/memory/samples.py',
+                'create_objects_example',
+                [1000],
+                dict(),
+            ),
+        ]
+    )

mage_ai/system/memory/samples.py

@@ -0,0 +1,13 @@
+def worker(*args, **kwargs):
+    # Example use of arguments inside the function
+    print(f'Worker received: some_arg={args}, another_arg={kwargs}')
+    # Simulate work with objects
+    return [i for i in range(1000000)]
+
+
+# Example function to track
+def create_objects_example(count, *args, **kwargs):
+    a = [1] * count
+    b = 'A string' * count
+    c = {'key': 'value' * count}
+    return a, b, c