fxn 0.0.42__py3-none-any.whl → 0.0.44__py3-none-any.whl

fxn/client.py

@@ -4,9 +4,9 @@
 #
 
 from json import loads, JSONDecodeError
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 from requests import request
-from typing import Any, Literal, Type, TypeVar
+from typing import AsyncGenerator, Literal, Type, TypeVar
 
 T = TypeVar("T", bound=BaseModel)
 
@@ -19,11 +19,20 @@ class FunctionClient:
     def request (
         self,
         *,
-        method: Literal["GET", "POST", "DELETE"],
+        method: Literal["GET", "POST", "PATCH", "DELETE"],
         path: str,
-        body: dict[str, Any]=None,
+        body: dict[str, object]=None,
         response_type: Type[T]=None
     ) -> T:
+        """
+        Make a request to a REST endpoint.
+
+        Parameters:
+            method (str): Request method.
+            path (str): Endpoint path.
+            body (dict): Request JSON body.
+            response_type (Type): Response type.
+        """
         response = request(
             method=method,
             url=f"{self.api_url}{path}",
@@ -40,6 +49,53 @@ class FunctionClient:
         else:
             error = _ErrorResponse(**data).errors[0].message if isinstance(data, dict) else data
             raise FunctionAPIError(error, response.status_code)
+        
+    async def stream (
+        self,
+        *,
+        method: Literal["GET", "POST", "PATCH", "DELETE"],
+        path: str,
+        body: dict[str, object]=None,
+        response_type: Type[T]=None
+    ) -> AsyncGenerator[T, None]:
+        """
+        Make a request to a REST endpoint and consume the response as a server-sent events stream.
+
+        Parameters:
+            method (str): Request method.
+            path (str): Endpoint path.
+            body (dict): Request JSON body.
+            response_type (Type): Response type.
+        """
+        response = request(
+            method=method,
+            url=f"{self.api_url}{path}",
+            json=body,
+            headers={
+                "Accept": "text/event-stream",
+                "Authorization": f"Bearer {self.access_key}"
+            },
+            stream=True
+        )
+        event = None
+        data: str = ""
+        for line in response.iter_lines(decode_unicode=True):
+            if line is None:
+                break
+            line: str = line.strip()
+            if line:
+                if line.startswith("event:"):
+                    event = line[len("event:"):].strip()
+                elif line.startswith("data:"):
+                    line_data = line[len("data:"):].strip()
+                    data = f"{data}\n{line_data}"
+                continue
+            if event is not None:
+                yield _parse_sse_event(event, data, response_type)
+            event = None
+            data = ""
+        if event or data:
+            yield _parse_sse_event(event, data, response_type)
 
 class FunctionAPIError (Exception):
 
@@ -55,4 +111,9 @@ class _APIError (BaseModel):
     message: str
 
 class _ErrorResponse (BaseModel):
-    errors: list[_APIError]
+    errors: list[_APIError]
+
+def _parse_sse_event (event: str, data: str, type: Type[T]=None) -> T:
+    result = { "event": event, "data": loads(data) }
+    result = TypeAdapter(type).validate_python(result) if type is not None else result
+    return result

fxn/{compile/compile.py → compile.py}

@@ -5,12 +5,16 @@
 
 from collections.abc import Callable
 from functools import wraps
+from inspect import isasyncgenfunction, iscoroutinefunction
 from pathlib import Path
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
+from types import ModuleType
+from typing import Literal
 
-from ..types import AccessMode, Signature
 from .sandbox import Sandbox
-from .signature import get_function_type, infer_function_signature, FunctionType
+from .types import AccessMode
+
+CompileTarget = Literal["android", "ios", "linux", "macos", "visionos", "wasm", "windows"]
 
 class PredictorSpec (BaseModel):
     """
@@ -19,21 +23,26 @@ class PredictorSpec (BaseModel):
     tag: str = Field(description="Predictor tag.")
     description: str = Field(description="Predictor description. MUST be less than 100 characters long.", min_length=4, max_length=100)
     sandbox: Sandbox = Field(description="Sandbox to compile the function.")
+    trace_modules: list[ModuleType] = Field(description="Modules to trace and compile.", exclude=True)
+    targets: list[str] | None = Field(description="Targets to compile this predictor for. Pass `None` to compile for our default targets.")
     access: AccessMode = Field(description="Predictor access.")
-    signature: Signature = Field(description="Predictor signature.")
     card: str | None = Field(default=None, description="Predictor card (markdown).")
     media: str | None = Field(default=None, description="Predictor media URL.")
     license: str | None = Field(default=None, description="Predictor license URL. This is required for public predictors.")
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow", frozen=True)
 
 def compile (
     tag: str,
     *,
     description: str,
     sandbox: Sandbox=None,
+    trace_modules: list[ModuleType]=[],
+    targets: list[CompileTarget]=None,
     access: AccessMode=AccessMode.Private,
     card: str | Path=None,
     media: Path=None,
     license: str=None,
+    **kwargs
 ):
     """
     Create a predictor by compiling a stateless function.
@@ -42,6 +51,8 @@ def compile (
         tag (str): Predictor tag.
         description (str): Predictor description. MUST be less than 100 characters long.
         sandbox (Sandbox): Sandbox to compile the function.
+        trace_modules (list): Modules to trace and compile.
+        targets (list): Targets to compile this predictor for. Pass `None` to compile for our default targets.
         access (AccessMode): Predictor access.
         card (str | Path): Predictor card markdown string or path to card.
         media (Path): Predictor thumbnail image (jpeg or png) path.
@@ -51,25 +62,20 @@ def compile (
         # Check type
         if not callable(func):
             raise TypeError("Cannot compile non-function objects")
-        func_type = get_function_type(func)
-        if func_type not in { FunctionType.Function, FunctionType.Generator }:
-            raise TypeError(f"Function '{func.__name__}' must be a regular function or generator")
+        if isasyncgenfunction(func) or iscoroutinefunction(func):
+            raise TypeError(f"Entrypoint function '{func.__name__}' must be a regular function or generator")            
         # Gather metadata
-        signature = infer_function_signature(func) # throws
-        if isinstance(card, Path):
-            with open(card_content, "r") as f:
-                card_content = f.read()
-        else:
-            card_content = card
         spec = PredictorSpec(
             tag=tag,
             description=description,
             sandbox=sandbox if sandbox is not None else Sandbox(),
+            trace_modules=trace_modules,
+            targets=targets,
             access=access,
-            signature=signature,
-            card=card_content,
+            card=card.read_text() if isinstance(card, Path) else card,
             media=None, # INCOMPLETE
-            license=license
+            license=license,
+            **kwargs
         )
         # Wrap
         @wraps(func)

fxn/function.py

@@ -30,11 +30,16 @@ class Function:
     predictions: PredictionService
     beta: BetaClient
 
-    def __init__ (self, access_key: str=None, api_url: str=None):
-        access_key = access_key or environ.get("FXN_ACCESS_KEY", None)
+    def __init__ (
+        self,
+        access_key: str=None,
+        *,
+        api_url: str=None
+    ):
+        access_key = access_key or environ.get("FXN_ACCESS_KEY")
         api_url = api_url or environ.get("FXN_API_URL")
         self.client = FunctionClient(access_key, api_url)
         self.users = UserService(self.client)
         self.predictors = PredictorService(self.client)
         self.predictions = PredictionService(self.client)
-        self.beta = BetaClient(self.client)
+        self.beta = BetaClient(self.client, predictions=self.predictions)

fxn/logging.py

@@ -0,0 +1,219 @@
+# 
+#   Function
+#   Copyright © 2025 NatML Inc. All Rights Reserved.
+#
+
+from contextvars import ContextVar
+from rich.console import Console, ConsoleOptions, RenderResult
+from rich.progress import Progress, ProgressColumn, SpinnerColumn, TextColumn
+from rich.text import Text
+from rich.traceback import Traceback
+from types import MethodType
+
+current_progress = ContextVar("current_progress", default=None)
+progress_task_stack = ContextVar("progress_task_stack", default=[])
+
+class CustomSpinnerColumn (SpinnerColumn):
+    
+    def __init__ (
+        self,
+        spinner_name="dots",
+        success_text="[bold green]✔[/bold green]",
+        failure_text="[bright_red]✘[/bright_red]",
+        style="",
+    ):
+        super().__init__(spinner_name=spinner_name, style=style)
+        self.success_text = success_text
+        self.failure_text = failure_text
+
+    def render (self, task):
+        done_text = (
+            self.failure_text
+            if task.fields.get("status") == "error"
+            else self.success_text
+        )
+        return done_text if task.finished else self.spinner
+
+class CustomTextColumn (TextColumn):
+    """Custom text column that changes color based on task status"""
+    
+    def __init__ (self, text_format="{task.description}"):
+        super().__init__(text_format)
+    
+    def render (self, task):
+        # Indent and color
+        description = task.description
+        indent_level = task.fields.get("indent_level", 0)
+        indent = self.__get_indent(indent_level)
+        task.description = f"{indent}{description}"
+        if task.fields.get("status") == "error":
+            task.description = f"[bright_red]{task.description}[/bright_red]"
+        # Render
+        text = super().render(task)
+        task.description = description        
+        # Return
+        return text
+    
+    def __get_indent (self, level: int) -> str:
+        if level == 0:
+            return ""
+        indicator = "└── "
+        return " " * len(indicator) * (level - 1) + indicator
+
+class CustomProgress(Progress):
+
+    def __init__ (
+        self,
+        *columns: ProgressColumn,
+        console=None,
+        auto_refresh=True,
+        refresh_per_second = 10,
+        speed_estimate_period=30,
+        transient=False,
+        redirect_stdout=True,
+        redirect_stderr=True,
+        get_time=None,
+        disable=False,
+        expand=False
+    ):
+        default_columns = list(columns) if len(columns) > 0 else [
+            CustomSpinnerColumn(),
+            CustomTextColumn("[progress.description]{task.description}"),
+        ]
+        super().__init__(
+            *default_columns,
+            console=console,
+            auto_refresh=auto_refresh,
+            refresh_per_second=refresh_per_second,
+            speed_estimate_period=speed_estimate_period,
+            transient=transient,
+            redirect_stdout=redirect_stdout,
+            redirect_stderr=redirect_stderr,
+            get_time=get_time,
+            disable=disable,
+            expand=expand
+        )
+        self.default_columns = default_columns
+
+    def __enter__ (self):
+        self._token = current_progress.set(self)
+        self._stack_token = progress_task_stack.set([])
+        return super().__enter__()
+
+    def __exit__ (self, exc_type, exc_val, exc_tb):
+        current_progress.reset(self._token)
+        progress_task_stack.reset(self._stack_token)
+        return super().__exit__(exc_type, exc_val, exc_tb)
+    
+    def get_renderables (self):
+        for task in self.tasks:
+            task_columns = task.fields.get("columns") or list()
+            self.columns = self.default_columns + task_columns
+            yield self.make_tasks_table([task])
+
+class CustomProgressTask:
+
+    def __init__ (
+        self,
+        *,
+        loading_text: str,
+        done_text: str=None,
+        columns: list[ProgressColumn]=None
+    ):
+        self.loading_text = loading_text
+        self.done_text = done_text
+        self.task_id = None
+        self.columns = columns
+
+    def __enter__ (self):
+        progress = current_progress.get()
+        if progress is not None:
+            self.task_id = progress.add_task(
+                self.loading_text,
+                total=1,
+                columns=self.columns,
+                indent_level=len(progress_task_stack.get())
+            )
+            current_stack = progress_task_stack.get()
+            progress_task_stack.set(current_stack + [self.task_id])
+        return self
+
+    def __exit__ (self, exc_type, exc_val, exc_tb):
+        progress = current_progress.get()
+        if progress is not None and self.task_id is not None:
+            current_task = progress._tasks[self.task_id]
+            progress.update(
+                self.task_id,
+                description=self.done_text or current_task.description,
+                completed=current_task.total,
+                status="error" if exc_type is not None else current_task.fields.get("status")
+            )
+            current_stack = progress_task_stack.get()
+            if current_stack:
+                progress_task_stack.set(current_stack[:-1])
+        self.task_id = None
+        return False
+
+    def update (self, **kwargs):
+        progress = current_progress.get()
+        if progress is None or self.task_id is None:
+            return
+        progress.update(self.task_id, **kwargs)
+
+    def finish (self, message: str):
+        self.done_text = message
+    
+class TracebackMarkupConsole (Console):
+
+    def print(
+        self,
+        *objects,
+        sep = " ",
+        end = "\n",
+        style = None,
+        justify = None,
+        overflow = None,
+        no_wrap = None,
+        emoji = None,
+        markup = None,
+        highlight = None,
+        width = None,
+        height = None,
+        crop = True,
+        soft_wrap = None,
+        new_line_start = False
+    ):
+        traceback = objects[0]
+        if isinstance(traceback, Traceback):
+            stack = traceback.trace.stacks[0]
+            original_rich_console = traceback.__rich_console__
+            def __rich_console__ (self: Traceback, console: Console, options: ConsoleOptions) -> RenderResult:
+                for renderable in original_rich_console(console, options):
+                    if (
+                        isinstance(renderable, Text) and
+                        any(part.startswith(f"{stack.exc_type}:") for part in renderable._text)
+                    ):
+                        yield Text.assemble(
+                            (f"{stack.exc_type}: ", "traceback.exc_type"),
+                            Text.from_markup(stack.exc_value)
+                        )
+                    else:
+                        yield renderable
+            traceback.__rich_console__ = MethodType(__rich_console__, traceback)
+        return super().print(
+            *objects,
+            sep=sep,
+            end=end,
+            style=style,
+            justify=justify,
+            overflow=overflow,
+            no_wrap=no_wrap,
+            emoji=emoji,
+            markup=markup,
+            highlight=highlight,
+            width=width,
+            height=height,
+            crop=crop,
+            soft_wrap=soft_wrap,
+            new_line_start=new_line_start
+        )

fxn/{compile/sandbox.py → sandbox.py}

@@ -4,13 +4,16 @@
 #
 
 from __future__ import annotations
+from abc import ABC, abstractmethod
 from hashlib import sha256
 from pathlib import Path
 from pydantic import BaseModel
 from requests import put
+from rich.progress import BarColumn, TextColumn
 from typing import Literal
 
-from ..function import Function
+from .function import Function
+from .logging import CustomProgressTask
 
 class WorkdirCommand (BaseModel):
     kind: Literal["workdir"] = "workdir"
@@ -20,17 +23,35 @@ class EnvCommand (BaseModel):
     kind: Literal["env"] = "env"
     env: dict[str, str]
 
-class UploadFileCommand (BaseModel):
-    kind: Literal["upload_file"] = "upload_file"
+class UploadableCommand (BaseModel, ABC):
     from_path: str
     to_path: str
     manifest: dict[str, str] | None = None
 
-class UploadDirectoryCommand (BaseModel):
+    @abstractmethod
+    def get_files (self) -> list[Path]:
+        pass
+
+class UploadFileCommand (UploadableCommand):
+    kind: Literal["upload_file"] = "upload_file"
+    
+    def get_files (self) -> list[Path]:
+        return [Path(self.from_path).resolve()]
+
+class UploadDirectoryCommand (UploadableCommand):
     kind: Literal["upload_dir"] = "upload_dir"
-    from_path: str
-    to_path: str
-    manifest: dict[str, str] | None = None
+
+    def get_files (self) -> list[Path]:
+        from_path = Path(self.from_path)
+        assert from_path.is_absolute(), "Cannot upload directory because directory path must be absolute"
+        return [file for file in from_path.rglob("*") if file.is_file()]
+    
+class EntrypointCommand (UploadableCommand):
+    kind: Literal["entrypoint"] = "entrypoint"
+    name: str
+
+    def get_files (self) -> list[Path]:
+        return [Path(self.from_path).resolve()]
 
 class PipInstallCommand (BaseModel):
     kind: Literal["pip_install"] = "pip_install"
@@ -40,10 +61,6 @@ class AptInstallCommand (BaseModel):
     kind: Literal["apt_install"] = "apt_install"
     packages: list[str]
 
-class EntrypointCommand (BaseModel):
-    kind: Literal["entrypoint"] = "entrypoint"
-    path: str
-
 Command = (
     WorkdirCommand          |
     EnvCommand              |
@@ -68,16 +85,14 @@ class Sandbox (BaseModel):
             path (str | Path): Path to change to.
         """
         command = WorkdirCommand(path=str(path))
-        self.commands.append(command)
-        return self
+        return Sandbox(commands=self.commands + [command])
 
     def env (self, **env: str) -> Sandbox:
         """
         Set environment variables in the sandbox.
         """
         command = EnvCommand(env=env)
-        self.commands.append(command)
-        return self
+        return Sandbox(commands=self.commands + [command])
 
     def upload_file (
         self,
@@ -92,8 +107,7 @@ class Sandbox (BaseModel):
             to_path (str | Path): Remote path to upload file to.
         """
         command = UploadFileCommand(from_path=str(from_path), to_path=str(to_path))
-        self.commands.append(command)
-        return self
+        return Sandbox(commands=self.commands + [command])
 
     def upload_directory (
         self,
@@ -108,8 +122,7 @@ class Sandbox (BaseModel):
             to_path (str | Path): Remote path to upload directory to.
         """
         command = UploadDirectoryCommand(from_path=str(from_path), to_path=str(to_path))
-        self.commands.append(command)
-        return self
+        return Sandbox(commands=self.commands + [command])
 
     def pip_install (self, *packages: str) -> Sandbox:
         """
@@ -119,8 +132,7 @@ class Sandbox (BaseModel):
             packages (list): Packages to install.
         """
         command = PipInstallCommand(packages=packages)
-        self.commands.append(command)
-        return self
+        return Sandbox(commands=self.commands + [command])
 
     def apt_install (self, *packages: str) -> Sandbox:
         """
@@ -130,24 +142,41 @@ class Sandbox (BaseModel):
             packages (list): Packages to install.
         """
         command = AptInstallCommand(packages=packages)
-        self.commands.append(command)
-        return self
-    
-    def populate (self, fxn: Function=None) -> Sandbox:
+        return Sandbox(commands=self.commands + [command])
+
+    def populate (self, fxn: Function=None) -> Sandbox: # CHECK # In place
         """
         Populate all metadata.
         """
         fxn = fxn if fxn is not None else Function()
+        entrypoint = next(cmd for cmd in self.commands if isinstance(cmd, EntrypointCommand))
+        entry_path = Path(entrypoint.from_path).resolve()
         for command in self.commands:
-            if isinstance(command, UploadFileCommand):
-                from_path = Path(command.from_path)
-                to_path = Path(command.to_path)
-                command.manifest = { str(to_path / from_path.name): self.__upload_file(from_path, fxn=fxn) }
-            elif isinstance(command, UploadDirectoryCommand):
+            if isinstance(command, UploadableCommand):
+                cwd = Path.cwd()
                 from_path = Path(command.from_path)
                 to_path = Path(command.to_path)
-                files = [file for file in from_path.rglob("*") if file.is_file()]
-                command.manifest = { str(to_path / file.relative_to(from_path)): self.__upload_file(file, fxn=fxn) for file in files }
+                if not from_path.is_absolute():
+                    from_path = (entry_path / from_path).resolve()
+                    command.from_path = str(from_path)
+                files = command.get_files()
+                name = from_path.relative_to(cwd) if from_path.is_relative_to(cwd) else from_path.resolve()
+                with CustomProgressTask(
+                    loading_text=f"Uploading [light_slate_blue]{name}[/light_slate_blue]...",
+                    done_text=f"Uploaded [light_slate_blue]{name}[/light_slate_blue]",
+                    columns=[
+                        BarColumn(),
+                        TextColumn("{task.completed}/{task.total}")
+                    ]
+                ) as task:
+                    manifest = { }
+                    for idx, file in enumerate(files):
+                        rel_file_path = file.relative_to(from_path) if from_path.is_dir() else file.name
+                        dst_path = to_path / rel_file_path
+                        checksum = self.__upload_file(file, fxn=fxn)
+                        manifest[str(dst_path)] = checksum
+                        task.update(total=len(files), completed=idx+1)
+                    command.manifest = manifest
         return self
 
     def __upload_file (self, path: Path, fxn: Function) -> str:
@@ -172,6 +201,6 @@ class Sandbox (BaseModel):
             for chunk in iter(lambda: f.read(4096), b""):
                 hash.update(chunk)
         return hash.hexdigest()
-    
+
 class _Resource (BaseModel):
     url: str