flyteplugins-omlx 0.0.0a0__py3-none-any.whl

flyteplugins/omlx/__init__.py

@@ -0,0 +1,3 @@
+__all__ = ["DEFAULT_OMLX_IMAGE", "OMLXAppEnvironment"]
+
+from flyteplugins.omlx._app_environment import DEFAULT_OMLX_IMAGE, OMLXAppEnvironment

flyteplugins/omlx/_app_environment.py

@@ -0,0 +1,228 @@
+from __future__ import annotations
+
+import shlex
+from dataclasses import dataclass, field, replace
+from typing import Any, Literal, Optional, Union
+
+import flyte.app
+import rich.repr
+from flyte import Environment, Image, Resources, SecretRequest
+from flyte.app._types import Port
+from flyte.models import SerializationContext
+
+from flyteplugins.omlx._constants import DEFAULT_OMLX_PORT
+
+
+def _shell_safe(args: list[str]) -> list[str]:
+    """Quote args that have to survive a trip through a shell.
+
+    `fserve` runs the app with `Popen(" ".join(args), shell=True)` (flyte/_bin/serve.py),
+    so any token carrying spaces or quotes -- a `--model-dir` under a user name with a
+    space, say -- reaches the server mangled unless it is quoted here. `shlex.quote` is
+    the identity function for ordinary tokens, so this is a no-op for everything else.
+
+    Tokens starting with `$` are left alone: `fserve` expands those against the
+    environment *before* joining, and quoting would turn the marker into a literal.
+    """
+    return [arg if arg.startswith("$") else shlex.quote(arg) for arg in args]
+
+
+# oMLX runs only on macOS with Apple Silicon: `mlx` has no Linux wheels, so this default
+# image only installs the plugin. The `omlx` binary itself is expected on the Mac's host
+# (Homebrew or pip-from-git), where the devbox's Apple Silicon bridge runs the app; the
+# `flyte-omlx` wrapper resolves it (probing Homebrew paths that host-execution's PATH
+# omits). See the plugin README.
+DEFAULT_OMLX_IMAGE = flyte.Image.from_debian_base(name="omlx-app-image").with_pip_packages(
+    "flyteplugins-omlx",
+    pre=True,
+)
+
+
+@rich.repr.auto
+@dataclass(kw_only=True, repr=True)
+class OMLXAppEnvironment(flyte.app.AppEnvironment):
+    """
+    App environment backed by [oMLX](https://github.com/jundot/omlx) for serving
+    LLMs on Apple Silicon Macs.
+
+    oMLX is an OpenAI-compatible inference server built on Apple's MLX framework.
+    It is **macOS + Apple Silicon only**. On a flyte devbox started with
+    `--apple-silicon`, this environment requests `gpu="AppleSilicon"` by default, so
+    the devbox's bridge runs `omlx serve` natively on the Mac (in an environment built
+    from the image's pip packages) and forwards the app's URL, port and logs.
+
+    ```bash
+    flyte start devbox --apple-silicon
+    flyte create config --devbox
+    python examples/accelerators/omlx_serve.py
+    ```
+
+    The user is responsible for installing `omlx` on the host machine:
+
+    ```bash
+    brew tap jundot/omlx https://github.com/jundot/omlx && brew install omlx
+    # or: pip install git+https://github.com/jundot/omlx
+    ```
+
+    Args:
+        name: The name of the application.
+        image: Container image. oMLX cannot run inside Linux containers, so this
+            mostly describes the pip *packages* that get rebuilt into the Mac's host
+            environment. Leave as default for the devbox workflow; you typically only
+            override it to add pip packages (e.g. `with_pip_packages("omlx @ git+...")`).
+        port: Port the oMLX server listens on. Defaults to 8000 (oMLX's own default).
+        resources: Compute resources. Defaults to `gpu="AppleSilicon"` so the app runs
+            on the Mac's GPU via the devbox bridge; other GPU types are rejected.
+        secrets: Secrets requested by the app.
+        env_vars: Environment variables for the server process.
+        scaling: Scaling configuration.
+        domain: Custom domain for the deployed app.
+        cluster_pool: Target cluster pool.
+        requires_auth: Whether the public URL requires Flyte auth.
+        type: App type tag.
+        model_dir: Local directory of MLX models, passed to `omlx serve --model-dir`.
+            Each subdirectory is exposed as a separate `model_id`. If unset, oMLX falls
+            back to `~/.omlx/models`.
+        model_id: Optional friendly id surfaced in the Flyte UI and usable as the
+            `model` in OpenAI client requests. oMLX derives actual served ids from
+            `--model-dir` subdirectories or HuggingFace repo ids.
+        api_key: Optional API key for the oMLX server (sent through `--api-key`). For
+            real secrets, prefer `secrets` + `env_vars`.
+        extra_args: Extra CLI args appended to `omlx serve`. See `omlx serve --help`
+            for the full list (`--max-model-memory`, `--max-concurrent-requests`,
+            `--paged-ssd-cache-dir`, `--hf-endpoint`, `--log-level`, etc.).
+    """
+
+    port: int | Port = DEFAULT_OMLX_PORT
+    type: str = "oMLX"
+    extra_args: str | list[str] = ""
+    model_dir: str = ""
+    model_id: str = ""
+    api_key: str = ""
+    image: str | Image | Literal["auto"] = DEFAULT_OMLX_IMAGE
+    _resolved_args: list[str] = field(default_factory=list, init=False, repr=False)
+
+    def __post_init__(self):
+        if self.env_vars is None:
+            self.env_vars = {}
+
+        if self._server is not None:
+            raise ValueError("server function cannot be set for OMLXAppEnvironment")
+
+        if self._on_startup is not None:
+            raise ValueError("on_startup function cannot be set for OMLXAppEnvironment")
+
+        if self._on_shutdown is not None:
+            raise ValueError("on_shutdown function cannot be set for OMLXAppEnvironment")
+
+        if self.args:
+            raise ValueError("args cannot be set for OMLXAppEnvironment. Use `extra_args` to add extra arguments.")
+
+        # oMLX is mac-only: default to running natively on the Mac's GPU through the
+        # devbox's Apple Silicon bridge, and refuse any other accelerator.
+        device = self.resources.get_device() if self.resources is not None else None
+        if device is None:
+            if self.resources is None:
+                self.resources = Resources(gpu="AppleSilicon")
+            else:
+                self.resources = replace(self.resources, gpu="AppleSilicon")
+        elif device.device_class != "APPLE_SILICON":
+            raise ValueError(
+                "OMLXAppEnvironment only runs on Apple Silicon (gpu='AppleSilicon') because oMLX uses MLX, "
+                "which is darwin/arm64 only; got resources with a non-Apple Silicon device."
+            )
+
+        if isinstance(self.extra_args, str):
+            extra_args = shlex.split(self.extra_args)
+        else:
+            extra_args = list(self.extra_args)
+
+        # oMLX is served locally on the Mac (devbox ports are bound to 127.0.0.1), so bind
+        # loopback by default; a --host in extra_args wins.
+        host_args = [] if any(a == "--host" for a in extra_args) else ["--host", "127.0.0.1"]
+
+        # We invoke the `flyte-omlx` wrapper (defined in this package) instead of `omlx`
+        # directly so the user gets an actionable install hint if oMLX is missing on the
+        # host. When oMLX is installed, the wrapper `execv`s straight into it.
+        cmd: list[str] = [
+            "flyte-omlx",
+            "serve",
+            *host_args,
+            "--port",
+            str(self.get_port().port),
+        ]
+        if self.model_dir:
+            cmd += ["--model-dir", self.model_dir]
+        if self.api_key:
+            cmd += ["--api-key", self.api_key]
+        cmd += extra_args
+
+        self.args = _shell_safe(cmd)
+        self._resolved_args = cmd
+
+        if self.image is None or self.image == "auto":
+            self.image = DEFAULT_OMLX_IMAGE
+
+        # Helpful default link to the OpenAI-compatible models endpoint so users can
+        # sanity-check the server from the Flyte UI.
+        self.links = [flyte.app.Link(path="/v1/models", title="oMLX Models", is_relative=True), *self.links]
+
+        super().__post_init__()
+
+    def container_args(self, serialization_context: SerializationContext) -> list[str]:
+        """Return the container arguments for oMLX."""
+        if isinstance(self.args, str):
+            return shlex.split(self.args)
+        return self.args or []
+
+    def clone_with(
+        self,
+        name: str,
+        image: Optional[Union[str, Image, Literal["auto"]]] = None,
+        resources: Optional[Resources] = None,
+        env_vars: Optional[dict[str, str]] = None,
+        secrets: Optional[SecretRequest] = None,
+        depends_on: Optional[list[Environment]] = None,
+        description: Optional[str] = None,
+        interruptible: Optional[bool] = None,
+        **kwargs: Any,
+    ) -> OMLXAppEnvironment:
+        port = kwargs.pop("port", None)
+        extra_args = kwargs.pop("extra_args", None)
+        model_dir = kwargs.pop("model_dir", None)
+        model_id = kwargs.pop("model_id", None)
+        api_key = kwargs.pop("api_key", None)
+
+        if kwargs:
+            raise TypeError(f"Unexpected keyword arguments: {list(kwargs.keys())}")
+
+        kwargs = self._get_kwargs()
+        kwargs["name"] = name
+        # Reset args so __post_init__ rebuilds them with the new fields.
+        kwargs["args"] = None
+        kwargs["parameters"] = None
+        if image is not None:
+            kwargs["image"] = image
+        if resources is not None:
+            kwargs["resources"] = resources
+        if env_vars is not None:
+            kwargs["env_vars"] = env_vars
+        if secrets is not None:
+            kwargs["secrets"] = secrets
+        if depends_on is not None:
+            kwargs["depends_on"] = depends_on
+        if description is not None:
+            kwargs["description"] = description
+        if interruptible is not None:
+            kwargs["interruptible"] = interruptible
+        if port is not None:
+            kwargs["port"] = port
+        if extra_args is not None:
+            kwargs["extra_args"] = extra_args
+        if model_dir is not None:
+            kwargs["model_dir"] = model_dir
+        if model_id is not None:
+            kwargs["model_id"] = model_id
+        if api_key is not None:
+            kwargs["api_key"] = api_key
+        return replace(self, **kwargs)

