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

mage_ai/kernels/models.py

@@ -0,0 +1,188 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from inspect import isawaitable
+from typing import Dict, List, Optional, Tuple
+
+from mage_ai.kernels.utils import (
+    find_ipykernel_launchers_info,
+    is_kernel_process_active,
+    terminate_process,
+)
+from mage_ai.shared.models import BaseDataClass
+
+
+@dataclass
+class PCPUTimes(BaseDataClass):
+    user: Optional[float] = None
+    system: Optional[float] = None
+    children_user: Optional[float] = None
+    children_system: Optional[float] = None
+    iowait: Optional[float] = None
+
+
+@dataclass
+class PMemoryInfo(BaseDataClass):
+    rss: Optional[int] = None
+    vms: Optional[int] = None
+    shared: Optional[int] = None
+    text: Optional[int] = None
+    lib: Optional[int] = None
+    data: Optional[int] = None
+    dirty: Optional[int] = None
+
+
+@dataclass
+class POpenFile(BaseDataClass):
+    path: Optional[str] = None
+    fd: Optional[int] = None
+    position: Optional[int] = None
+    mode: Optional[str] = None
+    flags: Optional[int] = None
+
+
+@dataclass
+class Addr(BaseDataClass):
+    ip: Optional[str] = None
+    port: Optional[int] = None
+
+
+@dataclass
+class PConn(BaseDataClass):
+    fd: Optional[int] = None
+    family: Optional[str] = None
+    type: Optional[str] = None
+    laddr: Optional[Addr] = None
+    raddr: Optional[Addr] = None
+    status: Optional[str] = None
+
+
+@dataclass
+class KernelProcess(BaseDataClass):
+    active: Optional[bool] = None
+    cmdline: Optional[str] = None
+    connection_file: Optional[str] = None
+    connections: Optional[List[PConn]] = None
+    cpu: Optional[float] = None
+    cpu_times: Optional[PCPUTimes] = None
+    create_time: Optional[float] = None
+    exe: Optional[str] = None
+    memory: Optional[float] = None
+    memory_info: Optional[PMemoryInfo] = None
+    name: Optional[str] = None
+    num_threads: Optional[str] = None
+    open_files: Optional[List[POpenFile]] = None
+    pid: Optional[int] = None
+    ppid: Optional[int] = None
+    status: Optional[str] = None
+    username: Optional[str] = None
+
+    def __post_init__(self):
+        self.serialize_attribute_class('memory_info', PMemoryInfo)
+        self.serialize_attribute_classes('connections', PConn)
+        self.serialize_attribute_classes('cpu_times', PCPUTimes)
+        self.serialize_attribute_classes('open_files', POpenFile)
+
+        self.cpu = (
+            sum([
+                float(val) if val is not None else 0
+                for val in [self.cpu_times.user, self.cpu_times.system]
+            ])
+            if self.cpu_times
+            else 0
+        )
+        self.memory = int(self.memory_info.rss) if self.memory_info and self.memory_info.rss else 0
+
+    @classmethod
+    def load_all(cls, check_active_status: bool = False) -> List['KernelProcess']:
+        return [
+            cls.load(**d)
+            for d in find_ipykernel_launchers_info(check_active_status=check_active_status)
+        ]
+
+    @classmethod
+    def terminate_inactive(
+        cls,
+        process_dicts: Optional[List[Dict]] = None,
+    ) -> Tuple[int, int]:
+        print('Terminating inactive kernels...')
+        memory_usage = []
+
+        if process_dicts:
+            arr = [cls.load(**d) for d in process_dicts]
+        else:
+            arr = cls.load_all(check_active_status=True)
+        if len(arr) >= 2:
+            for kernel_process in arr:
+                if not kernel_process.active:
+                    print(f'Terminating process {kernel_process.pid}...')
+                    if kernel_process.terminate():
+                        memory_usage.append(kernel_process.memory)
+
+        if len(memory_usage) == 0:
+            print('No inactive kernels found.')
+            return 0, 0
+
+        print(
+            f'{sum(memory_usage)} bytes of memory freed '
+            f'from {len(memory_usage)} inactive kernel(s).'
+        )
+        return len(memory_usage), sum(memory_usage)
+
+    def check_active_status(self) -> bool:
+        if not self.pid or not self.connection_file:
+            return False
+
+        self.active = is_kernel_process_active(self.pid, self.connection_file) or False
+        return self.active
+
+    def terminate(self) -> bool:
+        return self.pid is not None and terminate_process(self.pid)
+
+
+class KernelWrapper:
+    def __init__(self, kernel):
+        self.kernel = kernel
+        self.usage = None
+        self.active_kernels = None
+        self.inactive_kernels = None
+
+    async def prepare_usage(self):
+        try:
+            client = self.kernel.client()
+            session = self.kernel.session
+            control_channel = client.control_channel
+            usage_request = session.msg('usage_request', {})
+
+            control_channel.send(usage_request)
+            res = client.control_channel.get_msg(timeout=3)
+            if isawaitable(res):
+                # control_channel.get_msg may return a Future,
+                # depending on configured KernelManager class
+                res = await res
+            self.usage = res.get('content')
+            control_channel.stop()
+        except Exception:
+            pass
+
+    def is_ready(self) -> bool:
+        return self.usage is not None
+
+    def is_alive(self):
+        return self.kernel.is_alive()
+
+    @property
+    def kernel_id(self) -> str:
+        return self.kernel.kernel_id
+
+    @property
+    def kernel_name(self) -> str:
+        return self.kernel.kernel_name
+
+    def to_dict(self) -> Dict:
+        return {
+            'alive': self.is_alive(),
+            'id': self.kernel_id,
+            'name': self.kernel_name,
+            'usage': self.usage,
+        }

