moru 0.1.0__py3-none-any.whl

moru/template_async/build_api.py

@@ -0,0 +1,202 @@
+import asyncio
+from types import TracebackType
+from typing import Callable, Literal, Optional, List, Union
+
+import httpx
+
+from moru.api import handle_api_exception
+from moru.api.client.api.templates import (
+    post_v3_templates,
+    get_templates_template_id_files_hash,
+    post_v_2_templates_template_id_builds_build_id,
+    get_templates_template_id_builds_build_id_status,
+)
+from moru.api.client.client import AuthenticatedClient
+from moru.api.client.models import (
+    TemplateBuildRequestV3,
+    TemplateBuildStartV2,
+    TemplateBuildFileUpload,
+    TemplateBuild,
+    Error,
+)
+from moru.exceptions import BuildException, FileUploadException
+from moru.template.logger import LogEntry
+from moru.template.types import TemplateType
+from moru.template.utils import get_build_step_index, tar_file_stream
+
+
+async def request_build(
+    client: AuthenticatedClient, name: str, cpu_count: int, memory_mb: int
+):
+    res = await post_v3_templates.asyncio_detailed(
+        client=client,
+        body=TemplateBuildRequestV3(
+            alias=name,
+            cpu_count=cpu_count,
+            memory_mb=memory_mb,
+        ),
+    )
+
+    if res.status_code >= 300:
+        raise handle_api_exception(res, BuildException)
+
+    if isinstance(res.parsed, Error):
+        raise BuildException(f"API error: {res.parsed.message}")
+
+    if res.parsed is None:
+        raise BuildException("Failed to request build")
+
+    return res.parsed
+
+
+async def get_file_upload_link(
+    client: AuthenticatedClient,
+    template_id: str,
+    files_hash: str,
+    stack_trace: Optional[TracebackType] = None,
+) -> TemplateBuildFileUpload:
+    res = await get_templates_template_id_files_hash.asyncio_detailed(
+        template_id=template_id,
+        hash_=files_hash,
+        client=client,
+    )
+
+    if res.status_code >= 300:
+        raise handle_api_exception(res, FileUploadException, stack_trace)
+
+    if isinstance(res.parsed, Error):
+        raise FileUploadException(f"API error: {res.parsed.message}").with_traceback(
+            stack_trace
+        )
+
+    if res.parsed is None:
+        raise FileUploadException("Failed to get file upload link").with_traceback(
+            stack_trace
+        )
+
+    return res.parsed
+
+
+async def upload_file(
+    api_client: AuthenticatedClient,
+    file_name: str,
+    context_path: str,
+    url: str,
+    ignore_patterns: List[str],
+    resolve_symlinks: bool,
+    stack_trace: Optional[TracebackType],
+):
+    try:
+        tar_buffer = tar_file_stream(
+            file_name, context_path, ignore_patterns, resolve_symlinks
+        )
+
+        client = api_client.get_async_httpx_client()
+        response = await client.put(url, content=tar_buffer.getvalue())
+        response.raise_for_status()
+    except httpx.HTTPStatusError as e:
+        raise FileUploadException(f"Failed to upload file: {e}").with_traceback(
+            stack_trace
+        )
+    except Exception as e:
+        raise FileUploadException(f"Failed to upload file: {e}").with_traceback(
+            stack_trace
+        )
+
+
+async def trigger_build(
+    client: AuthenticatedClient,
+    template_id: str,
+    build_id: str,
+    template: TemplateType,
+) -> None:
+    # Convert template dict to TemplateBuildStartV2 model using from_dict
+    template_data = TemplateBuildStartV2.from_dict(template)
+
+    res = await post_v_2_templates_template_id_builds_build_id.asyncio_detailed(
+        template_id=template_id,
+        build_id=build_id,
+        client=client,
+        body=template_data,
+    )
+
+    if res.status_code >= 300:
+        raise handle_api_exception(res, BuildException)
+
+
+async def get_build_status(
+    client: AuthenticatedClient, template_id: str, build_id: str, logs_offset: int
+) -> TemplateBuild:
+    res = await get_templates_template_id_builds_build_id_status.asyncio_detailed(
+        template_id=template_id,
+        build_id=build_id,
+        client=client,
+        logs_offset=logs_offset,
+    )
+
+    if res.status_code >= 300:
+        raise handle_api_exception(res, BuildException)
+
+    if isinstance(res.parsed, Error):
+        raise BuildException(f"API error: {res.parsed.message}")
+
+    if res.parsed is None:
+        raise BuildException("Failed to get build status")
+
+    return res.parsed
+
+
+async def wait_for_build_finish(
+    client: AuthenticatedClient,
+    template_id: str,
+    build_id: str,
+    on_build_logs: Optional[Callable[[LogEntry], None]] = None,
+    logs_refresh_frequency: float = 0.2,
+    stack_traces: List[Union[TracebackType, None]] = [],
+):
+    logs_offset = 0
+    status: Literal["building", "waiting", "ready", "error"] = "building"
+
+    while status in ["building", "waiting"]:
+        build_status = await get_build_status(
+            client, template_id, build_id, logs_offset
+        )
+
+        logs_offset += len(build_status.log_entries)
+
+        for log_entry in build_status.log_entries:
+            if on_build_logs:
+                on_build_logs(
+                    LogEntry(
+                        timestamp=log_entry.timestamp,
+                        level=log_entry.level.value,
+                        message=log_entry.message,
+                    )
+                )
+
+        status = build_status.status.value
+
+        if status == "ready":
+            return
+
+        elif status == "waiting":
+            pass
+
+        elif status == "error":
+            traceback = None
+            if build_status.reason and build_status.reason.step:
+                # Find the corresponding stack trace for the failed step
+                step_index = get_build_step_index(
+                    build_status.reason.step, len(stack_traces)
+                )
+                if step_index < len(stack_traces):
+                    traceback = stack_traces[step_index]
+
+            raise BuildException(
+                build_status.reason.message if build_status.reason else "Build failed"
+            ).with_traceback(traceback)
+
+        # Wait for a short period before checking the status again
+        await asyncio.sleep(logs_refresh_frequency)
+
+    raise BuildException("Unknown build error occurred.")