flyteplugins/omlx/_cli.py

@@ -0,0 +1,79 @@
+"""Friendly preflight wrapper around the upstream `omlx` console command.
+
+An `OMLXAppEnvironment` serves via `omlx serve`. The local-serve / Apple Silicon
+host process invokes `flyte-omlx` instead of `omlx` directly so that, if oMLX isn't
+installed on the Mac, the user sees a clear install hint instead of a raw
+`FileNotFoundError: 'omlx'`.
+
+When oMLX *is* installed, this wrapper is transparent: it `execv`'s into the real
+binary (or imports `omlx.cli` if the package is installed but the console script
+isn't on PATH for some reason).
+
+`gpu="AppleSilicon"` workloads on the flyte devbox run their app natively on the Mac
+in a `uv` environment whose `PATH` only contains that environment's bin directory. A
+Homebrew-installed oMLX lives at `/opt/homebrew/bin` (Apple Silicon) or
+`/usr/local/bin` (Intel Macs and older setups), neither of which is on that `PATH`, so
+`shutil.which("omlx")` alone would miss it. We probe those locations explicitly.
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+import sys
+
+_INSTALL_HINT = """\
+[flyteplugins-omlx] oMLX is not installed on this machine.
+
+oMLX is macOS + Apple Silicon only and is not published to PyPI. Install it
+with one of:
+
+    brew tap jundot/omlx https://github.com/jundot/omlx && brew install omlx
+    pip install git+https://github.com/jundot/omlx
+
+Or download the .dmg from https://github.com/jundot/omlx/releases.
+
+Then re-run `flyte serve` / redeploy the app. See the plugin README for details.
+"""
+
+# Common Homebrew prefixes for the `omlx` binary. Checked in `resolve_omlx` after
+# `shutil.which`, because the Apple Silicon host-execution environment's PATH
+# omits Homebrew.
+_HOMEBREW_OMLX_CANDIDATES = ("/opt/homebrew/bin/omlx", "/usr/local/bin/omlx")
+
+
+def resolve_omlx() -> str | None:
+    """The path to the `omlx` binary, or None if it cannot be found anywhere."""
+    found = shutil.which("omlx")
+    if found:
+        return found
+    for candidate in _HOMEBREW_OMLX_CANDIDATES:
+        if os.path.isfile(candidate) and os.access(candidate, os.X_OK):
+            return candidate
+    return None
+
+
+def main() -> int:
+    binary = resolve_omlx()
+    if binary:
+        # Hand off to the real binary, preserving argv. execv replaces this
+        # process, so the real omlx becomes the long-running server and its
+        # signals/log streams belong to it.
+        os.execv(binary, [binary, *sys.argv[1:]])
+        return 0  # not reached
+
+    # Fall back to the installed Python package, in case the script wasn't
+    # placed on PATH (e.g. an unusual install layout).
+    try:
+        import omlx.cli as omlx_cli  # type: ignore[import-not-found]
+    except ImportError:
+        sys.stderr.write(_INSTALL_HINT)
+        return 127
+
+    sys.argv = ["omlx", *sys.argv[1:]]
+    rc = omlx_cli.main()
+    return int(rc) if rc is not None else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

flyteplugins/omlx/_constants.py

@@ -0,0 +1,9 @@
+"""Constants shared by the oMLX plugin."""
+
+OMLX_MIN_VERSION = (0, 1, 0)
+OMLX_MIN_VERSION_STR = ".".join(map(str, OMLX_MIN_VERSION))
+
+# oMLX's own default serve port. Kept distinct from the SDK's 8080 so that
+# "just construct an OMLXAppEnvironment" matches what `omlx serve` would bind
+# if you ran it by hand.
+DEFAULT_OMLX_PORT = 8000

flyteplugins_omlx-0.0.0a0.dist-info/METADATA

@@ -0,0 +1,151 @@
+Metadata-Version: 2.4
+Name: flyteplugins-omlx
+Version: 0.0.0a0
+Summary: oMLX plugin for flyte — serve LLMs on Apple Silicon
+Author: Flyte authors
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: flyte>=2.0.0b43
+
+# Flyte oMLX Plugin
+
+Serve LLMs locally on Apple Silicon using [oMLX](https://github.com/jundot/omlx)
+and Flyte Apps.
+
+This plugin provides the `OMLXAppEnvironment` class for defining an oMLX-backed
+app environment. oMLX is an OpenAI-compatible inference server built on Apple's
+MLX framework, macOS + Apple Silicon only.
+
+The plugin integrates with the **Apple Silicon flyte devbox**: an
+`OMLXAppEnvironment` requests `gpu="AppleSilicon"` by default, so the devbox's
+bridge runs `omlx serve` natively on the Mac, in an environment built from the
+image's pip packages, and forwards the app's URL, port and logs.
+
+> **Mac-only.** oMLX runs only on macOS 15+ with Apple Silicon (M1/M2/M3/M4).
+> `mlx` has no Linux wheels, so this plugin cannot run in a Linux container or on
+> a cloud cluster. It is designed for a devbox started with `--apple-silicon`.
+
+## Installation
+
+oMLX is **not on PyPI** — install it via Homebrew, the macOS `.dmg`, or from
+source. See the upstream
+[install guide](https://github.com/jundot/omlx#installation):
+
+```bash
+# oMLX itself (one of):
+brew tap jundot/omlx https://github.com/jundot/omlx && brew install omlx
+# or:
+pip install git+https://github.com/jundot/omlx
+
+# Then the Flyte plugin:
+pip install --pre flyteplugins-omlx
+```
+
+Verify the install with `omlx --help` before running the plugin.
+
+## Quickstart — serve an LLM from the Apple Silicon devbox
+
+`examples/accelerators/omlx_serve.py`:
+
+```python
+import flyte
+from flyteplugins.omlx import OMLXAppEnvironment
+
+omlx_app = OMLXAppEnvironment(
+    name="omlx-local",
+    # oMLX serves whatever is under ~/.omlx/models by default. Set model_dir to a
+    # different folder (an absolute, expanded path).
+    model_id="qwen3-0.6b",
+    extra_args="--log-level info",
+)
+```
+
+Run it against a devbox started with the Apple Silicon bridge:
+
+```bash
+flyte start devbox --apple-silicon
+flyte create config --devbox
+python examples/accelerators/omlx_serve.py
+```
+
+`OMLXAppEnvironment` defaults to `gpu="AppleSilicon"`, so the app runs natively
+on the Mac's GPU. From another shell, hit the OpenAI-compatible API (the app's
+URL is printed by the example):
+
+```python
+from openai import OpenAI
+
+client = OpenAI(base_url="http://omlx-local-...:30081/v1", api_key="not-used")
+resp = client.chat.completions.create(
+    model="qwen3-0.6b",
+    messages=[{"role": "user", "content": "Hello!"}],
+)
+print(resp.choices[0].message.content)
+```
+
+Useful endpoints (also surfaced as Flyte app links):
+
+| Endpoint           | Description                                  |
+|--------------------|----------------------------------------------|
+| `/v1/chat/completions` | OpenAI chat completions                  |
+| `/v1/completions`      | OpenAI completions                       |
+| `/v1/embeddings`       | OpenAI embeddings                        |
+| `/v1/models`           | List available models                    |
+
+## Configuration
+
+`OMLXAppEnvironment` constructs the `omlx serve --host 127.0.0.1 --port <port> ...`
+invocation that the devbox runs on the Mac. Recognised fields:
+
+| Field         | Description                                                                                       |
+|---------------|---------------------------------------------------------------------------------------------------|
+| `name`        | App name (required, lowercase + hyphens).                                                         |
+| `port`        | Bind port. Default `8000` (oMLX default).                                                         |
+| `model_dir`   | Local directory containing MLX-format model subdirectories, mapped to `--model-dir`. Use an expanded absolute path. |
+| `model_id`    | Friendly id surfaced in Flyte / OpenAI client requests; oMLX serves each subdir of `model_dir`.   |
+| `resources`   | `gpu="AppleSilicon"` by default (runs on the Mac via the devbox bridge). Other GPUs are rejected. |
+| `api_key`     | Optional `--api-key` for the oMLX server. For real secrets, prefer `secrets` + `env_vars`.        |
+| `extra_args`  | List or string of extra flags passed to `omlx serve`. See `omlx serve --help`.                    |
+| `env_vars`    | Environment variables for the server process (e.g. `HF_TOKEN`).                                   |
+| `image`       | Container image. Sets the pip packages rebuilt into the Mac's host environment.                   |
+| `scaling`     | `flyte.app.Scaling` config.                                                                       |
+
+### Example: tighten memory / cache config
+
+```python
+omlx_app = OMLXAppEnvironment(
+    name="omlx-local",
+    model_dir="/Users/me/.omlx/models",
+    extra_args=(
+        "--max-model-memory 24GB "
+        "--max-concurrent-requests 4 "
+        "--paged-ssd-cache-dir /Users/me/.omlx/ssd_cache "
+        "--paged-ssd-cache-max-size 50GB "
+        "--log-level info"
+    ),
+)
+```
+
+## How it works
+
+With `flyte start devbox --apple-silicon`, a `gpu="AppleSilicon"` app runs
+**natively on the Mac**. The devbox container (a Linux VM) has no Metal, so the
+bridge hands the app's entrypoint (`fserve`) to the Mac, which runs it in a `uv`
+environment built from the image's pip packages. `fserve` then runs the app's
+`args` — here `omlx serve ...` — on the Mac, and the devbox forwards the app's
+URL, port and logs.
+
+The `args` invoke a tiny `flyte-omlx` console-script wrapper (installed by this
+package) which `execv`s into the real `omlx` binary. Because the Apple Silicon
+host environment's PATH only contains its own `bin/` (Homebrew is not on it), the
+wrapper also probes `/opt/homebrew/bin` and `/usr/local/bin` so a
+Homebrew-installed oMLX is found. If oMLX isn't installed, the wrapper prints an
+actionable install hint instead of a raw `FileNotFoundError`.
+
+## Limitations
+
+- **Apple Silicon only.** oMLX uses `mlx`, which only ships for darwin/arm64. You
+  cannot run this plugin in a Linux container today.
+- oMLX must be installed on the Mac (Homebrew or pip-from-git). The default image
+  installs only the plugin, because `mlx` has no Linux wheels with which to build a
+  base image.

flyteplugins_omlx-0.0.0a0.dist-info/RECORD

@@ -0,0 +1,10 @@
+flyteplugins/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+flyteplugins/omlx/__init__.py,sha256=OVNvya0WcAvUOBczbjKmtRpG4nJhAqHgBOGFDfVOWPM,142
+flyteplugins/omlx/_app_environment.py,sha256=B_fK-bsfkHHk4lVD9LxzoNll1jXt23zs6CU4fdESMdk,9843
+flyteplugins/omlx/_cli.py,sha256=Xp1u1qn6EzJZ0q4rlJnYxULLDQD3nINaDk9VMV6Cplg,2881
+flyteplugins/omlx/_constants.py,sha256=ultyT7SMIPHoKC6GweyGRPui2IaknZ8UEgEg3wezvD4,337
+flyteplugins_omlx-0.0.0a0.dist-info/METADATA,sha256=9R_buwtJwiFMHmmkODGc4W6dB4Ehm2-MpAI6w7B5C0I,6335
+flyteplugins_omlx-0.0.0a0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+flyteplugins_omlx-0.0.0a0.dist-info/entry_points.txt,sha256=-oW89sQ9JlQ44d1g-GSG5sNyI7V2yPvzacH_sYXekss,59
+flyteplugins_omlx-0.0.0a0.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
+flyteplugins_omlx-0.0.0a0.dist-info/RECORD,,

flyteplugins_omlx-0.0.0a0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

flyteplugins_omlx-0.0.0a0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+flyte-omlx = flyteplugins.omlx._cli:main

flyteplugins_omlx-0.0.0a0.dist-info/top_level.txt

@@ -0,0 +1 @@
+flyteplugins