mage_ai/kernels/utils.py

@@ -0,0 +1,169 @@
+import asyncio
+import datetime
+import os
+from queue import Empty
+from typing import Any, Dict, List, Optional
+
+import jupyter_client
+import psutil
+
+
+def list_open_files(pid):
+    """
+    List open files by a process and their last modification time.
+
+    Parameters:
+    - pid (int): Process ID.
+    """
+    proc = psutil.Process(pid)
+    for item in proc.open_files():
+        file_path = item.path
+        try:
+            mod_time = os.path.getmtime(file_path)
+            print(f'File: {file_path}, Last Mod Time: {datetime.datetime.fromtimestamp(mod_time)}')
+        except Exception as e:
+            print(f'Could not retrieve modification time for {file_path}: {str(e)}')
+
+
+def is_cmdline_contains_ipykernel(cmdline, search_term='ipykernel_launcher'):
+    """Check if the command line (if exists) contains the specified search term."""
+    if cmdline and search_term in ' '.join(cmdline):
+        return True
+    return False
+
+
+async def find_ipykernel_launchers_info_async(timeout: int = 5) -> List[Dict]:
+    loop = asyncio.get_running_loop()
+    try:
+        # Run the synchronous function in a separate thread with a timeout
+        arr = await asyncio.wait_for(
+            loop.run_in_executor(None, lambda: find_ipykernel_launchers_info(True)), timeout
+        )
+        return arr
+    except asyncio.TimeoutError:
+        print(f'Operation timed out after {timeout} seconds')
+        return []
+
+
+def find_ipykernel_launchers_info(check_active_status: bool = False) -> List[Dict]:
+    """
+    ps aux | grep ipykernel_launcher
+    """
+    arr = []
+    for proc in psutil.process_iter(['pid', 'ppid', 'name', 'cmdline', 'memory_info']):
+        try:
+            # Check if cmdline is available and contains 'ipykernel_launcher'
+            cmdline = proc.info['cmdline']
+            if cmdline and 'ipykernel_launcher' in ' '.join(proc.info['cmdline']):
+                pid = proc.info['pid']
+                process = dict(pid=pid)
+                process.update(
+                    get_process_info(pid, check_active_status=check_active_status) or {}
+                )
+                arr.append(process)
+        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
+            pass  # Process has been terminated or access was denied
+    return arr
+
+
+def get_process_info(pid: int, check_active_status: bool = False) -> Optional[Dict]:
+    """
+    Retrieves information about a process given its PID.
+
+    Parameters:
+    - pid (int): Process ID.
+
+    Returns:
+    - dict: A dictionary containing process information, or None if the process does not exist.
+    """
+    try:
+        """
+        Most common to refer to the RSS (Resident Set Size) value
+        because it most accurately represents the amount of physical memory actually
+        occupied by the process
+        {
+            'memory_info': pmem(
+                rss=62554112,
+                vms=692289536,
+                shared=14155776,
+                text=4096,
+                lib=0,
+                data=130043904,
+                dirty=0,
+            ),
+        }
+        """
+        process = psutil.Process(int(pid))
+        cmdline = process.cmdline()
+        proc_info = {
+            'pid': process.pid,
+            'ppid': process.ppid(),
+            'name': process.name(),
+            'exe': process.exe(),
+            'cmdline': ' '.join(cmdline),
+            # If there are more than 1, kill the ones that are "sleeping"
+            # Kernels that are busy can still show status as "sleeping"
+            # ESTABLISHED
+            'status': process.status(),
+            'username': process.username(),
+            'create_time': process.create_time(),
+            'cpu_times': process.cpu_times(),
+            'memory_info': process.memory_info(),
+            'open_files': process.open_files(),
+            'connections': process.connections(),
+            'num_threads': process.num_threads(),
+        }
+
+        # Extract and add the connection file path
+        for arg in cmdline:
+            if arg and arg.endswith('.json'):
+                proc_info['connection_file'] = arg
+                break  # Break after finding the first .json argument
+
+        if check_active_status:
+            proc_info['active'] = is_kernel_process_active(pid, proc_info.get('connection_file'))
+        return proc_info
+    except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
+        return None
+
+
+def terminate_process(pid: int) -> bool:
+    """
+    Attempt to terminate the process with the given PID.
+    """
+    try:
+        process = psutil.Process(pid)
+        process.terminate()  # Requests the process to terminate
+        process.wait(timeout=3)  # Wait up to 3 seconds for the process to end
+        return True
+    except psutil.NoSuchProcess:
+        print(f'No process with PID {pid} exists.')
+        return False
+    except psutil.AccessDenied:
+        print(f'Permission denied to terminate process {pid}.')
+        return False
+    except psutil.TimeoutExpired:
+        print(f'Process {pid} did not terminate within the timeout period.')
+        # Optionally use process.kill() if you must forcibly stop the process
+        # return False
+    except Exception as e:
+        print(f'An unexpected error occurred: {e}')
+        return False
+    return False
+
+
+def is_kernel_process_active(pid: int, connection_file: Any) -> bool:
+    if not connection_file:
+        return False
+
+    km = jupyter_client.BlockingKernelClient()
+    km.load_connection_file(connection_file)
+
+    km.execute('1 + 1', connection_file)
+    try:
+        km.get_shell_msg(timeout=1)
+        return False
+    except Empty:
+        return True
+
+    return False