moru/template_async/main.py

@@ -0,0 +1,366 @@
+import os
+from datetime import datetime
+from typing import Callable, Optional
+
+from moru.api.client.client import AuthenticatedClient
+from moru.connection_config import ConnectionConfig
+from moru.template.consts import RESOLVE_SYMLINKS
+from moru.template.logger import LogEntry, LogEntryEnd, LogEntryStart
+from moru.template.main import TemplateBase, TemplateClass
+from moru.template.types import BuildInfo, InstructionType
+from moru.template.utils import read_dockerignore
+
+from .build_api import (
+    get_build_status,
+    get_file_upload_link,
+    request_build,
+    trigger_build,
+    upload_file,
+    wait_for_build_finish,
+)
+from moru.api.client_async import get_api_client
+
+
+class AsyncTemplate(TemplateBase):
+    """
+    Asynchronous template builder for Moru sandboxes.
+    """
+
+    @staticmethod
+    async def _build(
+        template: TemplateClass,
+        api_client: AuthenticatedClient,
+        alias: str,
+        cpu_count: int = 2,
+        memory_mb: int = 1024,
+        skip_cache: bool = False,
+        on_build_logs: Optional[Callable[[LogEntry], None]] = None,
+    ) -> BuildInfo:
+        """
+        Internal implementation of the template build process
+
+        :param template: The template to build
+        :param api_client: Authenticated API client
+        :param alias: Alias name for the template
+        :param cpu_count: Number of CPUs allocated to the sandbox
+        :param memory_mb: Amount of memory in MB allocated to the sandbox
+        :param skip_cache: If True, forces a complete rebuild ignoring cache
+        :param on_build_logs: Callback function to receive build logs during the build process
+        """
+        if skip_cache:
+            template._template._force = True
+
+        # Create template
+        if on_build_logs:
+            on_build_logs(
+                LogEntry(
+                    timestamp=datetime.now(),
+                    level="info",
+                    message=f"Requesting build for template: {alias}",
+                )
+            )
+
+        response = await request_build(
+            api_client,
+            name=alias,
+            cpu_count=cpu_count,
+            memory_mb=memory_mb,
+        )
+
+        template_id = response.template_id
+        build_id = response.build_id
+
+        if on_build_logs:
+            on_build_logs(
+                LogEntry(
+                    timestamp=datetime.now(),
+                    level="info",
+                    message=f"Template created with ID: {template_id}, Build ID: {build_id}",
+                )
+            )
+
+        instructions_with_hashes = template._template._instructions_with_hashes()
+
+        # Upload files
+        for index, file_upload in enumerate(instructions_with_hashes):
+            if file_upload["type"] != InstructionType.COPY:
+                continue
+
+            args = file_upload.get("args", [])
+            src = args[0] if len(args) > 0 else None
+            force_upload = file_upload.get("forceUpload")
+            files_hash = file_upload.get("filesHash", None)
+            resolve_symlinks = file_upload.get("resolveSymlinks", RESOLVE_SYMLINKS)
+
+            if src is None or files_hash is None:
+                raise ValueError("Source path and files hash are required")
+
+            stack_trace = None
+            if index + 1 < len(template._template._stack_traces):
+                stack_trace = template._template._stack_traces[index + 1]
+
+            file_info = await get_file_upload_link(
+                api_client, template_id, files_hash, stack_trace
+            )
+
+            if (force_upload and file_info.url) or (
+                file_info.present is False and file_info.url
+            ):
+                await upload_file(
+                    api_client,
+                    src,
+                    template._template._file_context_path,
+                    file_info.url,
+                    [
+                        *template._template._file_ignore_patterns,
+                        *read_dockerignore(template._template._file_context_path),
+                    ],
+                    resolve_symlinks,
+                    stack_trace,
+                )
+                if on_build_logs:
+                    on_build_logs(
+                        LogEntry(
+                            timestamp=datetime.now(),
+                            level="info",
+                            message=f"Uploaded '{src}'",
+                        )
+                    )
+            else:
+                if on_build_logs:
+                    on_build_logs(
+                        LogEntry(
+                            timestamp=datetime.now(),
+                            level="info",
+                            message=f"Skipping upload of '{src}', already cached",
+                        )
+                    )
+
+        if on_build_logs:
+            on_build_logs(
+                LogEntry(
+                    timestamp=datetime.now(),
+                    level="info",
+                    message="All file uploads completed",
+                )
+            )
+
+        # Start build
+        if on_build_logs:
+            on_build_logs(
+                LogEntry(
+                    timestamp=datetime.now(),
+                    level="info",
+                    message="Starting building...",
+                )
+            )
+
+        await trigger_build(
+            api_client,
+            template_id,
+            build_id,
+            template._template._serialize(instructions_with_hashes),
+        )
+
+        return BuildInfo(
+            alias=alias,
+            template_id=template_id,
+            build_id=build_id,
+        )
+
+    @staticmethod
+    async def build(
+        template: TemplateClass,
+        alias: str,
+        cpu_count: int = 2,
+        memory_mb: int = 1024,
+        skip_cache: bool = False,
+        on_build_logs: Optional[Callable[[LogEntry], None]] = None,
+        api_key: Optional[str] = None,
+        domain: Optional[str] = None,
+    ) -> BuildInfo:
+        """
+        Build and deploy a template to Moru infrastructure.
+
+        :param template: The template to build
+        :param alias: Alias name for the template
+        :param cpu_count: Number of CPUs allocated to the sandbox
+        :param memory_mb: Amount of memory in MB allocated to the sandbox
+        :param skip_cache: If True, forces a complete rebuild ignoring cache
+        :param on_build_logs: Callback function to receive build logs during the build process
+        :param api_key: Moru API key for authentication
+        :param domain: Domain of the Moru API
+
+        Example
+        ```python
+        from moru import AsyncTemplate
+
+        template = (
+            AsyncTemplate()
+            .from_python_image('3')
+            .copy('requirements.txt', '/home/user/')
+            .run_cmd('pip install -r /home/user/requirements.txt')
+        )
+
+        await AsyncTemplate.build(
+            template,
+            alias='my-python-env',
+            cpu_count=2,
+            memory_mb=1024
+        )
+        ```
+        """
+        try:
+            if on_build_logs:
+                on_build_logs(
+                    LogEntryStart(
+                        timestamp=datetime.now(),
+                        message="Build started",
+                    )
+                )
+
+            domain = domain or os.environ.get("MORU_DOMAIN", "moru.io")
+            config = ConnectionConfig(
+                domain=domain, api_key=api_key or os.environ.get("MORU_API_KEY")
+            )
+            api_client = get_api_client(
+                config,
+                require_api_key=True,
+                require_access_token=False,
+            )
+
+            data = await AsyncTemplate._build(
+                template,
+                api_client,
+                alias,
+                cpu_count,
+                memory_mb,
+                skip_cache,
+                on_build_logs,
+            )
+
+            if on_build_logs:
+                on_build_logs(
+                    LogEntry(
+                        timestamp=datetime.now(),
+                        level="info",
+                        message="Waiting for logs...",
+                    )
+                )
+
+            await wait_for_build_finish(
+                api_client,
+                data.template_id,
+                data.build_id,
+                on_build_logs,
+                logs_refresh_frequency=TemplateBase._logs_refresh_frequency,
+                stack_traces=template._template._stack_traces,
+            )
+
+            return data
+        finally:
+            if on_build_logs:
+                on_build_logs(
+                    LogEntryEnd(
+                        timestamp=datetime.now(),
+                        message="Build finished",
+                    )
+                )
+
+    @staticmethod
+    async def build_in_background(
+        template: TemplateClass,
+        alias: str,
+        cpu_count: int = 2,
+        memory_mb: int = 1024,
+        skip_cache: bool = False,
+        on_build_logs: Optional[Callable[[LogEntry], None]] = None,
+        api_key: Optional[str] = None,
+        domain: Optional[str] = None,
+    ) -> BuildInfo:
+        """
+        Build and deploy a template to Moru infrastructure without waiting for completion.
+
+        :param template: The template to build
+        :param alias: Alias name for the template
+        :param cpu_count: Number of CPUs allocated to the sandbox
+        :param memory_mb: Amount of memory in MB allocated to the sandbox
+        :param skip_cache: If True, forces a complete rebuild ignoring cache
+        :param api_key: Moru API key for authentication
+        :param domain: Domain of the Moru API
+        :return: BuildInfo containing the template ID and build ID
+
+        Example
+        ```python
+        from moru import AsyncTemplate
+
+        template = (
+            AsyncTemplate()
+            .from_python_image('3')
+            .run_cmd('echo "test"')
+            .set_start_cmd('echo "Hello"', 'sleep 1')
+        )
+
+        build_info = await AsyncTemplate.build_in_background(
+            template,
+            alias='my-python-env',
+            cpu_count=2,
+            memory_mb=1024
+        )
+        ```
+        """
+        domain = domain or os.environ.get("MORU_DOMAIN", "moru.io")
+        config = ConnectionConfig(
+            domain=domain, api_key=api_key or os.environ.get("MORU_API_KEY")
+        )
+        api_client = get_api_client(
+            config,
+            require_api_key=True,
+            require_access_token=False,
+        )
+
+        return await AsyncTemplate._build(
+            template,
+            api_client,
+            alias,
+            cpu_count,
+            memory_mb,
+            skip_cache,
+            on_build_logs,
+        )
+
+    @staticmethod
+    async def get_build_status(
+        build_info: BuildInfo,
+        logs_offset: int = 0,
+        api_key: Optional[str] = None,
+        domain: Optional[str] = None,
+    ):
+        """
+        Get the status of a build.
+
+        :param build_info: Build identifiers returned from build_in_background
+        :param logs_offset: Offset for fetching logs
+        :param api_key: Moru API key for authentication
+        :param domain: Domain of the Moru API
+        :return: TemplateBuild containing the build status and logs
+
+        Example
+        ```python
+        from moru import AsyncTemplate
+
+        build_info = await AsyncTemplate.build_in_background(template, alias='my-template')
+        status = await AsyncTemplate.get_build_status(build_info, logs_offset=0)
+        ```
+        """
+        domain = domain or os.environ.get("MORU_DOMAIN", "moru.io")
+        config = ConnectionConfig(
+            domain=domain, api_key=api_key or os.environ.get("MORU_API_KEY")
+        )
+        api_client = get_api_client(config)
+        return await get_build_status(
+            api_client,
+            build_info.template_id,
+            build_info.build_id,
+            logs_offset,
+        )