mage_ai/orchestration/concurrency.py

@@ -1,5 +1,9 @@
 from dataclasses import dataclass
 
+from mage_ai.settings.server import (
+    CONCURRENCY_CONFIG_BLOCK_RUN_LIMIT,
+    CONCURRENCY_CONFIG_PIPELINE_RUN_LIMIT,
+)
 from mage_ai.shared.config import BaseConfig
 
 
@@ -10,7 +14,7 @@ class OnLimitReached:
 
 @dataclass
 class ConcurrencyConfig(BaseConfig):
-    block_run_limit: int = None
-    pipeline_run_limit: int = None
+    block_run_limit: int = CONCURRENCY_CONFIG_BLOCK_RUN_LIMIT
+    pipeline_run_limit: int = CONCURRENCY_CONFIG_PIPELINE_RUN_LIMIT
     pipeline_run_limit_all_triggers: int = None
     on_pipeline_run_limit_reached: OnLimitReached = OnLimitReached.WAIT

mage_ai/orchestration/db/__init__.py

@@ -30,6 +30,7 @@ if OTEL_EXPORTER_OTLP_ENDPOINT:
 
 if is_test():
     db_connection_url = f'sqlite:///{TEST_DB}'
+    db_kwargs['connect_args']['check_same_thread'] = False
 elif not db_connection_url:
     pg_db_connection_url = get_postgres_connection_url()
 

mage_ai/orchestration/db/migrations/versions/0227396a216c_add_userproject_table.py

@@ -0,0 +1,38 @@
+"""Add UserProject table
+
+Revision ID: 0227396a216c
+Revises: 42a14d6143f1
+Create Date: 2024-04-01 11:31:34.680396
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '0227396a216c'
+down_revision = '42a14d6143f1'
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('user_project',
+    sa.Column('id', sa.Integer(), autoincrement=True, nullable=False),
+    sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=True),
+    sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=True),
+    sa.Column('user_id', sa.Integer(), nullable=True),
+    sa.Column('root_project_uuid', sa.String(length=255), nullable=True),
+    sa.Column('project_name', sa.String(length=255), nullable=True),
+    sa.Column('active', sa.Boolean(), nullable=True),
+    sa.ForeignKeyConstraint(['user_id'], ['user.id'], ),
+    sa.PrimaryKeyConstraint('id')
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('user_project')
+    # ### end Alembic commands ###

mage_ai/orchestration/db/migrations/versions/42a14d6143f1_update_token_column_type.py

@@ -0,0 +1,54 @@
+"""Update token column type
+
+Revision ID: 42a14d6143f1
+Revises: b9a2d6d0a2c7
+Create Date: 2024-04-12 15:19:52.639580
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '42a14d6143f1'
+down_revision = 'b9a2d6d0a2c7'
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('oauth2_access_token', schema=None) as batch_op:
+        batch_op.alter_column(
+            'token',
+            existing_type=sa.String(length=255),
+            type_=sa.Text(),
+            existing_nullable=True,
+        )
+        batch_op.alter_column(
+            'refresh_token',
+            existing_type=sa.String(length=255),
+            type_=sa.Text(),
+            existing_nullable=True,
+        )
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('oauth2_access_token', schema=None) as batch_op:
+        batch_op.alter_column(
+            'token',
+            existing_type=sa.Text(),
+            type_=sa.String(length=255),
+            existing_nullable=True,
+        )
+        batch_op.alter_column(
+            'refresh_token',
+            existing_type=sa.String(length=255),
+            type_=sa.Text(),
+            existing_nullable=True,
+        )
+
+    # ### end Alembic commands ###

mage_ai/orchestration/db/models/dynamic/controller.py

@@ -0,0 +1,67 @@
+from mage_ai.data_preparation.models.block.dynamic.models import DynamicRun
+
+
+def is_ready_to_process_data(pipeline, block_run, block_runs) -> bool:
+    """
+    Runs are “ready” can begin running when all their upstream blocks meet 1 of these conditions:
+        1. Condition set 1
+            - dynamic block or dynamic child
+            - not reduce_output
+            - status is running or all their upstream blocks are completed
+        2. Condition set 2
+            - not (dynamic block or dynamic child) or reduce_output
+            - block.all_upstream_blocks_completed(completed_block_uuids)
+    """
+    block = pipeline.get_block(block_run.block_uuid)
+    if not block:
+        raise Exception(f'Block {block_run.block_uuid} not found in pipeline {pipeline.uuid}')
+
+    if not block.upstream_blocks:
+        return True
+
+    mapping = {}
+    for br in block_runs:
+        block_inner = pipeline.get_block(br.block_uuid)
+        if block_inner.uuid not in mapping:
+            mapping[block_inner.uuid] = DynamicRun.load(block=block_inner)
+
+        run = mapping[block_inner.uuid]
+        run.add(br)
+
+    def __condition_1(upstream_block, block=block, mapping=mapping) -> bool:
+        run = mapping.get(upstream_block.uuid)
+        if run is None or run.block_run is None:
+            return False
+
+        return (
+            (
+                upstream_block.should_dynamically_generate_block(block)
+                or upstream_block.is_dynamic_child
+            )
+            and not upstream_block.should_reduce_output
+            and (
+                run.block_run.status
+                in [run.block_run.BlockRunStatus.COMPLETED, run.block_run.BlockRunStatus.RUNNING]
+            )
+        )
+
+    def __condition_2(upstream_block, block=block, mapping=mapping) -> bool:
+        run = mapping.get(upstream_block.uuid)
+        if run is None or run.block_run is None:
+            return False
+
+        return (
+            upstream_block.should_reduce_output
+            or not (
+                upstream_block.should_dynamically_generate_block(block)
+                or upstream_block.is_dynamic_child
+            )
+        ) and (run.block_run.status in [run.block_run.BlockRunStatus.COMPLETED])
+
+    results = [
+        (upstream_block.uuid, __condition_1(upstream_block), __condition_2(upstream_block))
+        for upstream_block in block.upstream_blocks
+    ]
+    ready = all([t[1] or t[2] for t in results])
+
+    return ready

mage_ai/orchestration/db/models/oauth.py

@@ -1,7 +1,7 @@
 import enum
 import re
 from datetime import datetime
-from typing import Callable, Dict, List, Union
+from typing import Callable, List, Union
 
 from sqlalchemy import (
     JSON,
@@ -12,6 +12,7 @@ from sqlalchemy import (
     ForeignKey,
     Integer,
     String,
+    Text,
     and_,
     asc,
     func,
@@ -29,7 +30,6 @@ from mage_ai.orchestration.constants import Entity
 from mage_ai.orchestration.db import db_connection, safe_db_query
 from mage_ai.orchestration.db.errors import ValidationError
 from mage_ai.orchestration.db.models.base import BaseModel
-from mage_ai.settings.repo import get_repo_path
 from mage_ai.shared.array import find
 from mage_ai.shared.environments import is_test
 from mage_ai.shared.hash import group_by, merge_dict
@@ -156,14 +156,8 @@ class User(BaseModel):
 
         return False
 
-    @property
-    def git_settings(self) -> Union[Dict, None]:
-        return self.get_git_settings()
-
-    def get_git_settings(self, repo_path: str = None):
+    def get_git_settings(self, repo_path: str):
         preferences = self.preferences or dict()
-        if not repo_path:
-            repo_path = get_repo_path()
         pref = preferences.get(repo_path)
         if pref:
             return (pref or {}).get('git_settings')
@@ -324,7 +318,7 @@ class Role(BaseModel):
     @classmethod
     @safe_db_query
     def create_default_roles(
-        self,
+        cls,
         entity: Entity = None,
         entity_id: str = None,
         name_func: Callable[[str], str] = None,
@@ -342,18 +336,18 @@ class Role(BaseModel):
             entity = Entity.GLOBAL
         permissions = Permission.create_default_permissions(entity=entity, entity_id=entity_id)
         mapping = {
-            self.DefaultRole.OWNER: Permission.Access.OWNER,
-            self.DefaultRole.ADMIN: Permission.Access.ADMIN,
-            self.DefaultRole.EDITOR: Permission.Access.EDITOR,
-            self.DefaultRole.VIEWER: Permission.Access.VIEWER,
+            cls.DefaultRole.OWNER: Permission.Access.OWNER,
+            cls.DefaultRole.ADMIN: Permission.Access.ADMIN,
+            cls.DefaultRole.EDITOR: Permission.Access.EDITOR,
+            cls.DefaultRole.VIEWER: Permission.Access.VIEWER,
         }
         for name, access in mapping.items():
             role_name = name
             if name_func is not None:
                 role_name = name_func(name)
-            role = self.query.filter(self.name == role_name).first()
+            role = cls.query.filter(Role.name == role_name).first()
             if not role:
-                self.create(
+                cls.create(
                     name=role_name,
                     permissions=[
                         Permission.query.filter(
@@ -880,10 +874,10 @@ class Oauth2AccessToken(BaseModel):
     expires = Column(DateTime(timezone=True))
     oauth2_application = relationship(Oauth2Application, back_populates='oauth2_access_tokens')
     oauth2_application_id = Column(Integer, ForeignKey('oauth2_application.id'))
-    token = Column(String(255), index=True, unique=True)
+    token = Column(Text, index=True, unique=True)
     user = relationship(User, back_populates='oauth2_access_tokens')
     user_id = Column(Integer, ForeignKey('user.id'))
-    refresh_token = Column(String(255))
+    refresh_token = Column(Text)
 
     def is_valid(self) -> bool:
         return self.token and \

mage_ai/orchestration/db/models/projects.py

@@ -0,0 +1,10 @@
+from sqlalchemy import Boolean, Column, ForeignKey, Integer, String
+
+from mage_ai.orchestration.db.models.base import BaseModel
+
+
+class UserProject(BaseModel):
+    user_id = Column(Integer, ForeignKey('user.id'))
+    root_project_uuid = Column(String(255), default=None)
+    project_name = Column(String(255), default=None)
+    active = Column(Boolean, default=False)