mirror of
https://github.com/exo-explore/exo.git
synced 2026-09-09 20:10:19 -04:00
Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Sami Khan
|
5b81d55d66 | ||
|
Sami Khan
|
3e9eb93f82 | ||
|
Sami Khan
|
ab622f79c3 |
No files matched your search
@@ -1,8 +1 @@
|
||||
use flake
|
||||
|
||||
# creates .venv if doesn't exist and loads its environment
|
||||
export VIRTUAL_ENV=".venv"
|
||||
if ! [ -d "./$VIRTUAL_ENV" ]; then
|
||||
uv venv
|
||||
fi
|
||||
layout python
|
||||
@@ -32,6 +32,7 @@ jobs:
|
||||
SPARKLE_ED25519_PRIVATE: ${{ secrets.SPARKLE_ED25519_PRIVATE }}
|
||||
SPARKLE_S3_BUCKET: ${{ secrets.SPARKLE_S3_BUCKET }}
|
||||
SPARKLE_S3_PREFIX: ${{ secrets.SPARKLE_S3_PREFIX }}
|
||||
EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT: ${{ secrets.EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT }}
|
||||
AWS_REGION: ${{ secrets.AWS_REGION }}
|
||||
EXO_BUILD_NUMBER: ${{ github.run_number }}
|
||||
EXO_LIBP2P_NAMESPACE: ${{ github.ref_name }}
|
||||
@@ -158,7 +159,7 @@ jobs:
|
||||
fi
|
||||
|
||||
- name: Install Homebrew packages
|
||||
run: brew install just awscli
|
||||
run: brew install just awscli macmon
|
||||
|
||||
- name: Install UV
|
||||
uses: astral-sh/setup-uv@v6
|
||||
@@ -238,92 +239,10 @@ jobs:
|
||||
# Export keychain path for other steps
|
||||
echo "BUILD_KEYCHAIN_PATH=$KEYCHAIN_PATH" >> $GITHUB_ENV
|
||||
|
||||
# ============================================================
|
||||
# Pre-flight credential / profile validation
|
||||
# Runs BEFORE the ~16 min build so auth/expiry failures surface in <1 min.
|
||||
# ============================================================
|
||||
|
||||
- name: Validate Apple notarization credentials
|
||||
env:
|
||||
APPLE_NOTARIZATION_USERNAME: ${{ secrets.APPLE_NOTARIZATION_USERNAME }}
|
||||
APPLE_NOTARIZATION_PASSWORD: ${{ secrets.APPLE_NOTARIZATION_PASSWORD }}
|
||||
APPLE_NOTARIZATION_TEAM: ${{ secrets.APPLE_NOTARIZATION_TEAM }}
|
||||
run: |
|
||||
# All-or-nothing: either all three creds are set, or none are.
|
||||
CRED_COUNT=0
|
||||
for v in "$APPLE_NOTARIZATION_USERNAME" "$APPLE_NOTARIZATION_PASSWORD" "$APPLE_NOTARIZATION_TEAM"; do
|
||||
[[ -n "$v" ]] && CRED_COUNT=$((CRED_COUNT + 1))
|
||||
done
|
||||
if [[ "$CRED_COUNT" -eq 0 ]]; then
|
||||
echo "No notarization credentials configured — skipping notarization for this build."
|
||||
exit 0
|
||||
fi
|
||||
if [[ "$CRED_COUNT" -ne 3 ]]; then
|
||||
echo "ERROR: partial notarization credentials set ($CRED_COUNT/3). Aborting before build."
|
||||
exit 1
|
||||
fi
|
||||
# Cheap, ~5s, auth-only call. Fails instantly with a clear message if
|
||||
# the app-specific password is stale, wrong team-id, etc.
|
||||
echo "Verifying Apple notarization credentials via notarytool history..."
|
||||
if ! xcrun notarytool history \
|
||||
--apple-id "$APPLE_NOTARIZATION_USERNAME" \
|
||||
--password "$APPLE_NOTARIZATION_PASSWORD" \
|
||||
--team-id "$APPLE_NOTARIZATION_TEAM" >/dev/null; then
|
||||
echo "ERROR: notarytool rejected the provided credentials. Fix before rerunning."
|
||||
echo "Common causes: app-specific password expired/revoked, wrong team-id,"
|
||||
echo "Apple ID not on the team, or 2FA not configured for this Apple ID."
|
||||
exit 1
|
||||
fi
|
||||
echo "Apple notarization credentials OK."
|
||||
|
||||
- name: Validate provisioning profile expiry
|
||||
run: |
|
||||
PROFILE="$HOME/Library/Developer/Xcode/UserData/Provisioning Profiles/EXO.provisionprofile"
|
||||
if [[ ! -f "$PROFILE" ]]; then
|
||||
echo "ERROR: provisioning profile not found at $PROFILE"
|
||||
exit 1
|
||||
fi
|
||||
EXPIRY=$(security cms -D -i "$PROFILE" | plutil -extract ExpirationDate raw -o - - 2>/dev/null || true)
|
||||
if [[ -z "$EXPIRY" ]]; then
|
||||
echo "WARNING: could not read ExpirationDate from provisioning profile; skipping expiry check."
|
||||
exit 0
|
||||
fi
|
||||
# Try a couple of known plutil date formats. If none parse, skip the check rather
|
||||
# than risk a false-positive "expired" block on a format we didn't anticipate.
|
||||
EXPIRY_EPOCH=""
|
||||
for fmt in "%Y-%m-%dT%H:%M:%SZ" "%Y-%m-%d %H:%M:%S %z" "%Y-%m-%d %H:%M:%S +0000"; do
|
||||
if parsed=$(date -j -f "$fmt" "$EXPIRY" +%s 2>/dev/null); then
|
||||
EXPIRY_EPOCH="$parsed"
|
||||
break
|
||||
fi
|
||||
done
|
||||
if [[ -z "$EXPIRY_EPOCH" ]]; then
|
||||
echo "WARNING: could not parse ExpirationDate '$EXPIRY'; skipping expiry check."
|
||||
exit 0
|
||||
fi
|
||||
NOW_EPOCH=$(date +%s)
|
||||
if [[ "$EXPIRY_EPOCH" -le "$NOW_EPOCH" ]]; then
|
||||
echo "ERROR: provisioning profile expired on $EXPIRY. Regenerate it before rerunning."
|
||||
exit 1
|
||||
fi
|
||||
DAYS_LEFT=$(( (EXPIRY_EPOCH - NOW_EPOCH) / 86400 ))
|
||||
echo "Provisioning profile valid until $EXPIRY ($DAYS_LEFT days remaining)."
|
||||
if [[ "$DAYS_LEFT" -lt 14 ]]; then
|
||||
echo "WARNING: profile expires in under 14 days — regenerate soon."
|
||||
fi
|
||||
|
||||
# ============================================================
|
||||
# Build the bundle
|
||||
# ============================================================
|
||||
|
||||
- name: Add pinned macmon to PATH
|
||||
run: |
|
||||
MACMON_DIR=$(nix develop --command sh -c 'dirname $(which macmon)')
|
||||
echo "Using macmon from: $MACMON_DIR"
|
||||
echo "$MACMON_DIR" >> $GITHUB_PATH
|
||||
# Remove any Homebrew macmon so PyInstaller can't accidentally pick it up
|
||||
brew uninstall macmon 2>/dev/null || true
|
||||
|
||||
- name: Build PyInstaller bundle
|
||||
run: uv run pyinstaller packaging/pyinstaller/exo.spec
|
||||
|
||||
@@ -346,6 +265,7 @@ jobs:
|
||||
EXO_BUILD_COMMIT="$GITHUB_SHA" \
|
||||
SPARKLE_FEED_URL="$SPARKLE_FEED_URL" \
|
||||
SPARKLE_ED25519_PUBLIC="$SPARKLE_ED25519_PUBLIC" \
|
||||
EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT="$EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT" \
|
||||
CODE_SIGNING_IDENTITY="$SIGNING_IDENTITY" \
|
||||
CODE_SIGN_INJECT_BASE_ENTITLEMENTS=YES
|
||||
mkdir -p ../../output
|
||||
@@ -378,41 +298,11 @@ jobs:
|
||||
APPLE_NOTARIZATION_PASSWORD: ${{ secrets.APPLE_NOTARIZATION_PASSWORD }}
|
||||
APPLE_NOTARIZATION_TEAM: ${{ secrets.APPLE_NOTARIZATION_TEAM }}
|
||||
run: |
|
||||
set -o pipefail
|
||||
cd output
|
||||
security unlock-keychain -p "$MACOS_CERTIFICATE_PASSWORD" "$BUILD_KEYCHAIN_PATH"
|
||||
SIGNING_IDENTITY=$(security find-identity -v -p codesigning "$BUILD_KEYCHAIN_PATH" | awk -F '"' '{print $2}')
|
||||
|
||||
# Fail fast if notarization creds are partial. All-or-nothing.
|
||||
CRED_COUNT=0
|
||||
for v in "$APPLE_NOTARIZATION_USERNAME" "$APPLE_NOTARIZATION_PASSWORD" "$APPLE_NOTARIZATION_TEAM"; do
|
||||
[[ -n "$v" ]] && CRED_COUNT=$((CRED_COUNT + 1))
|
||||
done
|
||||
if [[ "$CRED_COUNT" -ne 0 && "$CRED_COUNT" -ne 3 ]]; then
|
||||
echo "ERROR: partial Apple notarization credentials set ($CRED_COUNT/3). Aborting."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
/usr/bin/codesign --deep --force --timestamp --options runtime \
|
||||
--sign "$SIGNING_IDENTITY" EXO.app
|
||||
|
||||
# Pre-flight: verify the signed app BEFORE building DMG and submitting to Apple.
|
||||
# If this fails, notarization will fail too — cheap way to fail in seconds, not 15 minutes.
|
||||
echo "===== codesign --verify EXO.app ====="
|
||||
if ! /usr/bin/codesign --verify --deep --strict --verbose=2 EXO.app; then
|
||||
echo "ERROR: EXO.app failed codesign verification. Dumping signing status of every executable:"
|
||||
find EXO.app -type f \( -perm -111 -o -name "*.dylib" -o -name "*.so" -o -name "*.framework" \) -print0 |
|
||||
while IFS= read -r -d '' f; do
|
||||
printf -- '--- %s\n' "$f"
|
||||
/usr/bin/codesign -dv --verbose=2 "$f" 2>&1 | sed 's/^/ /' || true
|
||||
done
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Gatekeeper assessment. A failure here strongly predicts notarization rejection.
|
||||
echo "===== spctl assessment (predicts notarization outcome) ====="
|
||||
/usr/bin/spctl -a -vvv -t install EXO.app || echo "WARNING: spctl assessment failed — notarization is likely to fail too."
|
||||
|
||||
mkdir -p dmg-root
|
||||
cp -R EXO.app dmg-root/
|
||||
ln -s /Applications dmg-root/Applications
|
||||
@@ -420,22 +310,12 @@ jobs:
|
||||
hdiutil create -volname "EXO" -srcfolder dmg-root -ov -format UDZO "$DMG_NAME"
|
||||
/usr/bin/codesign --force --timestamp --options runtime \
|
||||
--sign "$SIGNING_IDENTITY" "$DMG_NAME"
|
||||
|
||||
echo "===== codesign --verify DMG ====="
|
||||
if ! /usr/bin/codesign --verify --verbose=2 "$DMG_NAME"; then
|
||||
echo "ERROR: DMG failed codesign verification."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [[ -n "$APPLE_NOTARIZATION_USERNAME" ]]; then
|
||||
echo "===== notarytool submit ====="
|
||||
# `|| true` so set -e doesn't abort before we can echo output / fetch the log.
|
||||
# We rely on the parsed STATUS below to decide pass/fail.
|
||||
SUBMISSION_OUTPUT=$(xcrun notarytool submit "$DMG_NAME" \
|
||||
--apple-id "$APPLE_NOTARIZATION_USERNAME" \
|
||||
--password "$APPLE_NOTARIZATION_PASSWORD" \
|
||||
--team-id "$APPLE_NOTARIZATION_TEAM" \
|
||||
--wait --timeout 15m 2>&1) || true
|
||||
--wait --timeout 15m 2>&1)
|
||||
echo "$SUBMISSION_OUTPUT"
|
||||
|
||||
SUBMISSION_ID=$(echo "$SUBMISSION_OUTPUT" | awk 'tolower($1)=="id:" && $2 ~ /^[0-9a-fA-F-]+$/ {print $2; exit}')
|
||||
|
||||
@@ -91,6 +91,9 @@ jobs:
|
||||
nix build .#metal-toolchain
|
||||
fi
|
||||
|
||||
# Build mlx (depends on metal-toolchain)
|
||||
nix build .#mlx
|
||||
|
||||
- name: Build all Nix outputs
|
||||
run: |
|
||||
nix flake show --json | jq -r '
|
||||
|
||||
@@ -38,5 +38,3 @@ bench/**/*.json
|
||||
|
||||
# tmp
|
||||
tmp/models
|
||||
/build/exo
|
||||
/.claude/skills
|
||||
Generated
+31
@@ -0,0 +1,31 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<module type="EMPTY_MODULE" version="4">
|
||||
<component name="FacetManager">
|
||||
<facet type="Python" name="Python facet">
|
||||
<configuration sdkName="Python 3.13 virtualenv at ~/Desktop/exo/.venv" />
|
||||
</facet>
|
||||
</component>
|
||||
<component name="Go" enabled="true" />
|
||||
<component name="NewModuleRootManager">
|
||||
<content url="file://$MODULE_DIR$">
|
||||
<sourceFolder url="file://$MODULE_DIR$/scripts/src" isTestSource="false" />
|
||||
<sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
|
||||
<sourceFolder url="file://$MODULE_DIR$/rust/exo_pyo3_bindings/src" isTestSource="false" />
|
||||
<sourceFolder url="file://$MODULE_DIR$/rust/exo_pyo3_bindings/tests" isTestSource="true" />
|
||||
<sourceFolder url="file://$MODULE_DIR$/rust/util/src" isTestSource="false" />
|
||||
<sourceFolder url="file://$MODULE_DIR$/rust/networking/examples" isTestSource="false" />
|
||||
<sourceFolder url="file://$MODULE_DIR$/rust/networking/src" isTestSource="false" />
|
||||
<sourceFolder url="file://$MODULE_DIR$/rust/networking/tests" isTestSource="true" />
|
||||
<sourceFolder url="file://$MODULE_DIR$/rust/system_custodian/src" isTestSource="false" />
|
||||
<excludeFolder url="file://$MODULE_DIR$/.venv" />
|
||||
<excludeFolder url="file://$MODULE_DIR$/.direnv" />
|
||||
<excludeFolder url="file://$MODULE_DIR$/build" />
|
||||
<excludeFolder url="file://$MODULE_DIR$/dist" />
|
||||
<excludeFolder url="file://$MODULE_DIR$/.go_cache" />
|
||||
<excludeFolder url="file://$MODULE_DIR$/rust/target" />
|
||||
</content>
|
||||
<orderEntry type="jdk" jdkName="Python 3.13 (exo)" jdkType="Python SDK" />
|
||||
<orderEntry type="sourceFolder" forTests="false" />
|
||||
<orderEntry type="library" name="Python 3.13 virtualenv at ~/Desktop/exo/.venv interpreter library" level="application" />
|
||||
</component>
|
||||
</module>
|
||||
Generated
+1
-1
@@ -1,6 +1,6 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="ExternalDependencies">
|
||||
<plugin id="al.aoli.intellijdirenv" />
|
||||
<plugin id="systems.fehn.intellijdirenv" />
|
||||
</component>
|
||||
</project>
|
||||
+14
@@ -0,0 +1,14 @@
|
||||
<component name="InspectionProjectProfileManager">
|
||||
<profile version="1.0">
|
||||
<option name="myName" value="Project Default" />
|
||||
<inspection_tool class="PyCompatibilityInspection" enabled="true" level="WARNING" enabled_by_default="true">
|
||||
<option name="ourVersions">
|
||||
<value>
|
||||
<list size="1">
|
||||
<item index="0" class="java.lang.String" itemvalue="3.14" />
|
||||
</list>
|
||||
</value>
|
||||
</option>
|
||||
</inspection_tool>
|
||||
</profile>
|
||||
</component>
|
||||
Generated
+3
@@ -4,4 +4,7 @@
|
||||
<option name="sdkName" value="Python 3.13 (exo)" />
|
||||
</component>
|
||||
<component name="ProjectRootManager" version="2" project-jdk-name="Python 3.13 (exo)" project-jdk-type="Python SDK" />
|
||||
<component name="PythonCompatibilityInspectionAdvertiser">
|
||||
<option name="version" value="3" />
|
||||
</component>
|
||||
</project>
|
||||
@@ -215,22 +215,6 @@ class StreamContext:
|
||||
traceback: object | None = ...,
|
||||
) -> None: ...
|
||||
|
||||
def device_info() -> dict[str, str | int]:
|
||||
"""
|
||||
Get information about the GPU device and system settings.
|
||||
|
||||
Currently returns:
|
||||
|
||||
* ``architecture``
|
||||
* ``max_buffer_size``
|
||||
* ``max_recommended_working_set_size``
|
||||
* ``memory_size``
|
||||
* ``resource_limit``
|
||||
|
||||
Returns:
|
||||
dict: A dictionary with string keys and string or integer values.
|
||||
"""
|
||||
|
||||
def abs(a: array, /, *, stream: Stream | Device | None = ...) -> array:
|
||||
"""
|
||||
Element-wise absolute value.
|
||||
@@ -1767,12 +1751,12 @@ def clip(
|
||||
array: The clipped array.
|
||||
"""
|
||||
|
||||
def compile[F: Callable[..., object]](
|
||||
fun: F,
|
||||
def compile(
|
||||
fun: Callable,
|
||||
inputs: object | None = ...,
|
||||
outputs: object | None = ...,
|
||||
shapeless: bool = ...,
|
||||
) -> F:
|
||||
) -> Callable:
|
||||
"""
|
||||
Returns a compiled function which produces the same output as ``fun``.
|
||||
|
||||
@@ -2396,7 +2380,7 @@ def degrees(a: array, /, *, stream: Stream | Device | None = ...) -> array:
|
||||
array: The angles in degrees.
|
||||
"""
|
||||
|
||||
def depends[T](inputs: T, dependencies: array | Sequence[array]) -> T:
|
||||
def depends(inputs: array | Sequence[array], dependencies: array | Sequence[array]):
|
||||
"""
|
||||
Insert dependencies between arrays in the graph. The outputs are
|
||||
identical to ``inputs`` but with dependencies on ``dependencies``.
|
||||
@@ -2915,8 +2899,8 @@ def gather_mm(
|
||||
a: array,
|
||||
b: array,
|
||||
/,
|
||||
lhs_indices: array | None = ...,
|
||||
rhs_indices: array | None = ...,
|
||||
lhs_indices: array,
|
||||
rhs_indices: array,
|
||||
*,
|
||||
sorted_indices: bool = ...,
|
||||
stream: Stream | Device | None = ...,
|
||||
@@ -4707,7 +4691,6 @@ def softmax(
|
||||
/,
|
||||
axis: int | Sequence[int] | None = ...,
|
||||
*,
|
||||
precise: bool = ...,
|
||||
stream: Stream | Device | None = ...,
|
||||
) -> array:
|
||||
"""
|
||||
|
||||
@@ -1,5 +1,9 @@
|
||||
from .layers import *
|
||||
from .utils import *
|
||||
"""
|
||||
This type stub file was generated by pyright.
|
||||
"""
|
||||
|
||||
from layers import *
|
||||
from utils import *
|
||||
|
||||
from . import init as init
|
||||
from . import losses as losses
|
||||
@@ -1,16 +1,20 @@
|
||||
from .activations import *
|
||||
from .base import *
|
||||
from .containers import *
|
||||
from .convolution import *
|
||||
from .convolution_transpose import *
|
||||
from .distributed import *
|
||||
from .dropout import *
|
||||
from .embedding import *
|
||||
from .linear import *
|
||||
from .normalization import *
|
||||
from .pooling import *
|
||||
from .positional_encoding import *
|
||||
from .quantized import *
|
||||
from .recurrent import *
|
||||
from .transformer import *
|
||||
from .upsample import *
|
||||
"""
|
||||
This type stub file was generated by pyright.
|
||||
"""
|
||||
|
||||
from activations import *
|
||||
from base import *
|
||||
from containers import *
|
||||
from convolution import *
|
||||
from convolution_transpose import *
|
||||
from distributed import *
|
||||
from dropout import *
|
||||
from embedding import *
|
||||
from linear import *
|
||||
from normalization import *
|
||||
from pooling import *
|
||||
from positional_encoding import *
|
||||
from quantized import *
|
||||
from recurrent import *
|
||||
from transformer import *
|
||||
from upsample import *
|
||||
@@ -53,14 +53,10 @@ class Module(dict):
|
||||
mx.eval(model.parameters())
|
||||
"""
|
||||
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
|
||||
__call__: Callable
|
||||
def __init__(self) -> None:
|
||||
"""Should be called by the subclasses of ``Module``."""
|
||||
|
||||
def __getitem__(self, key: str) -> mx.array | Module: ...
|
||||
def get(
|
||||
self, key: str, default: mx.array | Module | None = ...
|
||||
) -> mx.array | Module | None: ...
|
||||
@property
|
||||
def training(self): # -> bool:
|
||||
"""Boolean indicating if the model is in training mode."""
|
||||
@@ -204,7 +200,7 @@ class Module(dict):
|
||||
) -> mx.MX_ARRAY_TREE: # -> dict[Any, Any | dict[Any, Any | dict[Any, Any] | list[Any]] | dict[Any, Any] | list[Any]]:
|
||||
"""Return the submodules that do not contain other modules."""
|
||||
|
||||
def update(self, parameters: dict[str, Any], strict: bool = ...) -> Module:
|
||||
def update(self, parameters: dict, strict: bool = ...) -> Module:
|
||||
"""Replace the parameters of this Module with the provided ones in the
|
||||
dict of dicts and lists.
|
||||
|
||||
|
||||
@@ -32,7 +32,6 @@ class Conv1d(Module):
|
||||
"""
|
||||
|
||||
weight: mx.array
|
||||
bias: mx.array | None
|
||||
groups: int
|
||||
def __init__(
|
||||
self,
|
||||
|
||||
@@ -40,10 +40,6 @@ class Linear(Module):
|
||||
bias (bool, optional): If set to ``False`` then the layer will
|
||||
not use a bias. Default is ``True``.
|
||||
"""
|
||||
|
||||
weight: mx.array
|
||||
bias: mx.array | None
|
||||
|
||||
def __init__(self, input_dims: int, output_dims: int, bias: bool = ...) -> None: ...
|
||||
def __call__(self, x: mx.array) -> mx.array: ...
|
||||
def to_quantized(
|
||||
|
||||
@@ -88,9 +88,6 @@ class RMSNorm(Module):
|
||||
dims (int): The feature dimension of the input to normalize over
|
||||
eps (float): A small additive constant for numerical stability
|
||||
"""
|
||||
|
||||
weight: mx.array
|
||||
|
||||
def __init__(self, dims: int, eps: float = ...) -> None: ...
|
||||
def __call__(self, x) -> mx.array: ...
|
||||
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
This type stub file was generated by pyright.
|
||||
"""
|
||||
|
||||
from typing import Any, Callable, Optional, Union
|
||||
from typing import Callable, Optional, Union
|
||||
|
||||
import mlx.core as mx
|
||||
from base import Module
|
||||
@@ -13,10 +13,8 @@ def quantize(
|
||||
bits: int = ...,
|
||||
*,
|
||||
mode: str = ...,
|
||||
class_predicate: Optional[
|
||||
Callable[[str, Module], Union[bool, dict[str, Any]]]
|
||||
] = ...,
|
||||
) -> None:
|
||||
class_predicate: Optional[Callable[[str, Module], Union[bool, dict]]] = ...,
|
||||
): # -> None:
|
||||
"""Quantize the sub-modules of a module according to a predicate.
|
||||
|
||||
By default all layers that define a ``to_quantized(group_size, bits)``
|
||||
|
||||
@@ -7,10 +7,7 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
|
||||
from mlx.core import MX_ARRAY_TREE
|
||||
|
||||
def tree_map(
|
||||
fn: Callable[..., Any],
|
||||
tree: Any,
|
||||
*rest: Any,
|
||||
is_leaf: Callable[..., bool] | None = ...,
|
||||
fn: Callable, tree: Any, *rest: Any, is_leaf: Optional[Callable] = ...
|
||||
) -> Any:
|
||||
"""Applies ``fn`` to the leaves of the Python tree ``tree`` and
|
||||
returns a new collection with the results.
|
||||
@@ -47,11 +44,11 @@ def tree_map(
|
||||
"""
|
||||
|
||||
def tree_map_with_path(
|
||||
fn: Callable[..., Any],
|
||||
fn: Callable,
|
||||
tree: Any,
|
||||
*rest: Any,
|
||||
is_leaf: Callable[..., bool] | None = ...,
|
||||
path: str | None = ...,
|
||||
is_leaf: Optional[Callable] = ...,
|
||||
path: Optional[Any] = ...,
|
||||
) -> Any:
|
||||
"""Applies ``fn`` to the path and leaves of the Python tree ``tree`` and
|
||||
returns a new collection with the results.
|
||||
@@ -83,9 +80,9 @@ def tree_map_with_path(
|
||||
def tree_flatten(
|
||||
tree: Any,
|
||||
prefix: str = ...,
|
||||
is_leaf: Callable[..., bool] | None = ...,
|
||||
destination: list[tuple[str, Any]] | dict[str, Any] | None = ...,
|
||||
) -> list[tuple[str, Any]] | dict[str, Any]:
|
||||
is_leaf: Optional[Callable] = ...,
|
||||
destination: Optional[Union[List[Tuple[str, Any]], Dict[str, Any]]] = ...,
|
||||
) -> Union[List[Tuple[str, Any]], Dict[str, Any]]:
|
||||
"""Flattens a Python tree to a list of key, value tuples.
|
||||
|
||||
The keys are using the dot notation to define trees of arbitrary depth and
|
||||
@@ -121,7 +118,7 @@ def tree_flatten(
|
||||
the Python tree.
|
||||
"""
|
||||
|
||||
def tree_unflatten(tree: list[tuple[str, Any]] | dict[str, Any]) -> Any:
|
||||
def tree_unflatten(tree: Union[List[Tuple[str, Any]], Dict[str, Any]]) -> Any:
|
||||
"""Recreate a Python tree from its flat representation.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@@ -1,5 +0,0 @@
|
||||
"""
|
||||
This type stub file was generated by pyright.
|
||||
"""
|
||||
|
||||
__version__ = ...
|
||||
@@ -3,12 +3,13 @@ This type stub file was generated by pyright.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable, Generator, List, Optional, Tuple, Union
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
from dataclasses import dataclass
|
||||
from collections import deque
|
||||
from typing import Any, Callable, Generator, List, Optional, Sequence, Tuple, Union
|
||||
from transformers import PreTrainedTokenizer
|
||||
|
||||
from .tokenizer_utils import TokenizerWrapper
|
||||
|
||||
DEFAULT_PROMPT = ...
|
||||
@@ -28,9 +29,8 @@ def str2bool(string): # -> bool:
|
||||
...
|
||||
def setup_arg_parser(): # -> ArgumentParser:
|
||||
"""Set up and return the argument parser."""
|
||||
...
|
||||
|
||||
generation_stream: mx.Stream
|
||||
generation_stream = ...
|
||||
|
||||
@contextlib.contextmanager
|
||||
def wired_limit(
|
||||
@@ -43,7 +43,6 @@ def wired_limit(
|
||||
async eval could be running pass in the streams to synchronize with prior
|
||||
to exiting the context manager.
|
||||
"""
|
||||
...
|
||||
@dataclass
|
||||
class GenerationResponse:
|
||||
"""
|
||||
@@ -74,11 +73,9 @@ class GenerationResponse:
|
||||
finish_reason: Optional[str] = ...
|
||||
|
||||
def maybe_quantize_kv_cache(
|
||||
prompt_cache: Any,
|
||||
quantized_kv_start: int | None,
|
||||
kv_group_size: int | None,
|
||||
kv_bits: int | None,
|
||||
) -> None: ...
|
||||
prompt_cache, quantized_kv_start, kv_group_size, kv_bits
|
||||
): # -> None:
|
||||
...
|
||||
def generate_step(
|
||||
prompt: mx.array,
|
||||
model: nn.Module,
|
||||
@@ -92,7 +89,7 @@ def generate_step(
|
||||
kv_bits: Optional[int] = ...,
|
||||
kv_group_size: int = ...,
|
||||
quantized_kv_start: int = ...,
|
||||
prompt_progress_callback: Optional[Callable[[int, int], None]] = ...,
|
||||
prompt_progress_callback: Optional[Callable[[int], int]] = ...,
|
||||
input_embeddings: Optional[mx.array] = ...,
|
||||
) -> Generator[Tuple[mx.array, mx.array], None, None]:
|
||||
"""
|
||||
@@ -118,7 +115,7 @@ def generate_step(
|
||||
kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
|
||||
quantized_kv_start (int): Step to begin using a quantized KV cache.
|
||||
when ``kv_bits`` is non-None. Default: ``0``.
|
||||
prompt_progress_callback (Callable[[int, int], None]): A call-back which takes the
|
||||
prompt_progress_callback (Callable[[int], int]): A call-back which takes the
|
||||
prompt tokens processed so far and the total number of prompt tokens.
|
||||
input_embeddings (mx.array, optional): Input embeddings to use instead of or in
|
||||
conjunction with prompt tokens. Default: ``None``.
|
||||
@@ -126,7 +123,6 @@ def generate_step(
|
||||
Yields:
|
||||
Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
|
||||
"""
|
||||
...
|
||||
|
||||
def speculative_generate_step(
|
||||
prompt: mx.array,
|
||||
@@ -172,7 +168,6 @@ def speculative_generate_step(
|
||||
Tuple[mx.array, mx.array, bool]: One token, a vector of log probabilities,
|
||||
and a bool indicating if the token was generated by the draft model
|
||||
"""
|
||||
...
|
||||
|
||||
def stream_generate(
|
||||
model: nn.Module,
|
||||
@@ -180,7 +175,7 @@ def stream_generate(
|
||||
prompt: Union[str, mx.array, List[int]],
|
||||
max_tokens: int = ...,
|
||||
draft_model: Optional[nn.Module] = ...,
|
||||
**kwargs: Any,
|
||||
**kwargs: object,
|
||||
) -> Generator[GenerationResponse, None, None]:
|
||||
"""
|
||||
A generator producing text based on the given prompt from the model.
|
||||
@@ -202,7 +197,6 @@ def stream_generate(
|
||||
GenerationResponse: An instance containing the generated text segment and
|
||||
associated metadata. See :class:`GenerationResponse` for details.
|
||||
"""
|
||||
...
|
||||
|
||||
def generate(
|
||||
model: nn.Module,
|
||||
@@ -223,9 +217,6 @@ def generate(
|
||||
kwargs: The remaining options get passed to :func:`stream_generate`.
|
||||
See :func:`stream_generate` for more details.
|
||||
"""
|
||||
...
|
||||
|
||||
def _merge_caches(caches: List[List[Any]]) -> List[Any]: ...
|
||||
@dataclass
|
||||
class BatchStats:
|
||||
"""
|
||||
@@ -249,263 +240,10 @@ class BatchStats:
|
||||
generation_time: float = ...
|
||||
peak_memory: float = ...
|
||||
|
||||
class SequenceStateMachine:
|
||||
"""A state machine that uses one Aho-Corasick trie per state to efficiently
|
||||
track state across a generated sequence.
|
||||
|
||||
The transitions are provided as state -> [(sequence, new_state)].
|
||||
|
||||
Example:
|
||||
|
||||
sm = SequenceStateMachine(
|
||||
transitions={
|
||||
"normal": [
|
||||
(think_start_tokens, "reasoning"),
|
||||
(tool_start_tokens, "tool"),
|
||||
(eos, None),
|
||||
],
|
||||
"reasoning": [
|
||||
(think_end_tokens, "normal"),
|
||||
(eos, None),
|
||||
],
|
||||
"tool": [
|
||||
(tool_end_tokens, None),
|
||||
(eos, None)
|
||||
],
|
||||
},
|
||||
initial="normal"
|
||||
)
|
||||
"""
|
||||
def __init__(self, transitions=..., initial=...) -> None: ...
|
||||
def __deepcopy__(self, memo): # -> SequenceStateMachine:
|
||||
...
|
||||
def make_state(self): # -> tuple[str, Any, dict[Any, Any]]:
|
||||
...
|
||||
@staticmethod
|
||||
def match(state, x): # -> tuple[tuple[Any, Any | None, Any], Any | None, Any]:
|
||||
...
|
||||
|
||||
class PromptProcessingBatch:
|
||||
"""
|
||||
A batch processor for prompt tokens with support for incremental processing.
|
||||
|
||||
This class handles batched prompt processing, managing KV caches and preparing
|
||||
tokens for generation. It supports extending, filtering, and splitting batches.
|
||||
"""
|
||||
@dataclass
|
||||
class Response:
|
||||
uid: int
|
||||
progress: tuple
|
||||
end_of_segment: bool
|
||||
end_of_prompt: bool
|
||||
...
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: nn.Module,
|
||||
uids: List[int],
|
||||
caches: List[List[Any]],
|
||||
tokens: Optional[List[List[int]]] = ...,
|
||||
prefill_step_size: int = ...,
|
||||
samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
|
||||
fallback_sampler: Optional[Callable[[mx.array], mx.array]] = ...,
|
||||
logits_processors: Optional[
|
||||
List[List[Callable[[mx.array, mx.array], mx.array]]]
|
||||
] = ...,
|
||||
state_machines: Optional[List[SequenceStateMachine]] = ...,
|
||||
max_tokens: Optional[List[int]] = ...,
|
||||
) -> None: ...
|
||||
def __len__(self): # -> int:
|
||||
...
|
||||
def extract_cache(self, idx: int) -> List[Any]: ...
|
||||
def extend(self, batch): # -> None:
|
||||
...
|
||||
def split(self, indices: List[int]): # -> Self:
|
||||
...
|
||||
def filter(self, keep: List[int]): # -> None:
|
||||
...
|
||||
def prompt(self, tokens: List[List[int]]): # -> None:
|
||||
"""
|
||||
Process prompt tokens through the model.
|
||||
|
||||
Args:
|
||||
tokens: List of token sequences to process.
|
||||
"""
|
||||
...
|
||||
|
||||
def generate(self, tokens: List[List[int]]): # -> GenerationBatch:
|
||||
"""
|
||||
Transition from prompt processing to generation.
|
||||
|
||||
Args:
|
||||
tokens: Final tokens for each sequence to start generation.
|
||||
|
||||
Returns:
|
||||
A GenerationBatch ready for token generation.
|
||||
"""
|
||||
...
|
||||
|
||||
@classmethod
|
||||
def empty(
|
||||
cls,
|
||||
model: nn.Module,
|
||||
fallback_sampler: Callable[[mx.array], mx.array],
|
||||
prefill_step_size: int = ...,
|
||||
): # -> Self:
|
||||
...
|
||||
|
||||
class GenerationBatch:
|
||||
"""
|
||||
A batched token generator that manages multiple sequences in parallel.
|
||||
|
||||
This class handles the generation phase after prompt processing, managing
|
||||
KV caches, sampling, and stop sequence detection for multiple sequences.
|
||||
"""
|
||||
@dataclass
|
||||
class Response:
|
||||
uid: int
|
||||
token: int
|
||||
logprobs: mx.array
|
||||
finish_reason: Optional[str]
|
||||
current_state: Optional[str]
|
||||
match_sequence: Optional[List[int]]
|
||||
prompt_cache: Optional[List[Any]]
|
||||
all_tokens: Optional[List[int]]
|
||||
...
|
||||
|
||||
model: nn.Module
|
||||
uids: List[int]
|
||||
prompt_cache: List[Any]
|
||||
tokens: List[List[int]]
|
||||
samplers: Optional[List[Callable[[mx.array], mx.array]]]
|
||||
fallback_sampler: Callable[[mx.array], mx.array]
|
||||
logits_processors: Optional[List[List[Callable[[mx.array, mx.array], mx.array]]]]
|
||||
state_machines: List[SequenceStateMachine]
|
||||
max_tokens: List[int]
|
||||
_current_tokens: Optional[mx.array]
|
||||
_current_logprobs: mx.array | List[mx.array]
|
||||
_next_tokens: Optional[mx.array]
|
||||
_next_logprobs: mx.array | List[mx.array]
|
||||
_token_context: List[Any]
|
||||
_num_tokens: List[int]
|
||||
_matcher_states: List[Any]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: nn.Module,
|
||||
uids: List[int],
|
||||
inputs: mx.array,
|
||||
prompt_cache: List[Any],
|
||||
tokens: List[List[int]],
|
||||
samplers: Optional[List[Callable[[mx.array], mx.array]]],
|
||||
fallback_sampler: Callable[[mx.array], mx.array],
|
||||
logits_processors: Optional[
|
||||
List[List[Callable[[mx.array, mx.array], mx.array]]]
|
||||
],
|
||||
state_machines: List[SequenceStateMachine],
|
||||
max_tokens: List[int],
|
||||
) -> None: ...
|
||||
def __len__(self) -> int: ...
|
||||
def extend(self, batch: GenerationBatch) -> None: ...
|
||||
def extract_cache(self, idx: int) -> List[Any]: ...
|
||||
def filter(self, keep: List[int]) -> None: ...
|
||||
def _step(self) -> Tuple[List[int], List[mx.array]]: ...
|
||||
def next(self) -> List[Response]:
|
||||
"""
|
||||
Generate the next batch of tokens.
|
||||
|
||||
Returns:
|
||||
List of Response objects for each sequence in the batch.
|
||||
"""
|
||||
...
|
||||
|
||||
@classmethod
|
||||
def empty(
|
||||
cls, model: nn.Module, fallback_sampler: Callable[[mx.array], mx.array]
|
||||
): # -> Self:
|
||||
...
|
||||
|
||||
class BatchGenerator:
|
||||
"""
|
||||
A batch generator implements continuous batching.
|
||||
|
||||
This class provides automatic management of prompt processing and generation
|
||||
batches, handling the transition between the two.
|
||||
|
||||
It also allows for segmented prompt processing which guarantees that the
|
||||
generator will stop at these boundaries when processing an input.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
model: nn.Module,
|
||||
max_tokens: int = ...,
|
||||
stop_tokens: Optional[Sequence[Sequence[int]]] = ...,
|
||||
sampler: Optional[Callable[[mx.array], mx.array]] = ...,
|
||||
logits_processors: Optional[
|
||||
List[Callable[[mx.array, mx.array], mx.array]]
|
||||
] = ...,
|
||||
completion_batch_size: int = ...,
|
||||
prefill_batch_size: int = ...,
|
||||
prefill_step_size: int = ...,
|
||||
) -> None: ...
|
||||
def close(self) -> None: ...
|
||||
def __del__(self): # -> None:
|
||||
...
|
||||
@contextlib.contextmanager
|
||||
def stats(self, stats=...): # -> Generator[Any | BatchStats, Any, None]:
|
||||
...
|
||||
_unprocessed_sequences: deque[tuple[Any, ...]]
|
||||
_prompt_batch: PromptProcessingBatch
|
||||
_generation_batch: GenerationBatch
|
||||
_currently_processing: list[Any]
|
||||
_gen_tokens_counter: int
|
||||
_steps_counter: int
|
||||
def _next(
|
||||
self,
|
||||
) -> tuple[
|
||||
List[PromptProcessingBatch.Response], List[GenerationBatch.Response]
|
||||
]: ...
|
||||
def insert(
|
||||
self,
|
||||
prompts: List[List[int]],
|
||||
max_tokens: Optional[List[int]] = ...,
|
||||
caches: Optional[List[List[Any]]] = ...,
|
||||
all_tokens: Optional[List[List[int]]] = ...,
|
||||
samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
|
||||
logits_processors: Optional[
|
||||
List[List[Callable[[mx.array, mx.array], mx.array]]]
|
||||
] = ...,
|
||||
state_machines: Optional[List[SequenceStateMachine]] = ...,
|
||||
) -> List[int]: ...
|
||||
def insert_segments(
|
||||
self,
|
||||
segments: List[List[List[int]]],
|
||||
max_tokens: Optional[List[int]] = ...,
|
||||
caches: Optional[List[List[Any]]] = ...,
|
||||
all_tokens: Optional[List[List[int]]] = ...,
|
||||
samplers: Optional[List[Callable[[mx.array], mx.array]]] = ...,
|
||||
logits_processors: Optional[
|
||||
List[List[Callable[[mx.array, mx.array], mx.array]]]
|
||||
] = ...,
|
||||
state_machines: Optional[List[SequenceStateMachine]] = ...,
|
||||
) -> List[int]: ...
|
||||
def extract_cache(self, uids: List[int]) -> dict[int, Any]: ...
|
||||
def remove(
|
||||
self, uids: List[int], return_prompt_caches: bool = ...
|
||||
) -> dict[int, Any]: ...
|
||||
@property
|
||||
def prompt_cache_nbytes(self) -> int: ...
|
||||
def next(
|
||||
self,
|
||||
) -> tuple[
|
||||
List[PromptProcessingBatch.Response], List[GenerationBatch.Response]
|
||||
]: ...
|
||||
def next_generated(self) -> List[GenerationBatch.Response]: ...
|
||||
|
||||
@dataclass
|
||||
class BatchResponse:
|
||||
"""
|
||||
A data object to hold a batch generation response.
|
||||
An data object to hold a batch generation response.
|
||||
|
||||
Args:
|
||||
texts: (List[str]): The generated text for each prompt.
|
||||
@@ -514,18 +252,55 @@ class BatchResponse:
|
||||
|
||||
texts: List[str]
|
||||
stats: BatchStats
|
||||
caches: Optional[List[List[Any]]]
|
||||
...
|
||||
|
||||
@dataclass
|
||||
class Batch:
|
||||
uids: List[int]
|
||||
y: mx.array
|
||||
logprobs: mx.array
|
||||
max_tokens: List[int]
|
||||
num_tokens: List[int]
|
||||
cache: List[Any]
|
||||
def __len__(self): # -> int:
|
||||
...
|
||||
def filter(self, keep_idx: List[int]): # -> None:
|
||||
...
|
||||
def extend(self, other): # -> None:
|
||||
...
|
||||
|
||||
class BatchGenerator:
|
||||
@dataclass
|
||||
class Response:
|
||||
uid: int
|
||||
token: int
|
||||
logprobs: mx.array
|
||||
finish_reason: Optional[str]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model,
|
||||
max_tokens: int = ...,
|
||||
stop_tokens: Optional[set] = ...,
|
||||
sampler: Optional[Callable[[mx.array], mx.array]] = ...,
|
||||
completion_batch_size: int = ...,
|
||||
prefill_batch_size: int = ...,
|
||||
prefill_step_size: int = ...,
|
||||
) -> None: ...
|
||||
def insert(
|
||||
self, prompts, max_tokens: Union[List[int], int, None] = ...
|
||||
): # -> list[Any]:
|
||||
...
|
||||
def stats(self): # -> BatchStats:
|
||||
...
|
||||
def next(self): # -> list[Any]:
|
||||
...
|
||||
|
||||
def batch_generate(
|
||||
model,
|
||||
tokenizer,
|
||||
prompts: List[List[int]],
|
||||
prompt_caches: Optional[List[List[Any]]] = ...,
|
||||
prompts: List[int],
|
||||
max_tokens: Union[int, List[int]] = ...,
|
||||
verbose: bool = ...,
|
||||
return_prompt_caches: bool = ...,
|
||||
logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = ...,
|
||||
**kwargs,
|
||||
) -> BatchResponse:
|
||||
"""
|
||||
@@ -534,22 +309,14 @@ def batch_generate(
|
||||
Args:
|
||||
model (nn.Module): The language model.
|
||||
tokenizer (PreTrainedTokenizer): The tokenizer.
|
||||
prompts (List[List[int]]): The input prompts.
|
||||
prompt_caches (List[List[Any]], optional): Pre-computed prompt-caches
|
||||
for each input prompt. Note, unlike ``generate_step``, the caches
|
||||
won't be updated in-place.
|
||||
prompt (List[List[int]]): The input prompts.
|
||||
verbose (bool): If ``True``, print tokens and timing information.
|
||||
Default: ``False``.
|
||||
max_tokens (Union[int, List[int]): Maximum number of output tokens. This
|
||||
can be per prompt if a list is provided.
|
||||
return_prompt_caches (bool): Return the prompt caches in the batch
|
||||
responses. Default: ``False``.
|
||||
logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
|
||||
A list of functions that take tokens and logits and return the processed logits. Default: ``None``.
|
||||
kwargs: The remaining options get passed to :obj:`BatchGenerator`.
|
||||
See :obj:`BatchGenerator` for more details.
|
||||
"""
|
||||
...
|
||||
|
||||
def main(): # -> None:
|
||||
...
|
||||
|
||||
@@ -1,19 +0,0 @@
|
||||
"""
|
||||
This type stub file was generated by pyright.
|
||||
"""
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
from functools import partial
|
||||
|
||||
@partial(mx.compile, shapeless=True)
|
||||
def swiglu(gate, x): ...
|
||||
@partial(mx.compile, shapeless=True)
|
||||
def xielu(x, alpha_p, alpha_n, beta, eps): # -> array:
|
||||
...
|
||||
|
||||
class XieLU(nn.Module):
|
||||
def __init__(
|
||||
self, alpha_p_init=..., alpha_n_init=..., beta=..., eps=...
|
||||
) -> None: ...
|
||||
def __call__(self, x: mx.array) -> mx.array: ...
|
||||
@@ -3,7 +3,7 @@ This type stub file was generated by pyright.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
from typing import Optional
|
||||
|
||||
import mlx.core as mx
|
||||
|
||||
@@ -37,10 +37,10 @@ def quantized_scaled_dot_product_attention(
|
||||
bits: int = ...,
|
||||
) -> mx.array: ...
|
||||
def scaled_dot_product_attention(
|
||||
queries: mx.array,
|
||||
keys: mx.array,
|
||||
values: mx.array,
|
||||
cache: Optional[Any],
|
||||
queries,
|
||||
keys,
|
||||
values,
|
||||
cache,
|
||||
scale: float,
|
||||
mask: Optional[mx.array],
|
||||
sinks: Optional[mx.array] = ...,
|
||||
|
||||
@@ -16,7 +16,7 @@ class Cache(Protocol):
|
||||
self, keys: mx.array, values: mx.array
|
||||
) -> tuple[mx.array, mx.array]: ...
|
||||
@property
|
||||
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
def state(self) -> tuple[mx.array, mx.array]: ...
|
||||
@state.setter
|
||||
def state(self, v) -> None: ...
|
||||
|
||||
@@ -88,18 +88,17 @@ def create_attention_mask(
|
||||
) -> array | Literal["causal"] | None: ...
|
||||
|
||||
class _BaseCache(Cache):
|
||||
keys: mx.array | None
|
||||
values: mx.array | None
|
||||
keys: mx.array
|
||||
values: mx.array
|
||||
offset: int
|
||||
@property
|
||||
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
def state(self) -> tuple[mx.array, mx.array]: ...
|
||||
@state.setter
|
||||
def state(self, v) -> None: ...
|
||||
@property
|
||||
def meta_state(self) -> Literal[""]: ...
|
||||
@meta_state.setter
|
||||
def meta_state(self, v) -> None: ...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def is_trimmable(self) -> Literal[False]: ...
|
||||
@classmethod
|
||||
def from_state(cls, state, meta_state) -> Self: ...
|
||||
@@ -115,13 +114,15 @@ class ConcatenateKVCache(_BaseCache):
|
||||
def update_and_fetch(self, keys, values): # -> tuple[Any | array, Any | array]:
|
||||
...
|
||||
@property
|
||||
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
def state(self): # -> tuple[Any | array | None, Any | array | None]:
|
||||
...
|
||||
@state.setter
|
||||
def state(self, v): # -> None:
|
||||
...
|
||||
def is_trimmable(self): # -> Literal[True]:
|
||||
...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def trim(self, n): # -> int:
|
||||
...
|
||||
def make_mask(self, *args, **kwargs): # -> array | Literal['causal'] | None:
|
||||
...
|
||||
|
||||
@@ -131,7 +132,10 @@ class QuantizedKVCache(_BaseCache):
|
||||
def update_and_fetch(self, keys, values): # -> Any:
|
||||
...
|
||||
@property
|
||||
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
def state(
|
||||
self,
|
||||
): # -> tuple[Any | tuple[array, array, array] | None, Any | tuple[array, array, array] | None] | Any:
|
||||
...
|
||||
@state.setter
|
||||
def state(self, v): # -> None:
|
||||
...
|
||||
@@ -143,7 +147,8 @@ class QuantizedKVCache(_BaseCache):
|
||||
...
|
||||
def is_trimmable(self): # -> Literal[True]:
|
||||
...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def trim(self, n): # -> int:
|
||||
...
|
||||
def make_mask(self, *args, **kwargs): # -> array | Literal['causal'] | None:
|
||||
...
|
||||
|
||||
@@ -155,30 +160,22 @@ class KVCache(_BaseCache):
|
||||
@property
|
||||
def state(
|
||||
self,
|
||||
) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
) -> tuple[array, array]: ...
|
||||
@state.setter
|
||||
def state(self, v) -> None: ...
|
||||
def is_trimmable(self): # -> Literal[True]:
|
||||
...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def trim(self, n): # -> int:
|
||||
...
|
||||
def to_quantized(
|
||||
self, group_size: int = ..., bits: int = ...
|
||||
) -> QuantizedKVCache: ...
|
||||
def make_mask(
|
||||
self, *args: Any, **kwargs: Any
|
||||
) -> mx.array | Literal["causal"] | None: ...
|
||||
def make_mask(self, *args, **kwargs): # -> array | Literal['causal'] | None:
|
||||
...
|
||||
|
||||
class RotatingKVCache(_BaseCache):
|
||||
step = ...
|
||||
keys: mx.array | None
|
||||
values: mx.array | None
|
||||
keep: int
|
||||
max_size: int
|
||||
_idx: int
|
||||
def __init__(self, max_size, keep=...) -> None: ...
|
||||
def _trim(
|
||||
self, trim_size: int, v: mx.array, append: mx.array | None = ...
|
||||
) -> mx.array: ...
|
||||
def update_and_fetch(
|
||||
self, keys, values
|
||||
): # -> tuple[array | Any, array | Any] | tuple[array | Any, array | Any | None]:
|
||||
@@ -186,7 +183,8 @@ class RotatingKVCache(_BaseCache):
|
||||
@property
|
||||
def state(
|
||||
self,
|
||||
) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
): # -> tuple[Any | array, Any | array] | tuple[Any | array | None, Any | array | None]:
|
||||
...
|
||||
@state.setter
|
||||
def state(self, v): # -> None:
|
||||
...
|
||||
@@ -198,7 +196,8 @@ class RotatingKVCache(_BaseCache):
|
||||
...
|
||||
def is_trimmable(self): # -> bool:
|
||||
...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def trim(self, n): # -> int:
|
||||
...
|
||||
def to_quantized(
|
||||
self, group_size: int = ..., bits: int = ...
|
||||
) -> QuantizedKVCache: ...
|
||||
@@ -213,7 +212,8 @@ class ArraysCache(_BaseCache):
|
||||
...
|
||||
def __getitem__(self, idx): ...
|
||||
@property
|
||||
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
def state(self): # -> list[Any | array] | list[array]:
|
||||
...
|
||||
@state.setter
|
||||
def state(self, v): # -> None:
|
||||
...
|
||||
@@ -227,7 +227,8 @@ class ArraysCache(_BaseCache):
|
||||
In-place extend this cache with the other cache.
|
||||
"""
|
||||
|
||||
def make_mask(self, N: int) -> mx.array | None: ...
|
||||
def make_mask(self, N: int): # -> array | None:
|
||||
...
|
||||
|
||||
class MambaCache(ArraysCache):
|
||||
def __init__(self, left_padding: Optional[List[int]] = ...) -> None: ...
|
||||
@@ -238,7 +239,8 @@ class ChunkedKVCache(KVCache):
|
||||
...
|
||||
def update_and_fetch(self, keys, values): # -> tuple[array, array]:
|
||||
...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def trim(self, n): # -> int:
|
||||
...
|
||||
@property
|
||||
def meta_state(self): # -> tuple[str, ...]:
|
||||
...
|
||||
@@ -251,9 +253,10 @@ class CacheList(_BaseCache):
|
||||
def __getitem__(self, idx): ...
|
||||
def is_trimmable(self): # -> bool:
|
||||
...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def trim(self, n): ...
|
||||
@property
|
||||
def state(self) -> list[tuple[mx.array | None, mx.array | None]]: ...
|
||||
def state(self): # -> list[Any]:
|
||||
...
|
||||
@state.setter
|
||||
def state(self, v): # -> None:
|
||||
...
|
||||
@@ -268,14 +271,29 @@ class CacheList(_BaseCache):
|
||||
"""
|
||||
|
||||
class BatchKVCache(_BaseCache):
|
||||
step: int
|
||||
keys: array | None
|
||||
values: array | None
|
||||
offset: array
|
||||
left_padding: array
|
||||
_idx: int
|
||||
def __init__(self, left_padding: List[int]) -> None: ...
|
||||
def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
|
||||
step = ...
|
||||
def __init__(self, left_padding: List[int]) -> None:
|
||||
"""
|
||||
The BatchKV cache expects inputs to be left-padded.
|
||||
|
||||
E.g. the following prompts:
|
||||
|
||||
[1, 3, 5]
|
||||
[7]
|
||||
[2, 6, 8, 9]
|
||||
|
||||
Should be padded like so:
|
||||
|
||||
[0, 1, 3, 5]
|
||||
[0, 0, 0, 7]
|
||||
[2, 6, 8, 9]
|
||||
|
||||
And ``left_padding`` specifies the amount of padding for each.
|
||||
In this case, ``left_padding = [1, 3, 0]``.
|
||||
"""
|
||||
|
||||
def update_and_fetch(self, keys, values): # -> tuple[array | Any, array | Any]:
|
||||
...
|
||||
@property
|
||||
def state(
|
||||
self,
|
||||
@@ -301,21 +319,12 @@ class BatchKVCache(_BaseCache):
|
||||
"""
|
||||
|
||||
class BatchRotatingKVCache(_BaseCache):
|
||||
step: int
|
||||
keys: array | None
|
||||
values: array | None
|
||||
offset: array
|
||||
left_padding: array
|
||||
max_size: int
|
||||
_idx: int
|
||||
_offset: int
|
||||
rotated: bool
|
||||
_lengths: array | None
|
||||
def __init__(self, max_size: int, left_padding: List[int]) -> None: ...
|
||||
def _trim(self, trim_size: int, v: array, append: array | None = ...) -> array: ...
|
||||
def _update_in_place(self, keys: array, values: array) -> tuple[array, array]: ...
|
||||
def _update_concat(self, keys: array, values: array) -> tuple[array, array]: ...
|
||||
def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
|
||||
step = ...
|
||||
def __init__(self, max_size, left_padding: List[int]) -> None: ...
|
||||
def update_and_fetch(
|
||||
self, keys, values
|
||||
): # -> tuple[array | Any, array | Any] | tuple[array | Any, array | Any | None]:
|
||||
...
|
||||
@property
|
||||
def state(
|
||||
self,
|
||||
|
||||
@@ -1,280 +0,0 @@
|
||||
"""Type stubs for mlx_lm.models.deepseek_v4"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs
|
||||
from .cache import ArraysCache, RotatingKVCache
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
vocab_size: int
|
||||
hidden_size: int
|
||||
intermediate_size: int
|
||||
moe_intermediate_size: int
|
||||
num_hidden_layers: int
|
||||
num_attention_heads: int
|
||||
num_key_value_heads: int
|
||||
n_shared_experts: Optional[int]
|
||||
n_routed_experts: int
|
||||
num_experts_per_tok: int
|
||||
head_dim: int
|
||||
qk_rope_head_dim: int
|
||||
q_lora_rank: int
|
||||
o_lora_rank: int
|
||||
o_groups: int
|
||||
sliding_window: int
|
||||
hc_mult: int
|
||||
hc_sinkhorn_iters: int
|
||||
hc_eps: float
|
||||
compress_ratios: Optional[List[int]]
|
||||
compress_rope_theta: float
|
||||
rope_theta: float
|
||||
rope_scaling: Optional[Dict[str, Any]]
|
||||
rms_norm_eps: float
|
||||
swiglu_limit: float
|
||||
attention_bias: bool
|
||||
max_position_embeddings: int
|
||||
|
||||
class DeepseekV4RoPE(nn.Module):
|
||||
dims: int
|
||||
freqs: mx.array
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dims: int,
|
||||
base: float,
|
||||
scaling_config: Optional[Dict[str, Any]] = None,
|
||||
) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
offset: int = 0,
|
||||
inverse: bool = False,
|
||||
) -> mx.array: ...
|
||||
|
||||
class HyperConnection(nn.Module):
|
||||
dim: int
|
||||
hc_mult: int
|
||||
norm_eps: float
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
hc_mult: int,
|
||||
norm_eps: float,
|
||||
sinkhorn_iters: int,
|
||||
hc_eps: float,
|
||||
) -> None: ...
|
||||
|
||||
class HyperHead(nn.Module):
|
||||
dim: int
|
||||
hc_mult: int
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
hc_mult: int,
|
||||
norm_eps: float,
|
||||
hc_eps: float,
|
||||
) -> None: ...
|
||||
def __call__(self, x: mx.array) -> mx.array: ...
|
||||
|
||||
class Compressor(nn.Module):
|
||||
dim: int
|
||||
head_dim: int
|
||||
rope_head_dim: int
|
||||
compress_ratio: int
|
||||
overlap: bool
|
||||
wkv_gate: nn.Linear
|
||||
ape: mx.array
|
||||
norm: nn.RMSNorm
|
||||
rope: DeepseekV4RoPE
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
compress_ratio: int,
|
||||
head_dim: int,
|
||||
rope_head_dim: int,
|
||||
rms_norm_eps: float,
|
||||
rope: DeepseekV4RoPE,
|
||||
) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
cache: "DeepseekV4Cache",
|
||||
offset: Any,
|
||||
key: str = ...,
|
||||
) -> mx.array: ...
|
||||
|
||||
class Indexer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
args: ModelArgs,
|
||||
compress_ratio: int,
|
||||
rope: DeepseekV4RoPE,
|
||||
) -> None: ...
|
||||
|
||||
class _CompressorBranch:
|
||||
buffer_kv: Optional[mx.array]
|
||||
buffer_gate: Optional[mx.array]
|
||||
prev_kv: Optional[mx.array]
|
||||
prev_gate: Optional[mx.array]
|
||||
pool: Optional[mx.array]
|
||||
buffer_lengths: Optional[List[int]]
|
||||
pool_lengths: Optional[List[int]]
|
||||
buffer_count: int
|
||||
_new_pool_lengths: Optional[List[int]]
|
||||
|
||||
def __init__(self) -> None: ...
|
||||
|
||||
class DeepseekV4Cache:
|
||||
local: RotatingKVCache
|
||||
offset: int
|
||||
keys: Optional[mx.array]
|
||||
values: Optional[mx.array]
|
||||
state: Any
|
||||
meta_state: Any
|
||||
nbytes: int
|
||||
_branches: Dict[str, _CompressorBranch]
|
||||
_pending_lengths: Optional[List[int]]
|
||||
|
||||
def __init__(self, sliding_window: int) -> None: ...
|
||||
def update_and_fetch(
|
||||
self, keys: mx.array, values: mx.array
|
||||
) -> tuple[mx.array, mx.array]: ...
|
||||
def is_trimmable(self) -> bool: ...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def empty(self) -> bool: ...
|
||||
def size(self) -> int: ...
|
||||
def prepare(
|
||||
self,
|
||||
*,
|
||||
left_padding: Optional[List[int]] = None,
|
||||
lengths: Optional[List[int]] = None,
|
||||
right_padding: Optional[List[int]] = None,
|
||||
) -> None: ...
|
||||
def finalize(self) -> None: ...
|
||||
def filter(self, batch_indices: mx.array) -> None: ...
|
||||
def extend(self, other: "DeepseekV4Cache") -> None: ...
|
||||
def extract(self, idx: int) -> "DeepseekV4Cache": ...
|
||||
@classmethod
|
||||
def merge(cls, caches: List["DeepseekV4Cache"]) -> "DeepseekV4Cache": ...
|
||||
|
||||
class V4Attention(nn.Module):
|
||||
args: ModelArgs
|
||||
layer_id: int
|
||||
dim: int
|
||||
n_heads: int
|
||||
head_dim: int
|
||||
rope_head_dim: int
|
||||
nope_head_dim: int
|
||||
n_groups: int
|
||||
q_lora_rank: int
|
||||
o_lora_rank: int
|
||||
window: int
|
||||
eps: float
|
||||
scale: float
|
||||
compress_ratio: int
|
||||
wqkv_a: nn.Linear
|
||||
q_norm: nn.RMSNorm
|
||||
wq_b: nn.Linear
|
||||
kv_norm: nn.RMSNorm
|
||||
attn_sink: mx.array
|
||||
wo_a: nn.Linear
|
||||
wo_b: nn.Linear
|
||||
rope: DeepseekV4RoPE
|
||||
compressor: Compressor
|
||||
indexer: Indexer
|
||||
|
||||
def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class DeepseekV4MLP(nn.Module):
|
||||
gate_proj: nn.Linear
|
||||
up_proj: nn.Linear
|
||||
down_proj: nn.Linear
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
intermediate_size: int,
|
||||
swiglu_limit: float = 0.0,
|
||||
) -> None: ...
|
||||
def __call__(self, x: mx.array) -> mx.array: ...
|
||||
|
||||
class MoEGate(nn.Module):
|
||||
weight: mx.array
|
||||
|
||||
def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
|
||||
def __call__(
|
||||
self, x: mx.array, input_ids: mx.array
|
||||
) -> tuple[mx.array, mx.array]: ...
|
||||
|
||||
class DeepseekV4MoE(nn.Module):
|
||||
num_experts_per_tok: int
|
||||
switch_mlp: SwitchGLU
|
||||
gate: MoEGate
|
||||
shared_experts: DeepseekV4MLP
|
||||
|
||||
def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
|
||||
def __call__(self, x: mx.array, input_ids: mx.array) -> mx.array: ...
|
||||
|
||||
class DeepseekV4Block(nn.Module):
|
||||
attn_norm: nn.RMSNorm
|
||||
attn: V4Attention
|
||||
hc_attn: HyperConnection
|
||||
ffn_norm: nn.RMSNorm
|
||||
ffn: DeepseekV4MoE
|
||||
hc_ffn: HyperConnection
|
||||
|
||||
def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
h: mx.array,
|
||||
cache: Optional[Any],
|
||||
input_ids: mx.array,
|
||||
) -> mx.array: ...
|
||||
|
||||
class DeepseekV4Model(nn.Module):
|
||||
args: ModelArgs
|
||||
vocab_size: int
|
||||
embed_tokens: nn.Embedding
|
||||
layers: list[DeepseekV4Block]
|
||||
norm: nn.RMSNorm
|
||||
hc_head: HyperHead
|
||||
|
||||
def __init__(self, args: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[List[Any]] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class Model(nn.Module):
|
||||
args: ModelArgs
|
||||
model_type: str
|
||||
model: DeepseekV4Model
|
||||
lm_head: nn.Linear
|
||||
|
||||
def __init__(self, args: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[List[Any]] = None,
|
||||
) -> mx.array: ...
|
||||
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
|
||||
def make_cache(self) -> list[RotatingKVCache | DeepseekV4Cache]: ...
|
||||
@property
|
||||
def layers(self) -> list[DeepseekV4Block]: ...
|
||||
@@ -1,35 +0,0 @@
|
||||
from typing import Optional
|
||||
|
||||
import mlx.core as mx
|
||||
|
||||
def compute_g(A_log: mx.array, a: mx.array, dt_bias: mx.array) -> mx.array: ...
|
||||
def gated_delta_update(
|
||||
q: mx.array,
|
||||
k: mx.array,
|
||||
v: mx.array,
|
||||
a: mx.array,
|
||||
b: mx.array,
|
||||
A_log: mx.array,
|
||||
dt_bias: mx.array,
|
||||
state: Optional[mx.array] = ...,
|
||||
mask: Optional[mx.array] = ...,
|
||||
use_kernel: bool = ...,
|
||||
) -> tuple[mx.array, mx.array]: ...
|
||||
def gated_delta_ops(
|
||||
q: mx.array,
|
||||
k: mx.array,
|
||||
v: mx.array,
|
||||
g: mx.array,
|
||||
beta: mx.array,
|
||||
state: Optional[mx.array] = ...,
|
||||
mask: Optional[mx.array] = ...,
|
||||
) -> tuple[mx.array, mx.array]: ...
|
||||
def gated_delta_kernel(
|
||||
q: mx.array,
|
||||
k: mx.array,
|
||||
v: mx.array,
|
||||
g: mx.array,
|
||||
beta: mx.array,
|
||||
state: mx.array,
|
||||
mask: Optional[mx.array] = ...,
|
||||
) -> tuple[mx.array, mx.array]: ...
|
||||
@@ -1,31 +0,0 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from . import gemma4_text
|
||||
from .base import BaseModelArgs
|
||||
from .cache import KVCache, RotatingKVCache
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
text_config: Optional[dict[str, Any]]
|
||||
vocab_size: int
|
||||
|
||||
def __post_init__(self) -> None: ...
|
||||
|
||||
class Model(nn.Module):
|
||||
args: ModelArgs
|
||||
model_type: str
|
||||
language_model: gemma4_text.Model
|
||||
|
||||
def __init__(self, args: ModelArgs) -> None: ...
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
|
||||
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
|
||||
@property
|
||||
def layers(self) -> list[gemma4_text.DecoderLayer]: ...
|
||||
@property
|
||||
def quant_predicate(self) -> Any: ...
|
||||
def make_cache(self) -> list[KVCache | RotatingKVCache]: ...
|
||||
@@ -1,179 +0,0 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs
|
||||
from .cache import KVCache, RotatingKVCache
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
num_hidden_layers: int
|
||||
intermediate_size: int
|
||||
num_attention_heads: int
|
||||
head_dim: int
|
||||
global_head_dim: int
|
||||
global_partial_rotary_factor: float
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
vocab_size_per_layer_input: int
|
||||
num_key_value_heads: int
|
||||
num_global_key_value_heads: Optional[int]
|
||||
num_kv_shared_layers: int
|
||||
pad_token_id: int
|
||||
hidden_size_per_layer_input: int
|
||||
rope_traditional: bool
|
||||
partial_rotary_factor: float
|
||||
rope_parameters: Optional[Dict[str, Any]]
|
||||
sliding_window: int
|
||||
sliding_window_pattern: int
|
||||
max_position_embeddings: int
|
||||
attention_k_eq_v: bool
|
||||
final_logit_softcapping: float
|
||||
use_double_wide_mlp: bool
|
||||
enable_moe_block: bool
|
||||
num_experts: Optional[int]
|
||||
top_k_experts: Optional[int]
|
||||
moe_intermediate_size: Optional[int]
|
||||
layer_types: Optional[List[str]]
|
||||
tie_word_embeddings: bool
|
||||
|
||||
def __post_init__(self) -> None: ...
|
||||
|
||||
class MLP(nn.Module):
|
||||
gate_proj: nn.Linear
|
||||
down_proj: nn.Linear
|
||||
up_proj: nn.Linear
|
||||
|
||||
def __init__(self, config: ModelArgs, layer_idx: int = 0) -> None: ...
|
||||
def __call__(self, x: mx.array) -> mx.array: ...
|
||||
|
||||
class Router(nn.Module):
|
||||
proj: nn.Linear
|
||||
scale: mx.array
|
||||
per_expert_scale: mx.array
|
||||
|
||||
def __init__(self, config: ModelArgs) -> None: ...
|
||||
def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
|
||||
|
||||
class Experts(nn.Module):
|
||||
switch_glu: SwitchGLU
|
||||
|
||||
def __init__(self, config: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self, x: mx.array, top_k_indices: mx.array, top_k_weights: mx.array
|
||||
) -> mx.array: ...
|
||||
|
||||
class Attention(nn.Module):
|
||||
layer_idx: int
|
||||
layer_type: str
|
||||
is_sliding: bool
|
||||
head_dim: int
|
||||
n_heads: int
|
||||
n_kv_heads: int
|
||||
use_k_eq_v: bool
|
||||
scale: float
|
||||
q_proj: nn.Linear
|
||||
k_proj: nn.Linear
|
||||
v_proj: nn.Linear
|
||||
o_proj: nn.Linear
|
||||
q_norm: nn.Module
|
||||
k_norm: nn.Module
|
||||
v_norm: nn.Module
|
||||
rope: nn.Module
|
||||
|
||||
def __init__(self, config: ModelArgs, layer_idx: int) -> None: ...
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
layer_idx: int
|
||||
layer_type: str
|
||||
self_attn: Attention
|
||||
mlp: MLP
|
||||
enable_moe: bool
|
||||
router: Router
|
||||
experts: Experts
|
||||
input_layernorm: nn.Module
|
||||
post_attention_layernorm: nn.Module
|
||||
pre_feedforward_layernorm: nn.Module
|
||||
post_feedforward_layernorm: nn.Module
|
||||
post_feedforward_layernorm_1: nn.Module
|
||||
post_feedforward_layernorm_2: nn.Module
|
||||
pre_feedforward_layernorm_2: nn.Module
|
||||
hidden_size_per_layer_input: int
|
||||
per_layer_input_gate: Optional[nn.Linear]
|
||||
per_layer_projection: Optional[nn.Linear]
|
||||
post_per_layer_input_norm: Optional[nn.Module]
|
||||
layer_scalar: mx.array
|
||||
|
||||
def __init__(self, config: ModelArgs, layer_idx: int) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = ...,
|
||||
cache: Optional[Any] = ...,
|
||||
per_layer_input: Optional[mx.array] = ...,
|
||||
shared_kv: Optional[tuple[mx.array, mx.array]] = ...,
|
||||
offset: Optional[mx.array] = ...,
|
||||
) -> tuple[mx.array, tuple[mx.array, mx.array], mx.array]: ...
|
||||
|
||||
class Gemma4TextModel(nn.Module):
|
||||
config: ModelArgs
|
||||
vocab_size: int
|
||||
window_size: int
|
||||
sliding_window_pattern: int
|
||||
num_hidden_layers: int
|
||||
embed_tokens: nn.Embedding
|
||||
embed_scale: float
|
||||
layers: list[DecoderLayer]
|
||||
norm: nn.Module
|
||||
hidden_size_per_layer_input: int
|
||||
embed_tokens_per_layer: Optional[nn.Embedding]
|
||||
per_layer_model_projection: Optional[nn.Linear]
|
||||
per_layer_projection_norm: Optional[nn.Module]
|
||||
previous_kvs: list[int]
|
||||
|
||||
def __init__(self, config: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: Optional[mx.array] = ...,
|
||||
cache: Optional[list[Any]] = ...,
|
||||
input_embeddings: Optional[mx.array] = ...,
|
||||
per_layer_inputs: Optional[mx.array] = ...,
|
||||
) -> mx.array: ...
|
||||
def _get_per_layer_inputs(
|
||||
self,
|
||||
input_ids: Optional[mx.array],
|
||||
input_embeddings: Optional[mx.array] = ...,
|
||||
) -> mx.array: ...
|
||||
def _project_per_layer_inputs(
|
||||
self,
|
||||
input_embeddings: mx.array,
|
||||
per_layer_inputs: Optional[mx.array] = ...,
|
||||
) -> mx.array: ...
|
||||
def _make_masks(self, h: mx.array, cache: list[Any]) -> list[Any]: ...
|
||||
|
||||
class Model(nn.Module):
|
||||
args: ModelArgs
|
||||
model_type: str
|
||||
model: Gemma4TextModel
|
||||
final_logit_softcapping: float
|
||||
tie_word_embeddings: bool
|
||||
lm_head: nn.Linear
|
||||
|
||||
def __init__(self, args: ModelArgs) -> None: ...
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
|
||||
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
|
||||
@property
|
||||
def layers(self) -> list[DecoderLayer]: ...
|
||||
@property
|
||||
def head_dim(self) -> int: ...
|
||||
@property
|
||||
def n_kv_heads(self) -> int: ...
|
||||
@property
|
||||
def quant_predicate(self) -> Any: ...
|
||||
def make_cache(self) -> list[KVCache | RotatingKVCache]: ...
|
||||
@@ -1,46 +0,0 @@
|
||||
"""Type stubs for mlx_lm.models.glm_moe_dsa"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from .base import BaseModelArgs
|
||||
from .deepseek_v32 import Model as DSV32Model
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
vocab_size: int
|
||||
hidden_size: int
|
||||
index_head_dim: int
|
||||
index_n_heads: int
|
||||
index_topk: int
|
||||
intermediate_size: int
|
||||
moe_intermediate_size: int
|
||||
num_hidden_layers: int
|
||||
num_attention_heads: int
|
||||
num_key_value_heads: int
|
||||
n_shared_experts: Optional[int]
|
||||
n_routed_experts: Optional[int]
|
||||
routed_scaling_factor: float
|
||||
kv_lora_rank: int
|
||||
q_lora_rank: int
|
||||
qk_rope_head_dim: int
|
||||
v_head_dim: int
|
||||
qk_nope_head_dim: int
|
||||
topk_method: str
|
||||
scoring_func: str
|
||||
norm_topk_prob: bool
|
||||
n_group: int
|
||||
topk_group: int
|
||||
num_experts_per_tok: int
|
||||
moe_layer_freq: int
|
||||
first_k_dense_replace: int
|
||||
max_position_embeddings: int
|
||||
rms_norm_eps: float
|
||||
rope_parameters: Dict[str, Any]
|
||||
attention_bias: bool
|
||||
rope_scaling: Dict[str, Any] | None
|
||||
rope_theta: float | None
|
||||
|
||||
class Model(DSV32Model):
|
||||
def __init__(self, config: ModelArgs) -> None: ...
|
||||
@@ -1,103 +0,0 @@
|
||||
"""Type stubs for mlx_lm.models.gpt_oss"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, List, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs
|
||||
from .cache import KVCache
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
intermediate_size: int
|
||||
num_hidden_layers: int
|
||||
num_attention_heads: int
|
||||
num_key_value_heads: int
|
||||
num_local_experts: int
|
||||
num_experts_per_tok: int
|
||||
vocab_size: int
|
||||
rms_norm_eps: float
|
||||
sliding_window: int
|
||||
layer_types: Optional[List[str]]
|
||||
|
||||
def mlx_topk(a: mx.array, k: int, axis: int = -1) -> tuple[mx.array, mx.array]: ...
|
||||
|
||||
class AttentionBlock(nn.Module):
|
||||
head_dim: int
|
||||
num_attention_heads: int
|
||||
num_key_value_heads: int
|
||||
num_key_value_groups: int
|
||||
sinks: mx.array
|
||||
q_proj: nn.Linear
|
||||
k_proj: nn.Linear
|
||||
v_proj: nn.Linear
|
||||
o_proj: nn.Linear
|
||||
sm_scale: float
|
||||
rope: nn.Module
|
||||
|
||||
def __init__(self, config: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
self_attn: AttentionBlock
|
||||
mlp: MLPBlock
|
||||
|
||||
def __init__(self, config: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class MLPBlock(nn.Module):
|
||||
hidden_size: int
|
||||
num_local_experts: int
|
||||
num_experts_per_tok: int
|
||||
experts: SwitchGLU
|
||||
router: nn.Linear
|
||||
sharding_group: Optional[mx.distributed.Group]
|
||||
|
||||
def __init__(self, config: ModelArgs) -> None: ...
|
||||
def __call__(self, x: mx.array) -> mx.array: ...
|
||||
|
||||
class GptOssMoeModel(nn.Module):
|
||||
embed_tokens: nn.Embedding
|
||||
norm: nn.RMSNorm
|
||||
layer_types: List[str]
|
||||
layers: list[TransformerBlock]
|
||||
window_size: int
|
||||
swa_idx: int
|
||||
ga_idx: int
|
||||
|
||||
def __init__(self, args: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class Model(nn.Module):
|
||||
model_type: str
|
||||
model: GptOssMoeModel
|
||||
lm_head: nn.Linear
|
||||
|
||||
def __init__(self, args: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array: ...
|
||||
@property
|
||||
def layers(self) -> list[nn.Module]: ...
|
||||
def make_cache(self) -> list[KVCache]: ...
|
||||
@@ -1,94 +0,0 @@
|
||||
"""Type stubs for mlx_lm.models.minimax"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .base import BaseModelArgs
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
@dataclass
|
||||
class ModelArgs(BaseModelArgs):
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
intermediate_size: int
|
||||
num_hidden_layers: int
|
||||
num_attention_heads: int
|
||||
num_key_value_heads: int
|
||||
num_local_experts: int
|
||||
num_experts_per_tok: int
|
||||
max_position_embeddings: int
|
||||
|
||||
class MiniMaxAttention(nn.Module):
|
||||
num_heads: int
|
||||
num_attention_heads: int
|
||||
num_key_value_heads: int
|
||||
head_dim: int
|
||||
scale: float
|
||||
q_proj: nn.Linear
|
||||
k_proj: nn.Linear
|
||||
v_proj: nn.Linear
|
||||
o_proj: nn.Linear
|
||||
q_norm: nn.Module
|
||||
k_norm: nn.Module
|
||||
rope: nn.Module
|
||||
|
||||
def __init__(self, args: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class MiniMaxSparseMoeBlock(nn.Module):
|
||||
num_experts_per_tok: int
|
||||
gate: nn.Linear
|
||||
switch_mlp: SwitchGLU
|
||||
e_score_correction_bias: mx.array
|
||||
sharding_group: Optional[mx.distributed.Group]
|
||||
|
||||
def __init__(self, args: ModelArgs) -> None: ...
|
||||
def __call__(self, x: mx.array) -> mx.array: ...
|
||||
|
||||
class MiniMaxDecoderLayer(nn.Module):
|
||||
self_attn: MiniMaxAttention
|
||||
block_sparse_moe: MiniMaxSparseMoeBlock
|
||||
input_layernorm: nn.RMSNorm
|
||||
post_attention_layernorm: nn.RMSNorm
|
||||
|
||||
def __init__(self, args: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class MiniMaxModel(nn.Module):
|
||||
embed_tokens: nn.Embedding
|
||||
layers: list[MiniMaxDecoderLayer]
|
||||
norm: nn.RMSNorm
|
||||
|
||||
def __init__(self, args: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class Model(nn.Module):
|
||||
model_type: str
|
||||
model: MiniMaxModel
|
||||
lm_head: nn.Linear
|
||||
|
||||
def __init__(self, args: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array: ...
|
||||
@property
|
||||
def layers(self) -> list[MiniMaxDecoderLayer]: ...
|
||||
@@ -1,168 +0,0 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, List, Optional, Tuple
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .cache import ArraysCache, KVCache
|
||||
from .switch_layers import SwitchMLP
|
||||
|
||||
@dataclass
|
||||
class ModelArgs:
|
||||
model_type: str
|
||||
vocab_size: int
|
||||
hidden_size: int
|
||||
intermediate_size: int
|
||||
num_hidden_layers: int
|
||||
max_position_embeddings: int
|
||||
num_attention_heads: int
|
||||
num_key_value_heads: int
|
||||
attention_bias: bool
|
||||
mamba_num_heads: int
|
||||
mamba_head_dim: int
|
||||
mamba_proj_bias: bool
|
||||
ssm_state_size: int
|
||||
conv_kernel: int
|
||||
n_groups: int
|
||||
mlp_bias: bool
|
||||
layer_norm_epsilon: float
|
||||
use_bias: bool
|
||||
use_conv_bias: bool
|
||||
hybrid_override_pattern: List[str]
|
||||
head_dim: Optional[int]
|
||||
moe_intermediate_size: Optional[int]
|
||||
moe_shared_expert_intermediate_size: Optional[int]
|
||||
n_group: Optional[int]
|
||||
n_routed_experts: Optional[int]
|
||||
n_shared_experts: Optional[int]
|
||||
topk_group: Optional[int]
|
||||
num_experts_per_tok: Optional[int]
|
||||
norm_topk_prob: Optional[bool]
|
||||
routed_scaling_factor: Optional[float]
|
||||
time_step_limit: Optional[Tuple[float, float]]
|
||||
time_step_min: Optional[float]
|
||||
time_step_max: Optional[float]
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
|
||||
def __post_init__(self) -> None: ...
|
||||
|
||||
class NemotronHMamba2Mixer(nn.Module):
|
||||
num_heads: int
|
||||
hidden_size: int
|
||||
ssm_state_size: int
|
||||
conv_kernel_size: int
|
||||
intermediate_size: int
|
||||
n_groups: int
|
||||
head_dim: int
|
||||
conv_dim: int
|
||||
conv1d: nn.Conv1d
|
||||
in_proj: nn.Linear
|
||||
dt_bias: mx.array
|
||||
A_log: mx.array
|
||||
D: mx.array
|
||||
norm: nn.RMSNorm
|
||||
heads_per_group: int
|
||||
out_proj: nn.Linear
|
||||
|
||||
def __init__(self, args: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
hidden_states: mx.array,
|
||||
mask: Optional[mx.array],
|
||||
cache: Optional[ArraysCache] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class NemotronHAttention(nn.Module):
|
||||
hidden_size: int
|
||||
num_heads: int
|
||||
head_dim: int
|
||||
num_key_value_heads: int
|
||||
scale: float
|
||||
q_proj: nn.Linear
|
||||
k_proj: nn.Linear
|
||||
v_proj: nn.Linear
|
||||
o_proj: nn.Linear
|
||||
|
||||
def __init__(self, args: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[KVCache] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class MoEGate(nn.Module):
|
||||
config: ModelArgs
|
||||
top_k: int
|
||||
norm_topk_prob: bool
|
||||
weight: mx.array
|
||||
|
||||
def __init__(self, config: ModelArgs) -> None: ...
|
||||
def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
|
||||
|
||||
class NemotronHMLP(nn.Module):
|
||||
up_proj: nn.Linear
|
||||
down_proj: nn.Linear
|
||||
|
||||
def __init__(
|
||||
self, args: ModelArgs, intermediate_size: Optional[int] = None
|
||||
) -> None: ...
|
||||
def __call__(self, x: mx.array) -> mx.array: ...
|
||||
|
||||
class NemotronHMoE(nn.Module):
|
||||
config: ModelArgs
|
||||
num_experts_per_tok: int
|
||||
moe_latent_size: Optional[int]
|
||||
switch_mlp: SwitchMLP
|
||||
gate: MoEGate
|
||||
shared_experts: NemotronHMLP
|
||||
fc1_latent_proj: nn.Linear
|
||||
fc2_latent_proj: nn.Linear
|
||||
|
||||
def __init__(self, config: ModelArgs) -> None: ...
|
||||
def __call__(self, x: mx.array) -> mx.array: ...
|
||||
|
||||
class NemotronHBlock(nn.Module):
|
||||
block_type: str
|
||||
norm: nn.RMSNorm
|
||||
mixer: NemotronHMamba2Mixer | NemotronHAttention | NemotronHMLP | NemotronHMoE
|
||||
|
||||
def __init__(self, args: ModelArgs, block_type: str) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class NemotronHModel(nn.Module):
|
||||
embeddings: nn.Embedding
|
||||
layers: list[NemotronHBlock]
|
||||
norm_f: nn.RMSNorm
|
||||
fa_idx: int
|
||||
ssm_idx: int
|
||||
|
||||
def __init__(self, args: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class Model(nn.Module):
|
||||
args: ModelArgs
|
||||
backbone: NemotronHModel
|
||||
lm_head: nn.Linear
|
||||
model_type: str
|
||||
|
||||
def __init__(self, args: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array: ...
|
||||
@property
|
||||
def layers(self) -> list[NemotronHBlock]: ...
|
||||
def make_cache(self) -> list[ArraysCache | KVCache]: ...
|
||||
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
|
||||
@@ -1,153 +0,0 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .cache import ArraysCache, KVCache
|
||||
from .qwen3_next import (
|
||||
Qwen3NextAttention as Attention,
|
||||
Qwen3NextMLP as MLP,
|
||||
Qwen3NextRMSNormGated as RMSNormGated,
|
||||
Qwen3NextSparseMoeBlock,
|
||||
)
|
||||
|
||||
SparseMoeBlock = Qwen3NextSparseMoeBlock
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
@dataclass
|
||||
class TextModelArgs:
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
intermediate_size: int
|
||||
num_hidden_layers: int
|
||||
num_attention_heads: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
num_key_value_heads: int
|
||||
max_position_embeddings: int
|
||||
linear_num_value_heads: int
|
||||
linear_num_key_heads: int
|
||||
linear_key_head_dim: int
|
||||
linear_value_head_dim: int
|
||||
linear_conv_kernel_dim: int
|
||||
tie_word_embeddings: bool
|
||||
attention_bias: bool
|
||||
head_dim: Optional[int]
|
||||
full_attention_interval: int
|
||||
num_experts: int
|
||||
num_experts_per_tok: int
|
||||
decoder_sparse_step: int
|
||||
shared_expert_intermediate_size: int
|
||||
moe_intermediate_size: int
|
||||
norm_topk_prob: bool
|
||||
rope_parameters: Optional[dict[str, Any]]
|
||||
partial_rotary_factor: float
|
||||
rope_theta: float
|
||||
rope_scaling: Optional[dict[str, Any]]
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, params: dict[str, Any]) -> TextModelArgs: ...
|
||||
def __post_init__(self) -> None: ...
|
||||
|
||||
class GatedDeltaNet(nn.Module):
|
||||
hidden_size: int
|
||||
num_v_heads: int
|
||||
num_k_heads: int
|
||||
head_k_dim: int
|
||||
head_v_dim: int
|
||||
key_dim: int
|
||||
value_dim: int
|
||||
conv_kernel_size: int
|
||||
conv_dim: int
|
||||
conv1d: nn.Conv1d
|
||||
in_proj_qkv: nn.Linear
|
||||
in_proj_z: nn.Linear
|
||||
in_proj_b: nn.Linear
|
||||
in_proj_a: nn.Linear
|
||||
dt_bias: mx.array
|
||||
A_log: mx.array
|
||||
norm: RMSNormGated
|
||||
out_proj: nn.Linear
|
||||
|
||||
def __init__(self, config: TextModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
is_linear: bool
|
||||
linear_attn: GatedDeltaNet
|
||||
self_attn: Attention
|
||||
input_layernorm: nn.RMSNorm
|
||||
post_attention_layernorm: nn.RMSNorm
|
||||
mlp: MLP | SparseMoeBlock
|
||||
|
||||
def __init__(self, args: TextModelArgs, layer_idx: int) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class Qwen3_5TextModel(nn.Module):
|
||||
embed_tokens: nn.Embedding
|
||||
layers: list[DecoderLayer]
|
||||
norm: nn.RMSNorm
|
||||
ssm_idx: int
|
||||
fa_idx: int
|
||||
|
||||
def __init__(self, args: TextModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class TextModel(nn.Module):
|
||||
args: TextModelArgs
|
||||
model_type: str
|
||||
model: Qwen3_5TextModel
|
||||
lm_head: nn.Linear
|
||||
|
||||
def __init__(self, args: TextModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
) -> mx.array: ...
|
||||
@property
|
||||
def layers(self) -> list[DecoderLayer]: ...
|
||||
def make_cache(self) -> list[ArraysCache | KVCache]: ...
|
||||
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
|
||||
|
||||
@dataclass
|
||||
class ModelArgs:
|
||||
model_type: str
|
||||
text_config: dict[str, Any]
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
|
||||
|
||||
class Model(nn.Module):
|
||||
args: ModelArgs
|
||||
model_type: str
|
||||
language_model: TextModel
|
||||
|
||||
def __init__(self, args: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
) -> mx.array: ...
|
||||
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
|
||||
@property
|
||||
def layers(self) -> list[DecoderLayer]: ...
|
||||
def make_cache(self) -> list[ArraysCache | KVCache]: ...
|
||||
@@ -1,19 +0,0 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .cache import ArraysCache, KVCache
|
||||
from .qwen3_5 import DecoderLayer, Model as Qwen3_5Model, TextModel
|
||||
|
||||
@dataclass
|
||||
class ModelArgs:
|
||||
model_type: str
|
||||
text_config: dict[str, Any]
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
|
||||
|
||||
class Model(Qwen3_5Model):
|
||||
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
|
||||
@@ -5,18 +5,8 @@ from typing import Any, Optional
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .cache import ArraysCache, KVCache
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
class Qwen3NextRMSNormGated(nn.Module):
|
||||
eps: float
|
||||
weight: mx.array
|
||||
|
||||
def __init__(self, hidden_size: int, eps: float = ...) -> None: ...
|
||||
def __call__(
|
||||
self, hidden_states: mx.array, gate: mx.array | None = None
|
||||
) -> mx.array: ...
|
||||
|
||||
class Qwen3NextMLP(nn.Module):
|
||||
gate_proj: nn.Linear
|
||||
down_proj: nn.Linear
|
||||
@@ -71,7 +61,6 @@ class Qwen3NextAttention(nn.Module):
|
||||
class Qwen3NextSparseMoeBlock(nn.Module):
|
||||
norm_topk_prob: bool
|
||||
num_experts: int
|
||||
num_experts_per_tok: int
|
||||
top_k: int
|
||||
gate: nn.Linear
|
||||
switch_mlp: SwitchGLU
|
||||
@@ -101,8 +90,6 @@ class Qwen3NextModel(nn.Module):
|
||||
embed_tokens: nn.Embedding
|
||||
layers: list[Qwen3NextDecoderLayer]
|
||||
norm: nn.RMSNorm
|
||||
ssm_idx: int
|
||||
fa_idx: int
|
||||
|
||||
def __init__(self, args: Any) -> None: ...
|
||||
def __call__(
|
||||
@@ -125,4 +112,3 @@ class Model(nn.Module):
|
||||
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
|
||||
@property
|
||||
def layers(self) -> list[Qwen3NextDecoderLayer]: ...
|
||||
def make_cache(self) -> list[ArraysCache | KVCache]: ...
|
||||
@@ -1,51 +0,0 @@
|
||||
from typing import Any, Optional
|
||||
|
||||
import mlx.nn as nn
|
||||
|
||||
class YarnRoPE(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dims: int,
|
||||
traditional: bool = ...,
|
||||
max_position_embeddings: int = ...,
|
||||
base: float = ...,
|
||||
scaling_factor: float = ...,
|
||||
original_max_position_embeddings: int = ...,
|
||||
beta_fast: float = ...,
|
||||
beta_slow: float = ...,
|
||||
mscale: float = ...,
|
||||
mscale_all_dim: float = ...,
|
||||
) -> None: ...
|
||||
|
||||
class Llama3RoPE(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dims: int,
|
||||
traditional: bool = ...,
|
||||
max_position_embeddings: int = ...,
|
||||
base: float = ...,
|
||||
scaling_factor: float = ...,
|
||||
original_max_position_embeddings: int = ...,
|
||||
low_freq_factor: float = ...,
|
||||
high_freq_factor: float = ...,
|
||||
) -> None: ...
|
||||
|
||||
class SuScaledRoPE(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dims: int,
|
||||
traditional: bool = ...,
|
||||
max_position_embeddings: int = ...,
|
||||
base: float = ...,
|
||||
short_factor: Any = ...,
|
||||
long_factor: Any = ...,
|
||||
original_max_position_embeddings: int = ...,
|
||||
) -> None: ...
|
||||
|
||||
def initialize_rope(
|
||||
dims: int,
|
||||
base: float = ...,
|
||||
traditional: bool = ...,
|
||||
scaling_config: Optional[dict[str, Any]] = ...,
|
||||
max_position_embeddings: Optional[int] = ...,
|
||||
) -> nn.Module: ...
|
||||
@@ -73,9 +73,6 @@ class SwitchGLU(nn.Module):
|
||||
def __call__(self, x, indices) -> mx.array: ...
|
||||
|
||||
class SwitchMLP(nn.Module):
|
||||
fc1: SwitchLinear
|
||||
fc2: SwitchLinear
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_dims: int,
|
||||
|
||||
@@ -48,11 +48,7 @@ def make_logits_processors(
|
||||
logit_bias: Optional[Dict[int, float]] = ...,
|
||||
repetition_penalty: Optional[float] = ...,
|
||||
repetition_context_size: Optional[int] = ...,
|
||||
presence_penalty: Optional[float] = ...,
|
||||
presence_context_size: Optional[int] = ...,
|
||||
frequency_penalty: Optional[float] = ...,
|
||||
frequency_context_size: Optional[int] = ...,
|
||||
) -> list[Callable[[mx.array, mx.array], mx.array]]:
|
||||
): # -> list[Any]:
|
||||
"""
|
||||
Make logits processors for use with ``generate_step``.
|
||||
|
||||
|
||||
@@ -39,11 +39,11 @@ class StreamingDetokenizer:
|
||||
"""
|
||||
|
||||
__slots__ = ...
|
||||
def reset(self) -> None: ...
|
||||
def add_token(self, token: int) -> None: ...
|
||||
def finalize(self) -> None: ...
|
||||
def reset(self): ...
|
||||
def add_token(self, token): ...
|
||||
def finalize(self): ...
|
||||
@property
|
||||
def last_segment(self) -> str:
|
||||
def last_segment(self):
|
||||
"""Return the last segment of readable text since last time this property was accessed."""
|
||||
|
||||
class NaiveStreamingDetokenizer(StreamingDetokenizer):
|
||||
@@ -117,8 +117,6 @@ class TokenizerWrapper:
|
||||
think_end: str | None
|
||||
think_start_id: int | None
|
||||
think_end_id: int | None
|
||||
think_start_tokens: list[int] | None
|
||||
think_end_tokens: list[int] | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
||||
@@ -1,50 +0,0 @@
|
||||
"""
|
||||
This type stub file was generated by pyright.
|
||||
"""
|
||||
|
||||
import mlx.nn as nn
|
||||
|
||||
class DoRALinear(nn.Module):
|
||||
@staticmethod
|
||||
def from_base(
|
||||
linear: nn.Linear, r: int = ..., dropout: float = ..., scale: float = ...
|
||||
): # -> DoRALinear:
|
||||
...
|
||||
def fuse(self, dequantize: bool = ...): # -> QuantizedLinear | Linear:
|
||||
...
|
||||
def __init__(
|
||||
self,
|
||||
input_dims: int,
|
||||
output_dims: int,
|
||||
r: int = ...,
|
||||
dropout: float = ...,
|
||||
scale: float = ...,
|
||||
bias: bool = ...,
|
||||
) -> None: ...
|
||||
def set_linear(self, linear): # -> None:
|
||||
"""
|
||||
Set the self.linear layer and recompute self.m.
|
||||
"""
|
||||
...
|
||||
|
||||
def __call__(self, x): ...
|
||||
|
||||
class DoRAEmbedding(nn.Module):
|
||||
def from_base(
|
||||
embedding: nn.Embedding, r: int = ..., dropout: float = ..., scale: float = ...
|
||||
): # -> DoRAEmbedding:
|
||||
...
|
||||
def fuse(self, dequantize: bool = ...): # -> Embedding:
|
||||
...
|
||||
def __init__(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
dims: int,
|
||||
r: int = ...,
|
||||
dropout: float = ...,
|
||||
scale: float = ...,
|
||||
) -> None: ...
|
||||
def set_embedding(self, embedding: nn.Module): # -> None:
|
||||
...
|
||||
def __call__(self, x): ...
|
||||
def as_linear(self, x): ...
|
||||
@@ -1,66 +0,0 @@
|
||||
"""
|
||||
This type stub file was generated by pyright.
|
||||
"""
|
||||
|
||||
import mlx.nn as nn
|
||||
|
||||
class LoRALinear(nn.Module):
|
||||
@staticmethod
|
||||
def from_base(
|
||||
linear: nn.Linear, r: int = ..., dropout: float = ..., scale: float = ...
|
||||
): # -> LoRALinear:
|
||||
...
|
||||
def fuse(self, dequantize: bool = ...): # -> QuantizedLinear | Linear:
|
||||
...
|
||||
def __init__(
|
||||
self,
|
||||
input_dims: int,
|
||||
output_dims: int,
|
||||
r: int = ...,
|
||||
dropout: float = ...,
|
||||
scale: float = ...,
|
||||
bias: bool = ...,
|
||||
) -> None: ...
|
||||
def __call__(self, x): # -> array:
|
||||
...
|
||||
|
||||
class LoRASwitchLinear(nn.Module):
|
||||
@staticmethod
|
||||
def from_base(
|
||||
linear: nn.Module, r: int = ..., dropout: float = ..., scale: float = ...
|
||||
): # -> LoRASwitchLinear:
|
||||
...
|
||||
def fuse(self, dequantize: bool = ...): # -> QuantizedSwitchLinear | SwitchLinear:
|
||||
...
|
||||
def __init__(
|
||||
self,
|
||||
input_dims: int,
|
||||
output_dims: int,
|
||||
num_experts: int,
|
||||
r: int = ...,
|
||||
dropout: float = ...,
|
||||
scale: float = ...,
|
||||
bias: bool = ...,
|
||||
) -> None: ...
|
||||
def __call__(self, x, indices, sorted_indices=...): ...
|
||||
|
||||
class LoRAEmbedding(nn.Module):
|
||||
@staticmethod
|
||||
def from_base(
|
||||
embedding: nn.Embedding, r: int = ..., dropout: float = ..., scale: float = ...
|
||||
): # -> LoRAEmbedding:
|
||||
...
|
||||
def fuse(self, dequantize: bool = ...): # -> QuantizedEmbedding | Embedding:
|
||||
...
|
||||
def __init__(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
dims: int,
|
||||
r: int = ...,
|
||||
dropout: float = ...,
|
||||
scale: float = ...,
|
||||
) -> None: ...
|
||||
def __call__(self, x): # -> array:
|
||||
...
|
||||
def as_linear(self, x): # -> array:
|
||||
...
|
||||
@@ -1,57 +0,0 @@
|
||||
"""
|
||||
This type stub file was generated by pyright.
|
||||
"""
|
||||
|
||||
import mlx.nn as nn
|
||||
from typing import Dict
|
||||
|
||||
def build_schedule(schedule_config: Dict): # -> Any:
|
||||
"""
|
||||
Build a learning rate schedule from the given config.
|
||||
"""
|
||||
...
|
||||
|
||||
def linear_to_lora_layers(
|
||||
model: nn.Module, num_layers: int, config: Dict, use_dora: bool = ...
|
||||
): # -> None:
|
||||
"""
|
||||
Convert some of the models linear layers to lora layers.
|
||||
|
||||
Args:
|
||||
model (nn.Module): The neural network model.
|
||||
num_layers (int): The number of blocks to convert to lora layers
|
||||
starting from the last layer.
|
||||
config (dict): More configuration parameters for LoRA, including the
|
||||
rank, scale, and optional layer keys.
|
||||
use_dora (bool): If True, uses DoRA instead of LoRA.
|
||||
Default: ``False``
|
||||
"""
|
||||
...
|
||||
|
||||
def load_adapters(model: nn.Module, adapter_path: str) -> nn.Module:
|
||||
"""
|
||||
Load any fine-tuned adapters / layers.
|
||||
|
||||
Args:
|
||||
model (nn.Module): The neural network model.
|
||||
adapter_path (str): Path to the adapter configuration file.
|
||||
|
||||
Returns:
|
||||
nn.Module: The updated model with LoRA layers applied.
|
||||
"""
|
||||
...
|
||||
|
||||
def remove_lora_layers(model: nn.Module) -> nn.Module:
|
||||
"""
|
||||
Remove the LoRA layers from the model.
|
||||
|
||||
Args:
|
||||
model (nn.Module): The model with LoRA layers.
|
||||
|
||||
Returns:
|
||||
nn.Module: The model without LoRA layers.
|
||||
"""
|
||||
...
|
||||
|
||||
def print_trainable_parameters(model): # -> None:
|
||||
...
|
||||
Whitespace-only changes.
@@ -1,12 +0,0 @@
|
||||
from typing import Any
|
||||
|
||||
def get_message_json(
|
||||
model_name: str,
|
||||
prompt: str,
|
||||
role: str = "user",
|
||||
skip_image_token: bool = False,
|
||||
skip_audio_token: bool = False,
|
||||
num_images: int = 0,
|
||||
num_audios: int = 0,
|
||||
**kwargs: Any,
|
||||
) -> dict[str, Any]: ...
|
||||
@@ -1,15 +0,0 @@
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
class ImageProcessor:
|
||||
def preprocess(
|
||||
self, images: list[dict[str, Any]], **kwargs: Any
|
||||
) -> dict[str, Any]: ...
|
||||
def __call__(self, **kwargs: Any) -> dict[str, Any]: ...
|
||||
|
||||
def load_image_processor(
|
||||
model_path: str | Path, **kwargs: Any
|
||||
) -> ImageProcessor | None: ...
|
||||
def load_processor(
|
||||
model_path: str | Path, add_detokenizer: bool = ..., **kwargs: Any
|
||||
) -> ImageProcessor: ...
|
||||
@@ -1,8 +0,0 @@
|
||||
from typing import Any, Self
|
||||
|
||||
class safe_open:
|
||||
def __init__(self, filename: str, framework: str = "pt") -> None: ...
|
||||
def __enter__(self) -> Self: ...
|
||||
def __exit__(self, *args: Any) -> None: ...
|
||||
def keys(self) -> list[str]: ...
|
||||
def get_tensor(self, name: str) -> Any: ...
|
||||
Vendored
+1
-2
@@ -1,8 +1,7 @@
|
||||
{
|
||||
"recommendations": [
|
||||
"detachhead.basedpyright",
|
||||
"ms-python.python",
|
||||
"jnoortheen.nix-ide"
|
||||
"ms-python.python"
|
||||
],
|
||||
"unwantedRecommendations": [
|
||||
"ms-python.vscode-pylance",
|
||||
|
||||
Vendored
+1
-30
@@ -1,32 +1,3 @@
|
||||
{
|
||||
"files.associations": {
|
||||
"*.nix": "nix",
|
||||
},
|
||||
"nix.enableLanguageServer": true,
|
||||
"nix.serverPath": "nixd",
|
||||
"nix.serverSettings": {
|
||||
"nixd": {
|
||||
"formatting": {
|
||||
"command": ["nixpkgs-fmt"]
|
||||
},
|
||||
"nixpkgs": {
|
||||
"expr": "(builtins.getFlake \"path:${workspaceFolder}\").currentSystem.config._module.args.pkgs"
|
||||
},
|
||||
"options": {
|
||||
"flake-parts": {
|
||||
"expr": "(builtins.getFlake \"path:${workspaceFolder}\").debug.options"
|
||||
},
|
||||
"flake-parts-perSystem": {
|
||||
"expr": "(builtins.getFlake \"path:${workspaceFolder}\").currentSystem.options"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"[nix]": {
|
||||
"editor.defaultFormatter": "jnoortheen.nix-ide"
|
||||
},
|
||||
|
||||
"python.defaultInterpreterPath": "${workspaceFolder}/.venv/bin/python",
|
||||
"basedpyright.analysis.configFilePath": "${workspaceFolder}/pyproject.toml",
|
||||
"basedpyright.importStrategy": "fromEnvironment",
|
||||
"basedpyright.importStrategy": "fromEnvironment"
|
||||
}
|
||||
+1
-123
@@ -11,18 +11,9 @@ To run EXO from source:
|
||||
```bash
|
||||
brew install uv
|
||||
```
|
||||
- [rust](https://github.com/rust-lang/rustup) (to build Rust bindings, nightly for now)
|
||||
```bash
|
||||
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
|
||||
rustup toolchain install nightly
|
||||
```
|
||||
- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
|
||||
Use the pinned fork revision used by this repo instead of Homebrew `macmon`.
|
||||
```bash
|
||||
cargo install --git https://github.com/vladkens/macmon \
|
||||
--rev a1cd06b6cc0d5e61db24fd8832e74cd992097a7d \
|
||||
macmon \
|
||||
--force
|
||||
brew install macmon
|
||||
```
|
||||
|
||||
```bash
|
||||
@@ -48,119 +39,6 @@ Write pure functions where possible. When adding new code, prefer Rust unless th
|
||||
|
||||
Run `nix fmt` to auto-format your code before submitting.
|
||||
|
||||
## Model Cards
|
||||
|
||||
EXO uses TOML-based model cards to define model metadata and capabilities. Model cards are stored in:
|
||||
- `resources/inference_model_cards/` for text generation models
|
||||
- `resources/image_model_cards/` for image generation models
|
||||
- `~/.exo/custom_model_cards/` for user-added custom models
|
||||
|
||||
### Adding a Model Card
|
||||
|
||||
To add a new model, create a TOML file with the following structure:
|
||||
|
||||
```toml
|
||||
model_id = "mlx-community/Llama-3.2-1B-Instruct-4bit"
|
||||
n_layers = 16
|
||||
hidden_size = 2048
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "llama"
|
||||
quantization = "4bit"
|
||||
base_model = "Llama 3.2 1B"
|
||||
capabilities = ["text"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 729808896
|
||||
```
|
||||
|
||||
### Required Fields
|
||||
|
||||
- `model_id`: Hugging Face model identifier
|
||||
- `n_layers`: Number of transformer layers
|
||||
- `hidden_size`: Hidden dimension size
|
||||
- `supports_tensor`: Whether the model supports tensor parallelism
|
||||
- `tasks`: List of supported tasks (`TextGeneration`, `TextToImage`, `ImageToImage`)
|
||||
- `family`: Model family (e.g., "llama", "deepseek", "qwen")
|
||||
- `quantization`: Quantization level (e.g., "4bit", "8bit", "bf16")
|
||||
- `base_model`: Human-readable base model name
|
||||
- `capabilities`: List of capabilities (e.g., `["text"]`, `["text", "thinking"]`)
|
||||
|
||||
### Optional Fields
|
||||
|
||||
- `components`: For multi-component models (like image models with separate text encoders and transformers)
|
||||
- `uses_cfg`: Whether the model uses classifier-free guidance (for image models)
|
||||
- `trust_remote_code`: Whether to allow remote code execution (defaults to `false` for security)
|
||||
|
||||
### Capabilities
|
||||
|
||||
The `capabilities` field defines what the model can do:
|
||||
- `text`: Standard text generation
|
||||
- `thinking`: Model supports chain-of-thought reasoning
|
||||
- `thinking_toggle`: Thinking can be enabled/disabled via `enable_thinking` parameter
|
||||
- `image_edit`: Model supports image-to-image editing (FLUX.1-Kontext)
|
||||
|
||||
### Security Note
|
||||
|
||||
By default, `trust_remote_code` is set to `false` for security. Only enable it if the model explicitly requires remote code execution from the Hugging Face hub.
|
||||
|
||||
## API Adapters
|
||||
|
||||
EXO supports multiple API formats through an adapter pattern. Adapters convert API-specific request formats to the internal `TextGenerationTaskParams` format and convert internal token chunks back to API-specific responses.
|
||||
|
||||
### Adapter Architecture
|
||||
|
||||
All adapters live in `src/exo/master/adapters/` and follow the same pattern:
|
||||
|
||||
1. Convert API-specific requests to `TextGenerationTaskParams`
|
||||
2. Handle both streaming and non-streaming response generation
|
||||
3. Convert internal `TokenChunk` objects to API-specific formats
|
||||
4. Manage error handling and edge cases
|
||||
|
||||
### Existing Adapters
|
||||
|
||||
- `chat_completions.py`: OpenAI Chat Completions API
|
||||
- `claude.py`: Anthropic Claude Messages API
|
||||
- `responses.py`: OpenAI Responses API
|
||||
- `ollama.py`: Ollama API (for OpenWebUI compatibility)
|
||||
|
||||
### Adding a New API Adapter
|
||||
|
||||
To add support for a new API format:
|
||||
|
||||
1. Create a new adapter file in `src/exo/master/adapters/`
|
||||
2. Implement a request conversion function:
|
||||
```python
|
||||
def your_api_request_to_text_generation(
|
||||
request: YourAPIRequest,
|
||||
) -> TextGenerationTaskParams:
|
||||
# Convert API request to internal format
|
||||
pass
|
||||
```
|
||||
3. Implement streaming response generation:
|
||||
```python
|
||||
async def generate_your_api_stream(
|
||||
command_id: CommandId,
|
||||
chunk_stream: AsyncGenerator[TokenChunk | ErrorChunk | ToolCallChunk, None],
|
||||
) -> AsyncGenerator[str, None]:
|
||||
# Convert internal chunks to API-specific streaming format
|
||||
pass
|
||||
```
|
||||
4. Implement non-streaming response collection:
|
||||
```python
|
||||
async def collect_your_api_response(
|
||||
command_id: CommandId,
|
||||
chunk_stream: AsyncGenerator[TokenChunk | ErrorChunk | ToolCallChunk, None],
|
||||
) -> AsyncGenerator[str]:
|
||||
# Collect all chunks and return single response
|
||||
pass
|
||||
```
|
||||
5. Register the adapter endpoints in `src/exo/master/api.py`
|
||||
|
||||
The adapter pattern keeps API-specific logic isolated from core inference systems. Internal systems (worker, runner, event sourcing) only see `TextGenerationTaskParams` and `TokenChunk` objects - no API-specific types cross the adapter boundary.
|
||||
|
||||
For detailed API documentation, see [docs/api.md](docs/api.md).
|
||||
|
||||
## Testing
|
||||
|
||||
EXO relies heavily on manual testing at this point in the project, but this is evolving. Before submitting a change, test both before and after to demonstrate how your change improves behavior. Do the best you can with the hardware you have available - if you need help testing, ask and we'll do our best to assist. Add automated tests where possible - we're actively working to substantially improve our automated testing story.
|
||||
|
||||
Generated
+125
-35
@@ -141,6 +141,12 @@ version = "0.3.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb"
|
||||
|
||||
[[package]]
|
||||
name = "arrayvec"
|
||||
version = "0.7.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50"
|
||||
|
||||
[[package]]
|
||||
name = "asn1-rs"
|
||||
version = "0.7.1"
|
||||
@@ -216,28 +222,6 @@ dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-stream"
|
||||
version = "0.3.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476"
|
||||
dependencies = [
|
||||
"async-stream-impl",
|
||||
"futures-core",
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-stream-impl"
|
||||
version = "0.3.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.111",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-trait"
|
||||
version = "0.1.89"
|
||||
@@ -320,6 +304,19 @@ version = "1.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "55248b47b0caf0546f7988906588779981c43bb1bc9d0c44087278f80cdb44ba"
|
||||
|
||||
[[package]]
|
||||
name = "bigdecimal"
|
||||
version = "0.4.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "560f42649de9fa436b73517378a147ec21f6c997a546581df4b4b31677828934"
|
||||
dependencies = [
|
||||
"autocfg",
|
||||
"libm",
|
||||
"num-bigint",
|
||||
"num-integer",
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bimap"
|
||||
version = "0.6.3"
|
||||
@@ -519,6 +516,15 @@ version = "0.4.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2f421161cb492475f1661ddc9815a745a1c894592070661180fdec3d4872e9c3"
|
||||
|
||||
[[package]]
|
||||
name = "convert_case"
|
||||
version = "0.10.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "633458d4ef8c78b72454de2d54fd6ab2e60f9e02be22f3c6104cdc8a4e0fceb9"
|
||||
dependencies = [
|
||||
"unicode-segmentation",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "core-foundation"
|
||||
version = "0.9.4"
|
||||
@@ -740,6 +746,29 @@ dependencies = [
|
||||
"powerfmt",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "derive_more"
|
||||
version = "2.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "10b768e943bed7bf2cab53df09f4bc34bfd217cdb57d971e769874c9a6710618"
|
||||
dependencies = [
|
||||
"derive_more-impl",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "derive_more-impl"
|
||||
version = "2.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6d286bfdaf75e988b4a78e013ecd79c581e06399ab53fbacd2d916c2f904f30b"
|
||||
dependencies = [
|
||||
"convert_case",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"rustc_version",
|
||||
"syn 2.0.111",
|
||||
"unicode-xid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "digest"
|
||||
version = "0.10.7"
|
||||
@@ -910,17 +939,22 @@ name = "exo_pyo3_bindings"
|
||||
version = "0.0.1"
|
||||
dependencies = [
|
||||
"delegate",
|
||||
"derive_more",
|
||||
"env_logger",
|
||||
"extend",
|
||||
"futures-lite",
|
||||
"futures",
|
||||
"impl-trait-for-tuples",
|
||||
"libp2p",
|
||||
"log",
|
||||
"networking",
|
||||
"once_cell",
|
||||
"pin-project",
|
||||
"pyo3",
|
||||
"pyo3-async-runtimes",
|
||||
"pyo3-log",
|
||||
"pyo3-stub-gen",
|
||||
"thiserror 2.0.17",
|
||||
"thread_local",
|
||||
"tokio",
|
||||
"util",
|
||||
]
|
||||
@@ -936,12 +970,6 @@ dependencies = [
|
||||
"syn 2.0.111",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fastrand"
|
||||
version = "2.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be"
|
||||
|
||||
[[package]]
|
||||
name = "ff"
|
||||
version = "0.13.1"
|
||||
@@ -1050,10 +1078,7 @@ version = "2.6.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f78e10609fe0e0b3f4157ffab1876319b5b0db102a2c60dc4626306dc46b44ad"
|
||||
dependencies = [
|
||||
"fastrand",
|
||||
"futures-core",
|
||||
"futures-io",
|
||||
"parking",
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
@@ -1615,6 +1640,17 @@ dependencies = [
|
||||
"xmltree",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "impl-trait-for-tuples"
|
||||
version = "0.2.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a0eb5a3343abf848c0984fe4604b2b105da9539376e24fc0a3b0007411ae4fd9"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.111",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "indexmap"
|
||||
version = "2.12.1"
|
||||
@@ -1793,6 +1829,12 @@ version = "0.2.178"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "37c93d8daa9d8a012fd8ab92f088405fb202ea0b6ab73ee2482ae66af4f42091"
|
||||
|
||||
[[package]]
|
||||
name = "libm"
|
||||
version = "0.2.15"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de"
|
||||
|
||||
[[package]]
|
||||
name = "libp2p"
|
||||
version = "0.56.0"
|
||||
@@ -2781,16 +2823,17 @@ dependencies = [
|
||||
name = "networking"
|
||||
version = "0.0.1"
|
||||
dependencies = [
|
||||
"async-stream",
|
||||
"delegate",
|
||||
"derive_more",
|
||||
"either",
|
||||
"extend",
|
||||
"futures-lite",
|
||||
"futures",
|
||||
"futures-timer",
|
||||
"impl-trait-for-tuples",
|
||||
"keccak-const",
|
||||
"libp2p",
|
||||
"log",
|
||||
"pin-project",
|
||||
"thiserror 2.0.17",
|
||||
"tokio",
|
||||
"tracing-subscriber",
|
||||
"util",
|
||||
@@ -2875,6 +2918,17 @@ dependencies = [
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-rational"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824"
|
||||
dependencies = [
|
||||
"num-bigint",
|
||||
"num-integer",
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-traits"
|
||||
version = "0.2.19"
|
||||
@@ -3225,14 +3279,28 @@ version = "0.27.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ab53c047fcd1a1d2a8820fe84f05d6be69e9526be40cb03b73f86b6b03e6d87d"
|
||||
dependencies = [
|
||||
"bigdecimal",
|
||||
"either",
|
||||
"hashbrown 0.16.1",
|
||||
"indexmap",
|
||||
"indoc",
|
||||
"inventory",
|
||||
"libc",
|
||||
"lock_api",
|
||||
"memoffset",
|
||||
"num-bigint",
|
||||
"num-complex",
|
||||
"num-rational",
|
||||
"num-traits",
|
||||
"once_cell",
|
||||
"ordered-float",
|
||||
"parking_lot",
|
||||
"portable-atomic",
|
||||
"pyo3-build-config",
|
||||
"pyo3-ffi",
|
||||
"pyo3-macros",
|
||||
"rust_decimal",
|
||||
"smallvec",
|
||||
"unindent",
|
||||
]
|
||||
|
||||
@@ -3673,6 +3741,16 @@ dependencies = [
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rust_decimal"
|
||||
version = "1.39.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "35affe401787a9bd846712274d97654355d21b2a2c092a3139aabe31e9022282"
|
||||
dependencies = [
|
||||
"arrayvec",
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustc-hash"
|
||||
version = "1.1.0"
|
||||
@@ -4537,12 +4615,24 @@ version = "1.0.22"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5"
|
||||
|
||||
[[package]]
|
||||
name = "unicode-segmentation"
|
||||
version = "1.12.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493"
|
||||
|
||||
[[package]]
|
||||
name = "unicode-width"
|
||||
version = "0.2.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254"
|
||||
|
||||
[[package]]
|
||||
name = "unicode-xid"
|
||||
version = "0.2.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853"
|
||||
|
||||
[[package]]
|
||||
name = "unicode_names2"
|
||||
version = "1.3.0"
|
||||
|
||||
+35
-3
@@ -1,6 +1,10 @@
|
||||
[workspace]
|
||||
resolver = "3"
|
||||
members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"]
|
||||
members = [
|
||||
"rust/networking",
|
||||
"rust/exo_pyo3_bindings",
|
||||
"rust/util",
|
||||
]
|
||||
|
||||
[workspace.package]
|
||||
version = "0.0.1"
|
||||
@@ -22,21 +26,49 @@ opt-level = 3
|
||||
networking = { path = "rust/networking" }
|
||||
util = { path = "rust/util" }
|
||||
|
||||
# Proc-macro authoring tools
|
||||
syn = "2.0"
|
||||
quote = "1.0"
|
||||
proc-macro2 = "1.0"
|
||||
darling = "0.20"
|
||||
|
||||
# Macro dependecies
|
||||
extend = "1.2"
|
||||
delegate = "0.13"
|
||||
impl-trait-for-tuples = "0.2"
|
||||
clap = "4.5"
|
||||
derive_more = { version = "2.0.1", features = ["display"] }
|
||||
pin-project = "1"
|
||||
|
||||
# Utility dependencies
|
||||
itertools = "0.14"
|
||||
thiserror = "2"
|
||||
internment = "0.8"
|
||||
recursion = "0.5"
|
||||
regex = "1.11"
|
||||
once_cell = "1.21"
|
||||
thread_local = "1.1"
|
||||
bon = "3.4"
|
||||
generativity = "1.1"
|
||||
anyhow = "1.0"
|
||||
keccak-const = "0.2"
|
||||
|
||||
# Functional generics/lenses frameworks
|
||||
frunk_core = "0.4"
|
||||
frunk = "0.4"
|
||||
frunk_utils = "0.2"
|
||||
frunk-enum-core = "0.3"
|
||||
|
||||
# Async dependencies
|
||||
async-stream = "0.3"
|
||||
tokio = "1.46"
|
||||
futures-lite = "2.6.1"
|
||||
futures = "0.3"
|
||||
futures-util = "0.3"
|
||||
futures-timer = "3.0"
|
||||
|
||||
# Data structures
|
||||
either = "1.15"
|
||||
ordered-float = "5.0"
|
||||
ahash = "0.8"
|
||||
|
||||
# Tracing/logging
|
||||
log = "0.4"
|
||||
|
||||
@@ -26,8 +26,6 @@ exo connects all your devices into an AI cluster. Not only does exo enable runni
|
||||
- **Topology-Aware Auto Parallel**: exo figures out the best way to split your model across all available devices based on a realtime view of your device topology. It takes into account device resources and network latency/bandwidth between each link.
|
||||
- **Tensor Parallelism**: exo supports sharding models, for up to 1.8x speedup on 2 devices and 3.2x speedup on 4 devices.
|
||||
- **MLX Support**: exo uses [MLX](https://github.com/ml-explore/mlx) as an inference backend and [MLX distributed](https://ml-explore.github.io/mlx/build/html/usage/distributed.html) for distributed communication.
|
||||
- **Multiple API Compatibility**: Compatible with OpenAI Chat Completions API, Claude Messages API, OpenAI Responses API, and Ollama API - use your existing tools and clients.
|
||||
- **Custom Model Support**: Load custom models from HuggingFace hub to expand the range of available models.
|
||||
|
||||
## Dashboard
|
||||
|
||||
@@ -74,31 +72,18 @@ There are two ways to run exo:
|
||||
|
||||
### Run from Source (macOS)
|
||||
|
||||
If you have [Nix](https://nixos.org/) installed, you can skip most of the steps below and run exo directly:
|
||||
|
||||
```bash
|
||||
nix run .#exo
|
||||
```
|
||||
|
||||
**Note:** To accept the Cachix binary cache (and avoid the Xcode Metal ToolChain), add to `/etc/nix/nix.conf`:
|
||||
```
|
||||
trusted-users = root (or your username)
|
||||
experimental-features = nix-command flakes
|
||||
```
|
||||
Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-daemon`
|
||||
|
||||
**Prerequisites:**
|
||||
- [Xcode](https://developer.apple.com/xcode/) (provides the Metal ToolChain required for MLX compilation)
|
||||
- [brew](https://github.com/Homebrew/brew) (for simple package management on macOS)
|
||||
|
||||
|
||||
```bash
|
||||
/bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
|
||||
```
|
||||
- [uv](https://github.com/astral-sh/uv) (for Python dependency management)
|
||||
- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
|
||||
- [node](https://github.com/nodejs/node) (for building the dashboard)
|
||||
|
||||
|
||||
```bash
|
||||
brew install uv node
|
||||
brew install uv macmon node
|
||||
```
|
||||
- [rust](https://github.com/rust-lang/rustup) (to build Rust bindings, nightly for now)
|
||||
|
||||
@@ -106,17 +91,6 @@ Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-d
|
||||
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
|
||||
rustup toolchain install nightly
|
||||
```
|
||||
- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
|
||||
|
||||
Install the pinned fork revision used by this repo instead of Homebrew `macmon`.
|
||||
Homebrew `macmon 0.6.1` still crashes on Apple M5.
|
||||
|
||||
```bash
|
||||
cargo install --git https://github.com/vladkens/macmon \
|
||||
--rev a1cd06b6cc0d5e61db24fd8832e74cd992097a7d \
|
||||
macmon \
|
||||
--force
|
||||
```
|
||||
|
||||
Clone the repo, build the dashboard, and run exo:
|
||||
|
||||
@@ -208,8 +182,6 @@ exo follows the [XDG Base Directory Specification](https://specifications.freede
|
||||
- **Configuration files**: `~/.config/exo/` (or `$XDG_CONFIG_HOME/exo/`)
|
||||
- **Data files**: `~/.local/share/exo/` (or `$XDG_DATA_HOME/exo/`)
|
||||
- **Cache files**: `~/.cache/exo/` (or `$XDG_CACHE_HOME/exo/`)
|
||||
- **Log files**: `~/.cache/exo/exo_log/` (with automatic log rotation)
|
||||
- **Custom model cards**: `~/.local/share/exo/custom_model_cards/`
|
||||
|
||||
You can override these locations by setting the corresponding XDG environment variables.
|
||||
|
||||
@@ -289,51 +261,8 @@ After that, RDMA will be enabled in macOS and exo will take care of the rest.
|
||||
|
||||
---
|
||||
|
||||
## Environment Variables
|
||||
|
||||
exo supports several environment variables for configuration:
|
||||
|
||||
| Variable | Description | Default |
|
||||
|----------|-------------|---------|
|
||||
| `EXO_DEFAULT_MODELS_DIR` | Default directory for model downloads and caches. Always first in the writable dirs list. | `~/.local/share/exo/models` (Linux) or `~/.exo/models` (macOS) |
|
||||
| `EXO_MODELS_DIRS` | Colon-separated additional writable directories for model downloads. Checked in order after the default; first with enough free space is used. | None |
|
||||
| `EXO_MODELS_READ_ONLY_DIRS` | Colon-separated read-only directories to search for pre-downloaded models (e.g., NFS mounts, shared storage). Models here cannot be deleted. | None |
|
||||
| `EXO_OFFLINE` | Run without internet connection (uses only local models) | `false` |
|
||||
| `EXO_ENABLE_IMAGE_MODELS` | Enable image model support | `false` |
|
||||
| `EXO_LIBP2P_NAMESPACE` | Custom namespace for cluster isolation | None |
|
||||
| `EXO_FAST_SYNCH` | Control MLX_METAL_FAST_SYNCH behavior (for JACCL backend) | Auto |
|
||||
| `EXO_TRACING_ENABLED` | Enable distributed tracing for performance analysis | `false` |
|
||||
|
||||
**Example usage:**
|
||||
|
||||
```bash
|
||||
# Use pre-downloaded models from NFS mount (read-only)
|
||||
EXO_MODELS_READ_ONLY_DIRS=/mnt/nfs/models:/opt/ai-models uv run exo
|
||||
|
||||
# Download models to an external SSD (falls back to default dir if full)
|
||||
EXO_MODELS_DIRS=/Volumes/ExternalSSD/exo-models uv run exo
|
||||
|
||||
# Run in offline mode
|
||||
EXO_OFFLINE=true uv run exo
|
||||
|
||||
# Enable image models
|
||||
EXO_ENABLE_IMAGE_MODELS=true uv run exo
|
||||
|
||||
# Use custom namespace for cluster isolation
|
||||
EXO_LIBP2P_NAMESPACE=my-dev-cluster uv run exo
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Using the API
|
||||
|
||||
exo provides multiple API-compatible interfaces for maximum compatibility with existing tools:
|
||||
|
||||
- **OpenAI Chat Completions API** - Compatible with OpenAI clients
|
||||
- **Claude Messages API** - Compatible with Anthropic's Claude format
|
||||
- **OpenAI Responses API** - Compatible with OpenAI's Responses format
|
||||
- **Ollama API** - Compatible with Ollama and tools like OpenWebUI
|
||||
|
||||
If you prefer to interact with exo via the API, here is an example creating an instance of a small model (`mlx-community/Llama-3.2-1B-Instruct-4bit`), sending a chat completions request and deleting the instance.
|
||||
|
||||
---
|
||||
@@ -423,85 +352,14 @@ When you're done, delete the instance by its ID (find it via `/state` or `/insta
|
||||
curl -X DELETE http://localhost:52415/instance/YOUR_INSTANCE_ID
|
||||
```
|
||||
|
||||
### Claude Messages API Compatibility
|
||||
|
||||
Use the Claude Messages API format with the `/v1/messages` endpoint:
|
||||
|
||||
```bash
|
||||
curl -N -X POST http://localhost:52415/v1/messages \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"model": "mlx-community/Llama-3.2-1B-Instruct-4bit",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello"}
|
||||
],
|
||||
"max_tokens": 1024,
|
||||
"stream": true
|
||||
}'
|
||||
```
|
||||
|
||||
### OpenAI Responses API Compatibility
|
||||
|
||||
Use the OpenAI Responses API format with the `/v1/responses` endpoint:
|
||||
|
||||
```bash
|
||||
curl -N -X POST http://localhost:52415/v1/responses \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"model": "mlx-community/Llama-3.2-1B-Instruct-4bit",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello"}
|
||||
],
|
||||
"stream": true
|
||||
}'
|
||||
```
|
||||
|
||||
### Ollama API Compatibility
|
||||
|
||||
exo supports Ollama API endpoints for compatibility with tools like OpenWebUI:
|
||||
|
||||
```bash
|
||||
# Ollama chat
|
||||
curl -X POST http://localhost:52415/ollama/api/chat \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"model": "mlx-community/Llama-3.2-1B-Instruct-4bit",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello"}
|
||||
],
|
||||
"stream": false
|
||||
}'
|
||||
|
||||
# List models (Ollama format)
|
||||
curl http://localhost:52415/ollama/api/tags
|
||||
```
|
||||
|
||||
### Custom Model Loading from HuggingFace
|
||||
|
||||
You can add custom models from the HuggingFace hub:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:52415/models/add \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"model_id": "mlx-community/my-custom-model"
|
||||
}'
|
||||
```
|
||||
|
||||
**Security Note:**
|
||||
|
||||
Custom models requiring `trust_remote_code` in their configuration must be explicitly enabled (default is false) for security. Only enable this if you trust the model's remote code execution. Models are fetched from HuggingFace and stored locally as custom model cards.
|
||||
|
||||
**Other useful API endpoints*:**
|
||||
|
||||
- List all models: `curl http://localhost:52415/models`
|
||||
- List downloaded models only: `curl http://localhost:52415/models?status=downloaded`
|
||||
- Search HuggingFace: `curl "http://localhost:52415/models/search?query=llama&limit=10"`
|
||||
- Inspect instance IDs and deployment state: `curl http://localhost:52415/state`
|
||||
|
||||
For further details, see:
|
||||
|
||||
- API documentation in [docs/api.md](docs/api.md).
|
||||
- API basic documentation in [docs/api.md](docs/api.md).
|
||||
- API types and endpoints in [src/exo/master/api.py](src/exo/master/api.py).
|
||||
|
||||
---
|
||||
@@ -560,4 +418,4 @@ On macOS, exo uses the GPU. On Linux, exo currently runs on CPU. We are working
|
||||
|
||||
## Contributing
|
||||
|
||||
See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on how to contribute to exo.
|
||||
See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on how to contribute to exo.
|
||||
@@ -0,0 +1,628 @@
|
||||
// !$*UTF8*$!
|
||||
{
|
||||
archiveVersion = 1;
|
||||
classes = {
|
||||
};
|
||||
objectVersion = 77;
|
||||
objects = {
|
||||
|
||||
/* Begin PBXBuildFile section */
|
||||
E09D17522F44F359009C51A3 /* MLXLLM in Frameworks */ = {isa = PBXBuildFile; productRef = E09D17512F44F359009C51A3 /* MLXLLM */; };
|
||||
E09D17542F44F359009C51A3 /* MLXLMCommon in Frameworks */ = {isa = PBXBuildFile; productRef = E09D17532F44F359009C51A3 /* MLXLMCommon */; };
|
||||
/* End PBXBuildFile section */
|
||||
|
||||
/* Begin PBXContainerItemProxy section */
|
||||
E09D167D2F44CA20009C51A3 /* PBXContainerItemProxy */ = {
|
||||
isa = PBXContainerItemProxy;
|
||||
containerPortal = E09D16672F44CA1E009C51A3 /* Project object */;
|
||||
proxyType = 1;
|
||||
remoteGlobalIDString = E09D166E2F44CA1E009C51A3;
|
||||
remoteInfo = "EXO-iOS";
|
||||
};
|
||||
E09D16872F44CA20009C51A3 /* PBXContainerItemProxy */ = {
|
||||
isa = PBXContainerItemProxy;
|
||||
containerPortal = E09D16672F44CA1E009C51A3 /* Project object */;
|
||||
proxyType = 1;
|
||||
remoteGlobalIDString = E09D166E2F44CA1E009C51A3;
|
||||
remoteInfo = "EXO-iOS";
|
||||
};
|
||||
/* End PBXContainerItemProxy section */
|
||||
|
||||
/* Begin PBXFileReference section */
|
||||
E09D166F2F44CA1E009C51A3 /* EXO-iOS.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = "EXO-iOS.app"; sourceTree = BUILT_PRODUCTS_DIR; };
|
||||
E09D167C2F44CA20009C51A3 /* EXO-iOSTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = "EXO-iOSTests.xctest"; sourceTree = BUILT_PRODUCTS_DIR; };
|
||||
E09D16862F44CA20009C51A3 /* EXO-iOSUITests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = "EXO-iOSUITests.xctest"; sourceTree = BUILT_PRODUCTS_DIR; };
|
||||
/* End PBXFileReference section */
|
||||
|
||||
/* Begin PBXFileSystemSynchronizedBuildFileExceptionSet section */
|
||||
E09D169A2F44CA20009C51A3 /* Exceptions for "EXO-iOS" folder in "EXO-iOS" target */ = {
|
||||
isa = PBXFileSystemSynchronizedBuildFileExceptionSet;
|
||||
membershipExceptions = (
|
||||
Info.plist,
|
||||
);
|
||||
target = E09D166E2F44CA1E009C51A3 /* EXO-iOS */;
|
||||
};
|
||||
/* End PBXFileSystemSynchronizedBuildFileExceptionSet section */
|
||||
|
||||
/* Begin PBXFileSystemSynchronizedRootGroup section */
|
||||
E09D16712F44CA1E009C51A3 /* EXO-iOS */ = {
|
||||
isa = PBXFileSystemSynchronizedRootGroup;
|
||||
exceptions = (
|
||||
E09D169A2F44CA20009C51A3 /* Exceptions for "EXO-iOS" folder in "EXO-iOS" target */,
|
||||
);
|
||||
path = "EXO-iOS";
|
||||
sourceTree = "<group>";
|
||||
};
|
||||
E09D167F2F44CA20009C51A3 /* EXO-iOSTests */ = {
|
||||
isa = PBXFileSystemSynchronizedRootGroup;
|
||||
path = "EXO-iOSTests";
|
||||
sourceTree = "<group>";
|
||||
};
|
||||
E09D16892F44CA20009C51A3 /* EXO-iOSUITests */ = {
|
||||
isa = PBXFileSystemSynchronizedRootGroup;
|
||||
path = "EXO-iOSUITests";
|
||||
sourceTree = "<group>";
|
||||
};
|
||||
/* End PBXFileSystemSynchronizedRootGroup section */
|
||||
|
||||
/* Begin PBXFrameworksBuildPhase section */
|
||||
E09D166C2F44CA1E009C51A3 /* Frameworks */ = {
|
||||
isa = PBXFrameworksBuildPhase;
|
||||
buildActionMask = 2147483647;
|
||||
files = (
|
||||
E09D17542F44F359009C51A3 /* MLXLMCommon in Frameworks */,
|
||||
E09D17522F44F359009C51A3 /* MLXLLM in Frameworks */,
|
||||
);
|
||||
runOnlyForDeploymentPostprocessing = 0;
|
||||
};
|
||||
E09D16792F44CA20009C51A3 /* Frameworks */ = {
|
||||
isa = PBXFrameworksBuildPhase;
|
||||
buildActionMask = 2147483647;
|
||||
files = (
|
||||
);
|
||||
runOnlyForDeploymentPostprocessing = 0;
|
||||
};
|
||||
E09D16832F44CA20009C51A3 /* Frameworks */ = {
|
||||
isa = PBXFrameworksBuildPhase;
|
||||
buildActionMask = 2147483647;
|
||||
files = (
|
||||
);
|
||||
runOnlyForDeploymentPostprocessing = 0;
|
||||
};
|
||||
/* End PBXFrameworksBuildPhase section */
|
||||
|
||||
/* Begin PBXGroup section */
|
||||
E09D16662F44CA1E009C51A3 = {
|
||||
isa = PBXGroup;
|
||||
children = (
|
||||
E09D16712F44CA1E009C51A3 /* EXO-iOS */,
|
||||
E09D167F2F44CA20009C51A3 /* EXO-iOSTests */,
|
||||
E09D16892F44CA20009C51A3 /* EXO-iOSUITests */,
|
||||
E09D16702F44CA1E009C51A3 /* Products */,
|
||||
);
|
||||
sourceTree = "<group>";
|
||||
};
|
||||
E09D16702F44CA1E009C51A3 /* Products */ = {
|
||||
isa = PBXGroup;
|
||||
children = (
|
||||
E09D166F2F44CA1E009C51A3 /* EXO-iOS.app */,
|
||||
E09D167C2F44CA20009C51A3 /* EXO-iOSTests.xctest */,
|
||||
E09D16862F44CA20009C51A3 /* EXO-iOSUITests.xctest */,
|
||||
);
|
||||
name = Products;
|
||||
sourceTree = "<group>";
|
||||
};
|
||||
/* End PBXGroup section */
|
||||
|
||||
/* Begin PBXNativeTarget section */
|
||||
E09D166E2F44CA1E009C51A3 /* EXO-iOS */ = {
|
||||
isa = PBXNativeTarget;
|
||||
buildConfigurationList = E09D16902F44CA20009C51A3 /* Build configuration list for PBXNativeTarget "EXO-iOS" */;
|
||||
buildPhases = (
|
||||
E09D166B2F44CA1E009C51A3 /* Sources */,
|
||||
E09D166C2F44CA1E009C51A3 /* Frameworks */,
|
||||
E09D166D2F44CA1E009C51A3 /* Resources */,
|
||||
);
|
||||
buildRules = (
|
||||
);
|
||||
dependencies = (
|
||||
);
|
||||
fileSystemSynchronizedGroups = (
|
||||
E09D16712F44CA1E009C51A3 /* EXO-iOS */,
|
||||
);
|
||||
name = "EXO-iOS";
|
||||
packageProductDependencies = (
|
||||
E09D17512F44F359009C51A3 /* MLXLLM */,
|
||||
E09D17532F44F359009C51A3 /* MLXLMCommon */,
|
||||
);
|
||||
productName = "EXO-iOS";
|
||||
productReference = E09D166F2F44CA1E009C51A3 /* EXO-iOS.app */;
|
||||
productType = "com.apple.product-type.application";
|
||||
};
|
||||
E09D167B2F44CA20009C51A3 /* EXO-iOSTests */ = {
|
||||
isa = PBXNativeTarget;
|
||||
buildConfigurationList = E09D16932F44CA20009C51A3 /* Build configuration list for PBXNativeTarget "EXO-iOSTests" */;
|
||||
buildPhases = (
|
||||
E09D16782F44CA20009C51A3 /* Sources */,
|
||||
E09D16792F44CA20009C51A3 /* Frameworks */,
|
||||
E09D167A2F44CA20009C51A3 /* Resources */,
|
||||
);
|
||||
buildRules = (
|
||||
);
|
||||
dependencies = (
|
||||
E09D167E2F44CA20009C51A3 /* PBXTargetDependency */,
|
||||
);
|
||||
fileSystemSynchronizedGroups = (
|
||||
E09D167F2F44CA20009C51A3 /* EXO-iOSTests */,
|
||||
);
|
||||
name = "EXO-iOSTests";
|
||||
packageProductDependencies = (
|
||||
);
|
||||
productName = "EXO-iOSTests";
|
||||
productReference = E09D167C2F44CA20009C51A3 /* EXO-iOSTests.xctest */;
|
||||
productType = "com.apple.product-type.bundle.unit-test";
|
||||
};
|
||||
E09D16852F44CA20009C51A3 /* EXO-iOSUITests */ = {
|
||||
isa = PBXNativeTarget;
|
||||
buildConfigurationList = E09D16962F44CA20009C51A3 /* Build configuration list for PBXNativeTarget "EXO-iOSUITests" */;
|
||||
buildPhases = (
|
||||
E09D16822F44CA20009C51A3 /* Sources */,
|
||||
E09D16832F44CA20009C51A3 /* Frameworks */,
|
||||
E09D16842F44CA20009C51A3 /* Resources */,
|
||||
);
|
||||
buildRules = (
|
||||
);
|
||||
dependencies = (
|
||||
E09D16882F44CA20009C51A3 /* PBXTargetDependency */,
|
||||
);
|
||||
fileSystemSynchronizedGroups = (
|
||||
E09D16892F44CA20009C51A3 /* EXO-iOSUITests */,
|
||||
);
|
||||
name = "EXO-iOSUITests";
|
||||
packageProductDependencies = (
|
||||
);
|
||||
productName = "EXO-iOSUITests";
|
||||
productReference = E09D16862F44CA20009C51A3 /* EXO-iOSUITests.xctest */;
|
||||
productType = "com.apple.product-type.bundle.ui-testing";
|
||||
};
|
||||
/* End PBXNativeTarget section */
|
||||
|
||||
/* Begin PBXProject section */
|
||||
E09D16672F44CA1E009C51A3 /* Project object */ = {
|
||||
isa = PBXProject;
|
||||
attributes = {
|
||||
BuildIndependentTargetsInParallel = 1;
|
||||
LastSwiftUpdateCheck = 2620;
|
||||
LastUpgradeCheck = 2620;
|
||||
TargetAttributes = {
|
||||
E09D166E2F44CA1E009C51A3 = {
|
||||
CreatedOnToolsVersion = 26.2;
|
||||
};
|
||||
E09D167B2F44CA20009C51A3 = {
|
||||
CreatedOnToolsVersion = 26.2;
|
||||
TestTargetID = E09D166E2F44CA1E009C51A3;
|
||||
};
|
||||
E09D16852F44CA20009C51A3 = {
|
||||
CreatedOnToolsVersion = 26.2;
|
||||
TestTargetID = E09D166E2F44CA1E009C51A3;
|
||||
};
|
||||
};
|
||||
};
|
||||
buildConfigurationList = E09D166A2F44CA1E009C51A3 /* Build configuration list for PBXProject "EXO-iOS" */;
|
||||
developmentRegion = en;
|
||||
hasScannedForEncodings = 0;
|
||||
knownRegions = (
|
||||
en,
|
||||
Base,
|
||||
);
|
||||
mainGroup = E09D16662F44CA1E009C51A3;
|
||||
minimizedProjectReferenceProxies = 1;
|
||||
packageReferences = (
|
||||
E09D17502F44F359009C51A3 /* XCRemoteSwiftPackageReference "mlx-swift-lm" */,
|
||||
);
|
||||
preferredProjectObjectVersion = 77;
|
||||
productRefGroup = E09D16702F44CA1E009C51A3 /* Products */;
|
||||
projectDirPath = "";
|
||||
projectRoot = "";
|
||||
targets = (
|
||||
E09D166E2F44CA1E009C51A3 /* EXO-iOS */,
|
||||
E09D167B2F44CA20009C51A3 /* EXO-iOSTests */,
|
||||
E09D16852F44CA20009C51A3 /* EXO-iOSUITests */,
|
||||
);
|
||||
};
|
||||
/* End PBXProject section */
|
||||
|
||||
/* Begin PBXResourcesBuildPhase section */
|
||||
E09D166D2F44CA1E009C51A3 /* Resources */ = {
|
||||
isa = PBXResourcesBuildPhase;
|
||||
buildActionMask = 2147483647;
|
||||
files = (
|
||||
);
|
||||
runOnlyForDeploymentPostprocessing = 0;
|
||||
};
|
||||
E09D167A2F44CA20009C51A3 /* Resources */ = {
|
||||
isa = PBXResourcesBuildPhase;
|
||||
buildActionMask = 2147483647;
|
||||
files = (
|
||||
);
|
||||
runOnlyForDeploymentPostprocessing = 0;
|
||||
};
|
||||
E09D16842F44CA20009C51A3 /* Resources */ = {
|
||||
isa = PBXResourcesBuildPhase;
|
||||
buildActionMask = 2147483647;
|
||||
files = (
|
||||
);
|
||||
runOnlyForDeploymentPostprocessing = 0;
|
||||
};
|
||||
/* End PBXResourcesBuildPhase section */
|
||||
|
||||
/* Begin PBXSourcesBuildPhase section */
|
||||
E09D166B2F44CA1E009C51A3 /* Sources */ = {
|
||||
isa = PBXSourcesBuildPhase;
|
||||
buildActionMask = 2147483647;
|
||||
files = (
|
||||
);
|
||||
runOnlyForDeploymentPostprocessing = 0;
|
||||
};
|
||||
E09D16782F44CA20009C51A3 /* Sources */ = {
|
||||
isa = PBXSourcesBuildPhase;
|
||||
buildActionMask = 2147483647;
|
||||
files = (
|
||||
);
|
||||
runOnlyForDeploymentPostprocessing = 0;
|
||||
};
|
||||
E09D16822F44CA20009C51A3 /* Sources */ = {
|
||||
isa = PBXSourcesBuildPhase;
|
||||
buildActionMask = 2147483647;
|
||||
files = (
|
||||
);
|
||||
runOnlyForDeploymentPostprocessing = 0;
|
||||
};
|
||||
/* End PBXSourcesBuildPhase section */
|
||||
|
||||
/* Begin PBXTargetDependency section */
|
||||
E09D167E2F44CA20009C51A3 /* PBXTargetDependency */ = {
|
||||
isa = PBXTargetDependency;
|
||||
target = E09D166E2F44CA1E009C51A3 /* EXO-iOS */;
|
||||
targetProxy = E09D167D2F44CA20009C51A3 /* PBXContainerItemProxy */;
|
||||
};
|
||||
E09D16882F44CA20009C51A3 /* PBXTargetDependency */ = {
|
||||
isa = PBXTargetDependency;
|
||||
target = E09D166E2F44CA1E009C51A3 /* EXO-iOS */;
|
||||
targetProxy = E09D16872F44CA20009C51A3 /* PBXContainerItemProxy */;
|
||||
};
|
||||
/* End PBXTargetDependency section */
|
||||
|
||||
/* Begin XCBuildConfiguration section */
|
||||
E09D168E2F44CA20009C51A3 /* Debug */ = {
|
||||
isa = XCBuildConfiguration;
|
||||
buildSettings = {
|
||||
ALWAYS_SEARCH_USER_PATHS = NO;
|
||||
ASSETCATALOG_COMPILER_GENERATE_SWIFT_ASSET_SYMBOL_EXTENSIONS = YES;
|
||||
CLANG_ANALYZER_NONNULL = YES;
|
||||
CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE;
|
||||
CLANG_CXX_LANGUAGE_STANDARD = "gnu++20";
|
||||
CLANG_ENABLE_MODULES = YES;
|
||||
CLANG_ENABLE_OBJC_ARC = YES;
|
||||
CLANG_ENABLE_OBJC_WEAK = YES;
|
||||
CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES;
|
||||
CLANG_WARN_BOOL_CONVERSION = YES;
|
||||
CLANG_WARN_COMMA = YES;
|
||||
CLANG_WARN_CONSTANT_CONVERSION = YES;
|
||||
CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES;
|
||||
CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR;
|
||||
CLANG_WARN_DOCUMENTATION_COMMENTS = YES;
|
||||
CLANG_WARN_EMPTY_BODY = YES;
|
||||
CLANG_WARN_ENUM_CONVERSION = YES;
|
||||
CLANG_WARN_INFINITE_RECURSION = YES;
|
||||
CLANG_WARN_INT_CONVERSION = YES;
|
||||
CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES;
|
||||
CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;
|
||||
CLANG_WARN_OBJC_LITERAL_CONVERSION = YES;
|
||||
CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;
|
||||
CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;
|
||||
CLANG_WARN_RANGE_LOOP_ANALYSIS = YES;
|
||||
CLANG_WARN_STRICT_PROTOTYPES = YES;
|
||||
CLANG_WARN_SUSPICIOUS_MOVE = YES;
|
||||
CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;
|
||||
CLANG_WARN_UNREACHABLE_CODE = YES;
|
||||
CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
|
||||
COPY_PHASE_STRIP = NO;
|
||||
DEBUG_INFORMATION_FORMAT = dwarf;
|
||||
ENABLE_STRICT_OBJC_MSGSEND = YES;
|
||||
ENABLE_TESTABILITY = YES;
|
||||
ENABLE_USER_SCRIPT_SANDBOXING = YES;
|
||||
GCC_C_LANGUAGE_STANDARD = gnu17;
|
||||
GCC_DYNAMIC_NO_PIC = NO;
|
||||
GCC_NO_COMMON_BLOCKS = YES;
|
||||
GCC_OPTIMIZATION_LEVEL = 0;
|
||||
GCC_PREPROCESSOR_DEFINITIONS = (
|
||||
"DEBUG=1",
|
||||
"$(inherited)",
|
||||
);
|
||||
GCC_WARN_64_TO_32_BIT_CONVERSION = YES;
|
||||
GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR;
|
||||
GCC_WARN_UNDECLARED_SELECTOR = YES;
|
||||
GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE;
|
||||
GCC_WARN_UNUSED_FUNCTION = YES;
|
||||
GCC_WARN_UNUSED_VARIABLE = YES;
|
||||
IPHONEOS_DEPLOYMENT_TARGET = 26.2;
|
||||
LOCALIZATION_PREFERS_STRING_CATALOGS = YES;
|
||||
MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE;
|
||||
MTL_FAST_MATH = YES;
|
||||
ONLY_ACTIVE_ARCH = YES;
|
||||
SDKROOT = iphoneos;
|
||||
SWIFT_ACTIVE_COMPILATION_CONDITIONS = "DEBUG $(inherited)";
|
||||
SWIFT_OPTIMIZATION_LEVEL = "-Onone";
|
||||
};
|
||||
name = Debug;
|
||||
};
|
||||
E09D168F2F44CA20009C51A3 /* Release */ = {
|
||||
isa = XCBuildConfiguration;
|
||||
buildSettings = {
|
||||
ALWAYS_SEARCH_USER_PATHS = NO;
|
||||
ASSETCATALOG_COMPILER_GENERATE_SWIFT_ASSET_SYMBOL_EXTENSIONS = YES;
|
||||
CLANG_ANALYZER_NONNULL = YES;
|
||||
CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE;
|
||||
CLANG_CXX_LANGUAGE_STANDARD = "gnu++20";
|
||||
CLANG_ENABLE_MODULES = YES;
|
||||
CLANG_ENABLE_OBJC_ARC = YES;
|
||||
CLANG_ENABLE_OBJC_WEAK = YES;
|
||||
CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES;
|
||||
CLANG_WARN_BOOL_CONVERSION = YES;
|
||||
CLANG_WARN_COMMA = YES;
|
||||
CLANG_WARN_CONSTANT_CONVERSION = YES;
|
||||
CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES;
|
||||
CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR;
|
||||
CLANG_WARN_DOCUMENTATION_COMMENTS = YES;
|
||||
CLANG_WARN_EMPTY_BODY = YES;
|
||||
CLANG_WARN_ENUM_CONVERSION = YES;
|
||||
CLANG_WARN_INFINITE_RECURSION = YES;
|
||||
CLANG_WARN_INT_CONVERSION = YES;
|
||||
CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES;
|
||||
CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;
|
||||
CLANG_WARN_OBJC_LITERAL_CONVERSION = YES;
|
||||
CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;
|
||||
CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;
|
||||
CLANG_WARN_RANGE_LOOP_ANALYSIS = YES;
|
||||
CLANG_WARN_STRICT_PROTOTYPES = YES;
|
||||
CLANG_WARN_SUSPICIOUS_MOVE = YES;
|
||||
CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;
|
||||
CLANG_WARN_UNREACHABLE_CODE = YES;
|
||||
CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
|
||||
COPY_PHASE_STRIP = NO;
|
||||
DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
|
||||
ENABLE_NS_ASSERTIONS = NO;
|
||||
ENABLE_STRICT_OBJC_MSGSEND = YES;
|
||||
ENABLE_USER_SCRIPT_SANDBOXING = YES;
|
||||
GCC_C_LANGUAGE_STANDARD = gnu17;
|
||||
GCC_NO_COMMON_BLOCKS = YES;
|
||||
GCC_WARN_64_TO_32_BIT_CONVERSION = YES;
|
||||
GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR;
|
||||
GCC_WARN_UNDECLARED_SELECTOR = YES;
|
||||
GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE;
|
||||
GCC_WARN_UNUSED_FUNCTION = YES;
|
||||
GCC_WARN_UNUSED_VARIABLE = YES;
|
||||
IPHONEOS_DEPLOYMENT_TARGET = 26.2;
|
||||
LOCALIZATION_PREFERS_STRING_CATALOGS = YES;
|
||||
MTL_ENABLE_DEBUG_INFO = NO;
|
||||
MTL_FAST_MATH = YES;
|
||||
SDKROOT = iphoneos;
|
||||
SWIFT_COMPILATION_MODE = wholemodule;
|
||||
VALIDATE_PRODUCT = YES;
|
||||
};
|
||||
name = Release;
|
||||
};
|
||||
E09D16912F44CA20009C51A3 /* Debug */ = {
|
||||
isa = XCBuildConfiguration;
|
||||
buildSettings = {
|
||||
ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
|
||||
ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor;
|
||||
CODE_SIGN_STYLE = Automatic;
|
||||
CURRENT_PROJECT_VERSION = 1;
|
||||
DEVELOPMENT_TEAM = 3M3M67U93M;
|
||||
ENABLE_PREVIEWS = YES;
|
||||
GENERATE_INFOPLIST_FILE = YES;
|
||||
INFOPLIST_FILE = "EXO-iOS/Info.plist";
|
||||
INFOPLIST_KEY_UIApplicationSceneManifest_Generation = YES;
|
||||
INFOPLIST_KEY_UIApplicationSupportsIndirectInputEvents = YES;
|
||||
INFOPLIST_KEY_UILaunchScreen_Generation = YES;
|
||||
INFOPLIST_KEY_UISupportedInterfaceOrientations_iPad = "UIInterfaceOrientationPortrait UIInterfaceOrientationPortraitUpsideDown UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight";
|
||||
INFOPLIST_KEY_UISupportedInterfaceOrientations_iPhone = "UIInterfaceOrientationPortrait UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight";
|
||||
LD_RUNPATH_SEARCH_PATHS = (
|
||||
"$(inherited)",
|
||||
"@executable_path/Frameworks",
|
||||
);
|
||||
MARKETING_VERSION = 1.0;
|
||||
PRODUCT_BUNDLE_IDENTIFIER = "com.exo.EXO-iOS";
|
||||
PRODUCT_NAME = "$(TARGET_NAME)";
|
||||
STRING_CATALOG_GENERATE_SYMBOLS = YES;
|
||||
SWIFT_APPROACHABLE_CONCURRENCY = YES;
|
||||
SWIFT_DEFAULT_ACTOR_ISOLATION = MainActor;
|
||||
SWIFT_EMIT_LOC_STRINGS = YES;
|
||||
SWIFT_UPCOMING_FEATURE_MEMBER_IMPORT_VISIBILITY = YES;
|
||||
SWIFT_VERSION = 5.0;
|
||||
TARGETED_DEVICE_FAMILY = "1,2";
|
||||
};
|
||||
name = Debug;
|
||||
};
|
||||
E09D16922F44CA20009C51A3 /* Release */ = {
|
||||
isa = XCBuildConfiguration;
|
||||
buildSettings = {
|
||||
ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
|
||||
ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor;
|
||||
CODE_SIGN_STYLE = Automatic;
|
||||
CURRENT_PROJECT_VERSION = 1;
|
||||
DEVELOPMENT_TEAM = 3M3M67U93M;
|
||||
ENABLE_PREVIEWS = YES;
|
||||
GENERATE_INFOPLIST_FILE = YES;
|
||||
INFOPLIST_FILE = "EXO-iOS/Info.plist";
|
||||
INFOPLIST_KEY_UIApplicationSceneManifest_Generation = YES;
|
||||
INFOPLIST_KEY_UIApplicationSupportsIndirectInputEvents = YES;
|
||||
INFOPLIST_KEY_UILaunchScreen_Generation = YES;
|
||||
INFOPLIST_KEY_UISupportedInterfaceOrientations_iPad = "UIInterfaceOrientationPortrait UIInterfaceOrientationPortraitUpsideDown UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight";
|
||||
INFOPLIST_KEY_UISupportedInterfaceOrientations_iPhone = "UIInterfaceOrientationPortrait UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight";
|
||||
LD_RUNPATH_SEARCH_PATHS = (
|
||||
"$(inherited)",
|
||||
"@executable_path/Frameworks",
|
||||
);
|
||||
MARKETING_VERSION = 1.0;
|
||||
PRODUCT_BUNDLE_IDENTIFIER = "com.exo.EXO-iOS";
|
||||
PRODUCT_NAME = "$(TARGET_NAME)";
|
||||
STRING_CATALOG_GENERATE_SYMBOLS = YES;
|
||||
SWIFT_APPROACHABLE_CONCURRENCY = YES;
|
||||
SWIFT_DEFAULT_ACTOR_ISOLATION = MainActor;
|
||||
SWIFT_EMIT_LOC_STRINGS = YES;
|
||||
SWIFT_UPCOMING_FEATURE_MEMBER_IMPORT_VISIBILITY = YES;
|
||||
SWIFT_VERSION = 5.0;
|
||||
TARGETED_DEVICE_FAMILY = "1,2";
|
||||
};
|
||||
name = Release;
|
||||
};
|
||||
E09D16942F44CA20009C51A3 /* Debug */ = {
|
||||
isa = XCBuildConfiguration;
|
||||
buildSettings = {
|
||||
BUNDLE_LOADER = "$(TEST_HOST)";
|
||||
CODE_SIGN_STYLE = Automatic;
|
||||
CURRENT_PROJECT_VERSION = 1;
|
||||
GENERATE_INFOPLIST_FILE = YES;
|
||||
IPHONEOS_DEPLOYMENT_TARGET = 26.2;
|
||||
MARKETING_VERSION = 1.0;
|
||||
PRODUCT_BUNDLE_IDENTIFIER = "com.exo.EXO-iOSTests";
|
||||
PRODUCT_NAME = "$(TARGET_NAME)";
|
||||
STRING_CATALOG_GENERATE_SYMBOLS = NO;
|
||||
SWIFT_APPROACHABLE_CONCURRENCY = YES;
|
||||
SWIFT_EMIT_LOC_STRINGS = NO;
|
||||
SWIFT_UPCOMING_FEATURE_MEMBER_IMPORT_VISIBILITY = YES;
|
||||
SWIFT_VERSION = 5.0;
|
||||
TARGETED_DEVICE_FAMILY = "1,2";
|
||||
TEST_HOST = "$(BUILT_PRODUCTS_DIR)/EXO-iOS.app/$(BUNDLE_EXECUTABLE_FOLDER_PATH)/EXO-iOS";
|
||||
};
|
||||
name = Debug;
|
||||
};
|
||||
E09D16952F44CA20009C51A3 /* Release */ = {
|
||||
isa = XCBuildConfiguration;
|
||||
buildSettings = {
|
||||
BUNDLE_LOADER = "$(TEST_HOST)";
|
||||
CODE_SIGN_STYLE = Automatic;
|
||||
CURRENT_PROJECT_VERSION = 1;
|
||||
GENERATE_INFOPLIST_FILE = YES;
|
||||
IPHONEOS_DEPLOYMENT_TARGET = 26.2;
|
||||
MARKETING_VERSION = 1.0;
|
||||
PRODUCT_BUNDLE_IDENTIFIER = "com.exo.EXO-iOSTests";
|
||||
PRODUCT_NAME = "$(TARGET_NAME)";
|
||||
STRING_CATALOG_GENERATE_SYMBOLS = NO;
|
||||
SWIFT_APPROACHABLE_CONCURRENCY = YES;
|
||||
SWIFT_EMIT_LOC_STRINGS = NO;
|
||||
SWIFT_UPCOMING_FEATURE_MEMBER_IMPORT_VISIBILITY = YES;
|
||||
SWIFT_VERSION = 5.0;
|
||||
TARGETED_DEVICE_FAMILY = "1,2";
|
||||
TEST_HOST = "$(BUILT_PRODUCTS_DIR)/EXO-iOS.app/$(BUNDLE_EXECUTABLE_FOLDER_PATH)/EXO-iOS";
|
||||
};
|
||||
name = Release;
|
||||
};
|
||||
E09D16972F44CA20009C51A3 /* Debug */ = {
|
||||
isa = XCBuildConfiguration;
|
||||
buildSettings = {
|
||||
CODE_SIGN_STYLE = Automatic;
|
||||
CURRENT_PROJECT_VERSION = 1;
|
||||
GENERATE_INFOPLIST_FILE = YES;
|
||||
MARKETING_VERSION = 1.0;
|
||||
PRODUCT_BUNDLE_IDENTIFIER = "com.exo.EXO-iOSUITests";
|
||||
PRODUCT_NAME = "$(TARGET_NAME)";
|
||||
STRING_CATALOG_GENERATE_SYMBOLS = NO;
|
||||
SWIFT_APPROACHABLE_CONCURRENCY = YES;
|
||||
SWIFT_EMIT_LOC_STRINGS = NO;
|
||||
SWIFT_UPCOMING_FEATURE_MEMBER_IMPORT_VISIBILITY = YES;
|
||||
SWIFT_VERSION = 5.0;
|
||||
TARGETED_DEVICE_FAMILY = "1,2";
|
||||
TEST_TARGET_NAME = "EXO-iOS";
|
||||
};
|
||||
name = Debug;
|
||||
};
|
||||
E09D16982F44CA20009C51A3 /* Release */ = {
|
||||
isa = XCBuildConfiguration;
|
||||
buildSettings = {
|
||||
CODE_SIGN_STYLE = Automatic;
|
||||
CURRENT_PROJECT_VERSION = 1;
|
||||
GENERATE_INFOPLIST_FILE = YES;
|
||||
MARKETING_VERSION = 1.0;
|
||||
PRODUCT_BUNDLE_IDENTIFIER = "com.exo.EXO-iOSUITests";
|
||||
PRODUCT_NAME = "$(TARGET_NAME)";
|
||||
STRING_CATALOG_GENERATE_SYMBOLS = NO;
|
||||
SWIFT_APPROACHABLE_CONCURRENCY = YES;
|
||||
SWIFT_EMIT_LOC_STRINGS = NO;
|
||||
SWIFT_UPCOMING_FEATURE_MEMBER_IMPORT_VISIBILITY = YES;
|
||||
SWIFT_VERSION = 5.0;
|
||||
TARGETED_DEVICE_FAMILY = "1,2";
|
||||
TEST_TARGET_NAME = "EXO-iOS";
|
||||
};
|
||||
name = Release;
|
||||
};
|
||||
/* End XCBuildConfiguration section */
|
||||
|
||||
/* Begin XCConfigurationList section */
|
||||
E09D166A2F44CA1E009C51A3 /* Build configuration list for PBXProject "EXO-iOS" */ = {
|
||||
isa = XCConfigurationList;
|
||||
buildConfigurations = (
|
||||
E09D168E2F44CA20009C51A3 /* Debug */,
|
||||
E09D168F2F44CA20009C51A3 /* Release */,
|
||||
);
|
||||
defaultConfigurationIsVisible = 0;
|
||||
defaultConfigurationName = Release;
|
||||
};
|
||||
E09D16902F44CA20009C51A3 /* Build configuration list for PBXNativeTarget "EXO-iOS" */ = {
|
||||
isa = XCConfigurationList;
|
||||
buildConfigurations = (
|
||||
E09D16912F44CA20009C51A3 /* Debug */,
|
||||
E09D16922F44CA20009C51A3 /* Release */,
|
||||
);
|
||||
defaultConfigurationIsVisible = 0;
|
||||
defaultConfigurationName = Release;
|
||||
};
|
||||
E09D16932F44CA20009C51A3 /* Build configuration list for PBXNativeTarget "EXO-iOSTests" */ = {
|
||||
isa = XCConfigurationList;
|
||||
buildConfigurations = (
|
||||
E09D16942F44CA20009C51A3 /* Debug */,
|
||||
E09D16952F44CA20009C51A3 /* Release */,
|
||||
);
|
||||
defaultConfigurationIsVisible = 0;
|
||||
defaultConfigurationName = Release;
|
||||
};
|
||||
E09D16962F44CA20009C51A3 /* Build configuration list for PBXNativeTarget "EXO-iOSUITests" */ = {
|
||||
isa = XCConfigurationList;
|
||||
buildConfigurations = (
|
||||
E09D16972F44CA20009C51A3 /* Debug */,
|
||||
E09D16982F44CA20009C51A3 /* Release */,
|
||||
);
|
||||
defaultConfigurationIsVisible = 0;
|
||||
defaultConfigurationName = Release;
|
||||
};
|
||||
/* End XCConfigurationList section */
|
||||
|
||||
/* Begin XCRemoteSwiftPackageReference section */
|
||||
E09D17502F44F359009C51A3 /* XCRemoteSwiftPackageReference "mlx-swift-lm" */ = {
|
||||
isa = XCRemoteSwiftPackageReference;
|
||||
repositoryURL = "https://github.com/ml-explore/mlx-swift-lm";
|
||||
requirement = {
|
||||
kind = upToNextMajorVersion;
|
||||
minimumVersion = 2.30.3;
|
||||
};
|
||||
};
|
||||
/* End XCRemoteSwiftPackageReference section */
|
||||
|
||||
/* Begin XCSwiftPackageProductDependency section */
|
||||
E09D17512F44F359009C51A3 /* MLXLLM */ = {
|
||||
isa = XCSwiftPackageProductDependency;
|
||||
package = E09D17502F44F359009C51A3 /* XCRemoteSwiftPackageReference "mlx-swift-lm" */;
|
||||
productName = MLXLLM;
|
||||
};
|
||||
E09D17532F44F359009C51A3 /* MLXLMCommon */ = {
|
||||
isa = XCSwiftPackageProductDependency;
|
||||
package = E09D17502F44F359009C51A3 /* XCRemoteSwiftPackageReference "mlx-swift-lm" */;
|
||||
productName = MLXLMCommon;
|
||||
};
|
||||
/* End XCSwiftPackageProductDependency section */
|
||||
};
|
||||
rootObject = E09D16672F44CA1E009C51A3 /* Project object */;
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<Workspace
|
||||
version = "1.0">
|
||||
<FileRef
|
||||
location = "self:">
|
||||
</FileRef>
|
||||
</Workspace>
|
||||
+60
@@ -0,0 +1,60 @@
|
||||
{
|
||||
"originHash" : "facc0ac7c70363ea20f6cd1235de91dea6b06f0d00190946045a6c8ae753abc2",
|
||||
"pins" : [
|
||||
{
|
||||
"identity" : "mlx-swift",
|
||||
"kind" : "remoteSourceControl",
|
||||
"location" : "https://github.com/ml-explore/mlx-swift",
|
||||
"state" : {
|
||||
"revision" : "6ba4827fb82c97d012eec9ab4b2de21f85c3b33d",
|
||||
"version" : "0.30.6"
|
||||
}
|
||||
},
|
||||
{
|
||||
"identity" : "mlx-swift-lm",
|
||||
"kind" : "remoteSourceControl",
|
||||
"location" : "https://github.com/ml-explore/mlx-swift-lm",
|
||||
"state" : {
|
||||
"revision" : "360c5052b81cc154b04ee0933597a4ad6db4b8ae",
|
||||
"version" : "2.30.3"
|
||||
}
|
||||
},
|
||||
{
|
||||
"identity" : "swift-collections",
|
||||
"kind" : "remoteSourceControl",
|
||||
"location" : "https://github.com/apple/swift-collections.git",
|
||||
"state" : {
|
||||
"revision" : "7b847a3b7008b2dc2f47ca3110d8c782fb2e5c7e",
|
||||
"version" : "1.3.0"
|
||||
}
|
||||
},
|
||||
{
|
||||
"identity" : "swift-jinja",
|
||||
"kind" : "remoteSourceControl",
|
||||
"location" : "https://github.com/huggingface/swift-jinja.git",
|
||||
"state" : {
|
||||
"revision" : "d81197f35f41445bc10e94600795e68c6f5e94b0",
|
||||
"version" : "2.3.1"
|
||||
}
|
||||
},
|
||||
{
|
||||
"identity" : "swift-numerics",
|
||||
"kind" : "remoteSourceControl",
|
||||
"location" : "https://github.com/apple/swift-numerics",
|
||||
"state" : {
|
||||
"revision" : "0c0290ff6b24942dadb83a929ffaaa1481df04a2",
|
||||
"version" : "1.1.1"
|
||||
}
|
||||
},
|
||||
{
|
||||
"identity" : "swift-transformers",
|
||||
"kind" : "remoteSourceControl",
|
||||
"location" : "https://github.com/huggingface/swift-transformers",
|
||||
"state" : {
|
||||
"revision" : "573e5c9036c2f136b3a8a071da8e8907322403d0",
|
||||
"version" : "1.1.6"
|
||||
}
|
||||
}
|
||||
],
|
||||
"version" : 3
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
{
|
||||
"colors" : [
|
||||
{
|
||||
"color" : {
|
||||
"color-space" : "srgb",
|
||||
"components" : {
|
||||
"alpha" : "1.000",
|
||||
"blue" : "0x00",
|
||||
"green" : "0xD7",
|
||||
"red" : "0xFF"
|
||||
}
|
||||
},
|
||||
"idiom" : "universal"
|
||||
}
|
||||
],
|
||||
"info" : {
|
||||
"author" : "xcode",
|
||||
"version" : 1
|
||||
}
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 10 KiB |
@@ -0,0 +1,38 @@
|
||||
{
|
||||
"images" : [
|
||||
{
|
||||
"filename" : "AppIcon.png",
|
||||
"idiom" : "universal",
|
||||
"platform" : "ios",
|
||||
"size" : "1024x1024"
|
||||
},
|
||||
{
|
||||
"appearances" : [
|
||||
{
|
||||
"appearance" : "luminosity",
|
||||
"value" : "dark"
|
||||
}
|
||||
],
|
||||
"filename" : "AppIcon.png",
|
||||
"idiom" : "universal",
|
||||
"platform" : "ios",
|
||||
"size" : "1024x1024"
|
||||
},
|
||||
{
|
||||
"appearances" : [
|
||||
{
|
||||
"appearance" : "luminosity",
|
||||
"value" : "tinted"
|
||||
}
|
||||
],
|
||||
"filename" : "AppIcon.png",
|
||||
"idiom" : "universal",
|
||||
"platform" : "ios",
|
||||
"size" : "1024x1024"
|
||||
}
|
||||
],
|
||||
"info" : {
|
||||
"author" : "xcode",
|
||||
"version" : 1
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
{
|
||||
"info" : {
|
||||
"author" : "xcode",
|
||||
"version" : 1
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
{
|
||||
"images" : [
|
||||
{
|
||||
"filename" : "exo-logo.png",
|
||||
"idiom" : "universal",
|
||||
"scale" : "1x"
|
||||
},
|
||||
{
|
||||
"idiom" : "universal",
|
||||
"scale" : "2x"
|
||||
},
|
||||
{
|
||||
"idiom" : "universal",
|
||||
"scale" : "3x"
|
||||
}
|
||||
],
|
||||
"info" : {
|
||||
"author" : "xcode",
|
||||
"version" : 1
|
||||
}
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.6 KiB |
@@ -0,0 +1,8 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
|
||||
<plist version="1.0">
|
||||
<dict>
|
||||
<key>com.apple.developer.kernel.increased-memory-limit</key>
|
||||
<true/>
|
||||
</dict>
|
||||
</plist>
|
||||
@@ -0,0 +1,67 @@
|
||||
import SwiftUI
|
||||
import UIKit
|
||||
|
||||
@main
|
||||
struct EXO_iOSApp: App {
|
||||
@State private var clusterService = ClusterService()
|
||||
@State private var discoveryService = DiscoveryService()
|
||||
@State private var localInferenceService = LocalInferenceService()
|
||||
@State private var chatService: ChatService?
|
||||
|
||||
init() {
|
||||
let darkGray = UIColor(red: 0x1F / 255.0, green: 0x1F / 255.0, blue: 0x1F / 255.0, alpha: 1)
|
||||
let yellow = UIColor(red: 0xFF / 255.0, green: 0xD7 / 255.0, blue: 0x00 / 255.0, alpha: 1)
|
||||
|
||||
let navAppearance = UINavigationBarAppearance()
|
||||
navAppearance.configureWithOpaqueBackground()
|
||||
navAppearance.backgroundColor = darkGray
|
||||
navAppearance.titleTextAttributes = [
|
||||
.foregroundColor: yellow,
|
||||
.font: UIFont.monospacedSystemFont(ofSize: 17, weight: .semibold),
|
||||
]
|
||||
navAppearance.largeTitleTextAttributes = [
|
||||
.foregroundColor: yellow,
|
||||
.font: UIFont.monospacedSystemFont(ofSize: 34, weight: .bold),
|
||||
]
|
||||
|
||||
UINavigationBar.appearance().standardAppearance = navAppearance
|
||||
UINavigationBar.appearance().compactAppearance = navAppearance
|
||||
UINavigationBar.appearance().scrollEdgeAppearance = navAppearance
|
||||
UINavigationBar.appearance().tintColor = yellow
|
||||
}
|
||||
|
||||
var body: some Scene {
|
||||
WindowGroup {
|
||||
if let chatService {
|
||||
RootView()
|
||||
.environment(clusterService)
|
||||
.environment(discoveryService)
|
||||
.environment(chatService)
|
||||
.environment(localInferenceService)
|
||||
.preferredColorScheme(.dark)
|
||||
.task {
|
||||
await clusterService.attemptAutoReconnect()
|
||||
discoveryService.startBrowsing()
|
||||
await localInferenceService.prepareModel()
|
||||
}
|
||||
.onChange(of: discoveryService.discoveredClusters) { _, clusters in
|
||||
guard !clusterService.isConnected,
|
||||
case .disconnected = clusterService.connectionState,
|
||||
let first = clusters.first
|
||||
else { return }
|
||||
Task {
|
||||
await clusterService.connectToDiscoveredCluster(
|
||||
first, using: discoveryService)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
Color.exoBlack.onAppear {
|
||||
chatService = ChatService(
|
||||
clusterService: clusterService,
|
||||
localInferenceService: localInferenceService
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
|
||||
<plist version="1.0">
|
||||
<dict>
|
||||
<key>UIUserInterfaceStyle</key>
|
||||
<string>Dark</string>
|
||||
<key>CFBundleDisplayName</key>
|
||||
<string>EXO</string>
|
||||
<key>NSLocalNetworkUsageDescription</key>
|
||||
<string>EXO needs local network access to connect to your EXO cluster.</string>
|
||||
<key>NSBonjourServices</key>
|
||||
<array>
|
||||
<string>_exo._tcp</string>
|
||||
<string>_p2p._tcp</string>
|
||||
<string>_p2p._udp</string>
|
||||
<string>_libp2p._udp</string>
|
||||
</array>
|
||||
</dict>
|
||||
</plist>
|
||||
@@ -0,0 +1,129 @@
|
||||
import Foundation
|
||||
|
||||
// MARK: - Request
|
||||
|
||||
struct ChatCompletionRequest: Encodable {
|
||||
let model: String
|
||||
let messages: [ChatCompletionMessageParam]
|
||||
let stream: Bool
|
||||
let maxTokens: Int?
|
||||
let temperature: Double?
|
||||
|
||||
enum CodingKeys: String, CodingKey {
|
||||
case model, messages, stream, temperature
|
||||
case maxTokens = "max_tokens"
|
||||
}
|
||||
}
|
||||
|
||||
struct ChatCompletionMessageParam: Encodable {
|
||||
let role: String
|
||||
let content: String
|
||||
}
|
||||
|
||||
// MARK: - Streaming Response
|
||||
|
||||
struct ChatCompletionChunk: Decodable {
|
||||
let id: String
|
||||
let model: String?
|
||||
let choices: [StreamingChoice]
|
||||
let usage: ChunkUsage?
|
||||
|
||||
init(id: String, model: String?, choices: [StreamingChoice], usage: ChunkUsage?) {
|
||||
self.id = id
|
||||
self.model = model
|
||||
self.choices = choices
|
||||
self.usage = usage
|
||||
}
|
||||
}
|
||||
|
||||
struct StreamingChoice: Decodable {
|
||||
let index: Int
|
||||
let delta: Delta
|
||||
let finishReason: String?
|
||||
|
||||
enum CodingKeys: String, CodingKey {
|
||||
case index, delta
|
||||
case finishReason = "finish_reason"
|
||||
}
|
||||
|
||||
init(index: Int, delta: Delta, finishReason: String?) {
|
||||
self.index = index
|
||||
self.delta = delta
|
||||
self.finishReason = finishReason
|
||||
}
|
||||
}
|
||||
|
||||
struct Delta: Decodable {
|
||||
let role: String?
|
||||
let content: String?
|
||||
|
||||
init(role: String?, content: String?) {
|
||||
self.role = role
|
||||
self.content = content
|
||||
}
|
||||
}
|
||||
|
||||
struct ChunkUsage: Decodable {
|
||||
let promptTokens: Int?
|
||||
let completionTokens: Int?
|
||||
let totalTokens: Int?
|
||||
|
||||
enum CodingKeys: String, CodingKey {
|
||||
case promptTokens = "prompt_tokens"
|
||||
case completionTokens = "completion_tokens"
|
||||
case totalTokens = "total_tokens"
|
||||
}
|
||||
|
||||
init(promptTokens: Int?, completionTokens: Int?, totalTokens: Int?) {
|
||||
self.promptTokens = promptTokens
|
||||
self.completionTokens = completionTokens
|
||||
self.totalTokens = totalTokens
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Non-Streaming Response
|
||||
|
||||
struct ChatCompletionResponse: Decodable {
|
||||
let id: String
|
||||
let model: String?
|
||||
let choices: [ResponseChoice]
|
||||
}
|
||||
|
||||
struct ResponseChoice: Decodable {
|
||||
let index: Int
|
||||
let message: ResponseMessage
|
||||
let finishReason: String?
|
||||
|
||||
enum CodingKeys: String, CodingKey {
|
||||
case index, message
|
||||
case finishReason = "finish_reason"
|
||||
}
|
||||
}
|
||||
|
||||
struct ResponseMessage: Decodable {
|
||||
let role: String?
|
||||
let content: String?
|
||||
}
|
||||
|
||||
// MARK: - Models List
|
||||
|
||||
struct ModelListResponse: Decodable {
|
||||
let data: [ModelInfo]
|
||||
}
|
||||
|
||||
struct ModelInfo: Decodable, Identifiable {
|
||||
let id: String
|
||||
let name: String?
|
||||
}
|
||||
|
||||
// MARK: - Error
|
||||
|
||||
struct APIErrorResponse: Decodable {
|
||||
let error: APIErrorInfo
|
||||
}
|
||||
|
||||
struct APIErrorInfo: Decodable {
|
||||
let message: String
|
||||
let type: String?
|
||||
let code: Int?
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
import Foundation
|
||||
|
||||
struct ChatMessage: Identifiable, Equatable {
|
||||
let id: UUID
|
||||
let role: Role
|
||||
var content: String
|
||||
let timestamp: Date
|
||||
var isStreaming: Bool
|
||||
|
||||
enum Role: String, Codable {
|
||||
case user
|
||||
case assistant
|
||||
case system
|
||||
}
|
||||
|
||||
init(
|
||||
id: UUID = UUID(), role: Role, content: String, timestamp: Date = Date(),
|
||||
isStreaming: Bool = false
|
||||
) {
|
||||
self.id = id
|
||||
self.role = role
|
||||
self.content = content
|
||||
self.timestamp = timestamp
|
||||
self.isStreaming = isStreaming
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
import Foundation
|
||||
|
||||
struct ConnectionInfo: Codable, Equatable {
|
||||
let host: String
|
||||
let port: Int
|
||||
let nodeId: String?
|
||||
|
||||
var baseURL: URL { URL(string: "http://\(host):\(port)")! }
|
||||
|
||||
static let defaultPort = 52415
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
import Foundation
|
||||
|
||||
struct Conversation: Identifiable, Codable, Equatable {
|
||||
let id: UUID
|
||||
var title: String
|
||||
var messages: [StoredMessage]
|
||||
var modelId: String?
|
||||
let createdAt: Date
|
||||
|
||||
init(
|
||||
id: UUID = UUID(), title: String = "New Chat", messages: [StoredMessage] = [],
|
||||
modelId: String? = nil, createdAt: Date = Date()
|
||||
) {
|
||||
self.id = id
|
||||
self.title = title
|
||||
self.messages = messages
|
||||
self.modelId = modelId
|
||||
self.createdAt = createdAt
|
||||
}
|
||||
}
|
||||
|
||||
struct StoredMessage: Identifiable, Codable, Equatable {
|
||||
let id: UUID
|
||||
let role: String
|
||||
var content: String
|
||||
let timestamp: Date
|
||||
|
||||
init(id: UUID = UUID(), role: String, content: String, timestamp: Date = Date()) {
|
||||
self.id = id
|
||||
self.role = role
|
||||
self.content = content
|
||||
self.timestamp = timestamp
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,227 @@
|
||||
import Foundation
|
||||
|
||||
@Observable
|
||||
@MainActor
|
||||
final class ChatService {
|
||||
var conversations: [Conversation] = []
|
||||
var activeConversationId: UUID?
|
||||
private(set) var isGenerating: Bool = false
|
||||
private var currentGenerationTask: Task<Void, Never>?
|
||||
|
||||
private let clusterService: ClusterService
|
||||
private let localInferenceService: LocalInferenceService
|
||||
|
||||
var canSendMessage: Bool {
|
||||
clusterService.isConnected || localInferenceService.isAvailable
|
||||
}
|
||||
|
||||
var activeConversation: Conversation? {
|
||||
guard let id = activeConversationId else { return nil }
|
||||
return conversations.first { $0.id == id }
|
||||
}
|
||||
|
||||
var activeMessages: [ChatMessage] {
|
||||
guard let conversation = activeConversation else { return [] }
|
||||
return conversation.messages.map { stored in
|
||||
ChatMessage(
|
||||
id: stored.id,
|
||||
role: ChatMessage.Role(rawValue: stored.role) ?? .user,
|
||||
content: stored.content,
|
||||
timestamp: stored.timestamp
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
init(clusterService: ClusterService, localInferenceService: LocalInferenceService) {
|
||||
self.clusterService = clusterService
|
||||
self.localInferenceService = localInferenceService
|
||||
loadConversations()
|
||||
}
|
||||
|
||||
// MARK: - Conversation Management
|
||||
|
||||
func createConversation(modelId: String? = nil) {
|
||||
let conversation = Conversation(
|
||||
modelId: modelId ?? clusterService.availableModels.first?.id)
|
||||
conversations.insert(conversation, at: 0)
|
||||
activeConversationId = conversation.id
|
||||
saveConversations()
|
||||
}
|
||||
|
||||
func deleteConversation(id: UUID) {
|
||||
conversations.removeAll { $0.id == id }
|
||||
if activeConversationId == id {
|
||||
activeConversationId = conversations.first?.id
|
||||
}
|
||||
saveConversations()
|
||||
}
|
||||
|
||||
func setActiveConversation(id: UUID) {
|
||||
activeConversationId = id
|
||||
}
|
||||
|
||||
func setModelForActiveConversation(_ modelId: String) {
|
||||
guard let index = conversations.firstIndex(where: { $0.id == activeConversationId }) else {
|
||||
return
|
||||
}
|
||||
conversations[index].modelId = modelId
|
||||
saveConversations()
|
||||
}
|
||||
|
||||
// MARK: - Messaging
|
||||
|
||||
func sendMessage(_ text: String) {
|
||||
guard !text.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { return }
|
||||
|
||||
if activeConversation == nil {
|
||||
createConversation()
|
||||
}
|
||||
|
||||
guard let index = conversations.firstIndex(where: { $0.id == activeConversationId }) else {
|
||||
return
|
||||
}
|
||||
|
||||
let userMessage = StoredMessage(role: "user", content: text)
|
||||
conversations[index].messages.append(userMessage)
|
||||
|
||||
if conversations[index].title == "New Chat" {
|
||||
let preview = String(text.prefix(40))
|
||||
conversations[index].title = preview + (text.count > 40 ? "..." : "")
|
||||
}
|
||||
|
||||
let modelId: String
|
||||
if clusterService.isConnected {
|
||||
guard
|
||||
let clusterId = conversations[index].modelId
|
||||
?? clusterService.availableModels.first?.id
|
||||
else {
|
||||
let errorMessage = StoredMessage(
|
||||
role: "assistant", content: "No model selected. Please select a model first.")
|
||||
conversations[index].messages.append(errorMessage)
|
||||
saveConversations()
|
||||
return
|
||||
}
|
||||
modelId = clusterId
|
||||
} else if localInferenceService.isAvailable {
|
||||
modelId = localInferenceService.defaultModelId
|
||||
} else {
|
||||
let errorMessage = StoredMessage(
|
||||
role: "assistant",
|
||||
content: "Not connected to a cluster and local model is not available.")
|
||||
conversations[index].messages.append(errorMessage)
|
||||
saveConversations()
|
||||
return
|
||||
}
|
||||
|
||||
conversations[index].modelId = modelId
|
||||
|
||||
let assistantMessageId = UUID()
|
||||
let assistantMessage = StoredMessage(
|
||||
id: assistantMessageId, role: "assistant", content: "", timestamp: Date())
|
||||
conversations[index].messages.append(assistantMessage)
|
||||
|
||||
let messagesForAPI = conversations[index].messages.dropLast().map { stored in
|
||||
ChatCompletionMessageParam(role: stored.role, content: stored.content)
|
||||
}
|
||||
|
||||
let request = ChatCompletionRequest(
|
||||
model: modelId,
|
||||
messages: Array(messagesForAPI),
|
||||
stream: true,
|
||||
maxTokens: 4096,
|
||||
temperature: nil
|
||||
)
|
||||
|
||||
let conversationId = conversations[index].id
|
||||
|
||||
isGenerating = true
|
||||
currentGenerationTask = Task { [weak self] in
|
||||
guard let self else { return }
|
||||
await self.performStreaming(
|
||||
request: request, conversationId: conversationId,
|
||||
assistantMessageId: assistantMessageId)
|
||||
}
|
||||
|
||||
saveConversations()
|
||||
}
|
||||
|
||||
func cancelGeneration() {
|
||||
currentGenerationTask?.cancel()
|
||||
currentGenerationTask = nil
|
||||
localInferenceService.cancelGeneration()
|
||||
isGenerating = false
|
||||
}
|
||||
|
||||
// MARK: - Streaming
|
||||
|
||||
private func performStreaming(
|
||||
request: ChatCompletionRequest, conversationId: UUID, assistantMessageId: UUID
|
||||
) async {
|
||||
defer {
|
||||
isGenerating = false
|
||||
currentGenerationTask = nil
|
||||
saveConversations()
|
||||
}
|
||||
|
||||
do {
|
||||
let stream =
|
||||
clusterService.isConnected
|
||||
? clusterService.streamChatCompletion(request: request)
|
||||
: localInferenceService.streamChatCompletion(request: request)
|
||||
for try await chunk in stream {
|
||||
guard !Task.isCancelled else { return }
|
||||
guard let content = chunk.choices.first?.delta.content, !content.isEmpty else {
|
||||
continue
|
||||
}
|
||||
|
||||
if let convIndex = conversations.firstIndex(where: { $0.id == conversationId }),
|
||||
let msgIndex = conversations[convIndex].messages.firstIndex(where: {
|
||||
$0.id == assistantMessageId
|
||||
})
|
||||
{
|
||||
conversations[convIndex].messages[msgIndex].content += content
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
if !Task.isCancelled {
|
||||
if let convIndex = conversations.firstIndex(where: { $0.id == conversationId }),
|
||||
let msgIndex = conversations[convIndex].messages.firstIndex(where: {
|
||||
$0.id == assistantMessageId
|
||||
})
|
||||
{
|
||||
if conversations[convIndex].messages[msgIndex].content.isEmpty {
|
||||
conversations[convIndex].messages[msgIndex].content =
|
||||
"Error: \(error.localizedDescription)"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Persistence
|
||||
|
||||
private static var storageURL: URL {
|
||||
let documents = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask)
|
||||
.first!
|
||||
return documents.appendingPathComponent("exo_conversations.json")
|
||||
}
|
||||
|
||||
private func saveConversations() {
|
||||
do {
|
||||
let data = try JSONEncoder().encode(conversations)
|
||||
try data.write(to: Self.storageURL, options: .atomic)
|
||||
} catch {
|
||||
// Save failed silently
|
||||
}
|
||||
}
|
||||
|
||||
private func loadConversations() {
|
||||
do {
|
||||
let data = try Data(contentsOf: Self.storageURL)
|
||||
conversations = try JSONDecoder().decode([Conversation].self, from: data)
|
||||
activeConversationId = conversations.first?.id
|
||||
} catch {
|
||||
conversations = []
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,246 @@
|
||||
import Foundation
|
||||
|
||||
enum ConnectionState: Equatable {
|
||||
case disconnected
|
||||
case connecting
|
||||
case connected(ConnectionInfo)
|
||||
}
|
||||
|
||||
struct ModelOption: Identifiable, Equatable {
|
||||
let id: String
|
||||
let displayName: String
|
||||
}
|
||||
|
||||
@Observable
|
||||
@MainActor
|
||||
final class ClusterService {
|
||||
private(set) var connectionState: ConnectionState = .disconnected
|
||||
private(set) var availableModels: [ModelOption] = []
|
||||
private(set) var lastError: String?
|
||||
|
||||
private let session: URLSession
|
||||
private let decoder: JSONDecoder
|
||||
private var pollingTask: Task<Void, Never>?
|
||||
private var heartbeatTask: Task<Void, Never>?
|
||||
|
||||
private static let connectionInfoKey = "exo_last_connection_info"
|
||||
|
||||
var isConnected: Bool {
|
||||
if case .connected = connectionState { return true }
|
||||
return false
|
||||
}
|
||||
|
||||
var currentConnection: ConnectionInfo? {
|
||||
if case .connected(let info) = connectionState { return info }
|
||||
return nil
|
||||
}
|
||||
|
||||
init(session: URLSession = .shared) {
|
||||
self.session = session
|
||||
let decoder = JSONDecoder()
|
||||
self.decoder = decoder
|
||||
}
|
||||
|
||||
// MARK: - Connection
|
||||
|
||||
func connect(to info: ConnectionInfo) async {
|
||||
connectionState = .connecting
|
||||
lastError = nil
|
||||
|
||||
do {
|
||||
let url = info.baseURL.appendingPathComponent("node_id")
|
||||
var request = URLRequest(url: url)
|
||||
request.timeoutInterval = 5
|
||||
request.cachePolicy = .reloadIgnoringLocalCacheData
|
||||
let (_, response) = try await session.data(for: request)
|
||||
|
||||
guard let httpResponse = response as? HTTPURLResponse,
|
||||
(200..<300).contains(httpResponse.statusCode)
|
||||
else {
|
||||
throw URLError(.badServerResponse)
|
||||
}
|
||||
|
||||
connectionState = .connected(info)
|
||||
persistConnection(info)
|
||||
startPolling()
|
||||
startHeartbeat()
|
||||
await fetchModels(baseURL: info.baseURL)
|
||||
} catch {
|
||||
connectionState = .disconnected
|
||||
lastError = "Could not connect to \(info.host):\(info.port)"
|
||||
}
|
||||
}
|
||||
|
||||
func connectToDiscoveredCluster(
|
||||
_ cluster: DiscoveredCluster, using discoveryService: DiscoveryService
|
||||
) async {
|
||||
guard case .disconnected = connectionState else { return }
|
||||
connectionState = .connecting
|
||||
lastError = nil
|
||||
|
||||
guard let info = await discoveryService.resolve(cluster) else {
|
||||
connectionState = .disconnected
|
||||
lastError = "Could not resolve \(cluster.name)"
|
||||
return
|
||||
}
|
||||
connectionState = .disconnected // reset so connect() can proceed
|
||||
await connect(to: info)
|
||||
}
|
||||
|
||||
func disconnect() {
|
||||
stopPolling()
|
||||
stopHeartbeat()
|
||||
connectionState = .disconnected
|
||||
availableModels = []
|
||||
lastError = nil
|
||||
}
|
||||
|
||||
func attemptAutoReconnect() async {
|
||||
guard case .disconnected = connectionState,
|
||||
let info = loadPersistedConnection()
|
||||
else { return }
|
||||
await connect(to: info)
|
||||
}
|
||||
|
||||
// MARK: - Polling
|
||||
|
||||
private func startPolling(interval: TimeInterval = 2.0) {
|
||||
stopPolling()
|
||||
pollingTask = Task { [weak self] in
|
||||
while !Task.isCancelled {
|
||||
try? await Task.sleep(for: .seconds(interval))
|
||||
guard let self, !Task.isCancelled else { return }
|
||||
guard let connection = self.currentConnection else { return }
|
||||
await self.fetchModels(baseURL: connection.baseURL)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func stopPolling() {
|
||||
pollingTask?.cancel()
|
||||
pollingTask = nil
|
||||
}
|
||||
|
||||
// MARK: - Heartbeat
|
||||
|
||||
private func startHeartbeat(interval: TimeInterval = 10.0) {
|
||||
stopHeartbeat()
|
||||
heartbeatTask = Task { [weak self] in
|
||||
// Send immediately, then on interval
|
||||
if let self, let connection = self.currentConnection {
|
||||
await self.sendHeartbeat(baseURL: connection.baseURL)
|
||||
}
|
||||
while !Task.isCancelled {
|
||||
try? await Task.sleep(for: .seconds(interval))
|
||||
guard let self, !Task.isCancelled else { return }
|
||||
guard let connection = self.currentConnection else { return }
|
||||
await self.sendHeartbeat(baseURL: connection.baseURL)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func stopHeartbeat() {
|
||||
heartbeatTask?.cancel()
|
||||
heartbeatTask = nil
|
||||
}
|
||||
|
||||
private func sendHeartbeat(baseURL: URL) async {
|
||||
do {
|
||||
let deviceInfo = DeviceInfoService.gather()
|
||||
let url = baseURL.appendingPathComponent("v1/lite_node/heartbeat")
|
||||
var request = URLRequest(url: url)
|
||||
request.httpMethod = "POST"
|
||||
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
|
||||
request.timeoutInterval = 5
|
||||
request.httpBody = try JSONEncoder().encode(deviceInfo)
|
||||
let (_, response) = try await session.data(for: request)
|
||||
|
||||
if let httpResponse = response as? HTTPURLResponse,
|
||||
!(200..<300).contains(httpResponse.statusCode) {
|
||||
// Heartbeat failed silently — will retry on next interval
|
||||
}
|
||||
} catch {
|
||||
// Heartbeat failed silently — will retry on next interval
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - API
|
||||
|
||||
private func fetchModels(baseURL: URL) async {
|
||||
do {
|
||||
let url = baseURL.appendingPathComponent("models")
|
||||
var request = URLRequest(url: url)
|
||||
request.cachePolicy = .reloadIgnoringLocalCacheData
|
||||
let (data, response) = try await session.data(for: request)
|
||||
|
||||
guard let httpResponse = response as? HTTPURLResponse,
|
||||
(200..<300).contains(httpResponse.statusCode)
|
||||
else { return }
|
||||
|
||||
let list = try decoder.decode(ModelListResponse.self, from: data)
|
||||
availableModels = list.data.map {
|
||||
ModelOption(id: $0.id, displayName: $0.name ?? $0.id)
|
||||
}
|
||||
} catch {
|
||||
// Models fetch failed silently — will retry on next poll
|
||||
}
|
||||
}
|
||||
|
||||
func streamChatCompletion(request body: ChatCompletionRequest) -> AsyncThrowingStream<
|
||||
ChatCompletionChunk, Error
|
||||
> {
|
||||
AsyncThrowingStream { continuation in
|
||||
let task = Task { [weak self] in
|
||||
guard let self, let connection = self.currentConnection else {
|
||||
continuation.finish(throwing: URLError(.notConnectedToInternet))
|
||||
return
|
||||
}
|
||||
|
||||
do {
|
||||
let url = connection.baseURL.appendingPathComponent("v1/chat/completions")
|
||||
var request = URLRequest(url: url)
|
||||
request.httpMethod = "POST"
|
||||
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
|
||||
request.httpBody = try JSONEncoder().encode(body)
|
||||
|
||||
let (bytes, response) = try await self.session.bytes(for: request)
|
||||
|
||||
guard let httpResponse = response as? HTTPURLResponse,
|
||||
(200..<300).contains(httpResponse.statusCode)
|
||||
else {
|
||||
continuation.finish(throwing: URLError(.badServerResponse))
|
||||
return
|
||||
}
|
||||
|
||||
let parser = SSEStreamParser<ChatCompletionChunk>(
|
||||
bytes: bytes, decoder: self.decoder)
|
||||
for try await chunk in parser {
|
||||
continuation.yield(chunk)
|
||||
}
|
||||
continuation.finish()
|
||||
} catch {
|
||||
continuation.finish(throwing: error)
|
||||
}
|
||||
}
|
||||
|
||||
continuation.onTermination = { _ in
|
||||
task.cancel()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Persistence
|
||||
|
||||
private func persistConnection(_ info: ConnectionInfo) {
|
||||
if let data = try? JSONEncoder().encode(info) {
|
||||
UserDefaults.standard.set(data, forKey: Self.connectionInfoKey)
|
||||
}
|
||||
}
|
||||
|
||||
private func loadPersistedConnection() -> ConnectionInfo? {
|
||||
guard let data = UserDefaults.standard.data(forKey: Self.connectionInfoKey) else {
|
||||
return nil
|
||||
}
|
||||
return try? JSONDecoder().decode(ConnectionInfo.self, from: data)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
import Foundation
|
||||
import os
|
||||
import UIKit
|
||||
|
||||
struct DeviceInfo: Encodable {
|
||||
let nodeId: String
|
||||
let model: String
|
||||
let chip: String
|
||||
let osVersion: String
|
||||
let friendlyName: String
|
||||
let ramTotal: Int
|
||||
let ramAvailable: Int
|
||||
|
||||
enum CodingKeys: String, CodingKey {
|
||||
case nodeId = "node_id"
|
||||
case model
|
||||
case chip
|
||||
case osVersion = "os_version"
|
||||
case friendlyName = "friendly_name"
|
||||
case ramTotal = "ram_total"
|
||||
case ramAvailable = "ram_available"
|
||||
}
|
||||
}
|
||||
|
||||
enum DeviceInfoService {
|
||||
private static let liteNodeIdKey = "exo_lite_node_id"
|
||||
|
||||
static var liteNodeId: String {
|
||||
if let existing = UserDefaults.standard.string(forKey: liteNodeIdKey) {
|
||||
return existing
|
||||
}
|
||||
let newId = UUID().uuidString.lowercased()
|
||||
UserDefaults.standard.set(newId, forKey: liteNodeIdKey)
|
||||
return newId
|
||||
}
|
||||
|
||||
static func gather() -> DeviceInfo {
|
||||
let model = modelName()
|
||||
let chip = chipName(for: model)
|
||||
|
||||
let totalRam = Int(ProcessInfo.processInfo.physicalMemory)
|
||||
let availableRam = availableMemory()
|
||||
|
||||
return DeviceInfo(
|
||||
nodeId: liteNodeId,
|
||||
model: model,
|
||||
chip: chip,
|
||||
osVersion: UIDevice.current.systemVersion,
|
||||
friendlyName: UIDevice.current.name,
|
||||
ramTotal: totalRam,
|
||||
ramAvailable: availableRam
|
||||
)
|
||||
}
|
||||
|
||||
// MARK: - Private
|
||||
|
||||
private static func modelName() -> String {
|
||||
var systemInfo = utsname()
|
||||
uname(&systemInfo)
|
||||
let machine = withUnsafePointer(to: &systemInfo.machine) {
|
||||
$0.withMemoryRebound(to: CChar.self, capacity: 1) {
|
||||
String(cString: $0)
|
||||
}
|
||||
}
|
||||
return modelMapping[machine] ?? machine
|
||||
}
|
||||
|
||||
private static func chipName(for model: String) -> String {
|
||||
let lower = model.lowercased()
|
||||
if lower.contains("iphone 16 pro") || lower.contains("iphone 16 pro max") {
|
||||
return "Apple A18 Pro"
|
||||
} else if lower.contains("iphone 16") {
|
||||
return "Apple A18"
|
||||
} else if lower.contains("iphone 15 pro") || lower.contains("iphone 15 pro max") {
|
||||
return "Apple A17 Pro"
|
||||
} else if lower.contains("iphone 15") {
|
||||
return "Apple A16 Bionic"
|
||||
} else if lower.contains("iphone 14 pro") || lower.contains("iphone 14 pro max") {
|
||||
return "Apple A16 Bionic"
|
||||
} else if lower.contains("iphone 14") {
|
||||
return "Apple A15 Bionic"
|
||||
}
|
||||
return "Apple Silicon"
|
||||
}
|
||||
|
||||
private static func availableMemory() -> Int {
|
||||
return Int(os_proc_available_memory())
|
||||
}
|
||||
|
||||
private static let modelMapping: [String: String] = [
|
||||
// iPhone 16 series
|
||||
"iPhone17,1": "iPhone 16 Pro",
|
||||
"iPhone17,2": "iPhone 16 Pro Max",
|
||||
"iPhone17,3": "iPhone 16",
|
||||
"iPhone17,4": "iPhone 16 Plus",
|
||||
// iPhone 15 series
|
||||
"iPhone16,1": "iPhone 15 Pro",
|
||||
"iPhone16,2": "iPhone 15 Pro Max",
|
||||
"iPhone15,4": "iPhone 15",
|
||||
"iPhone15,5": "iPhone 15 Plus",
|
||||
// iPhone 14 series
|
||||
"iPhone15,2": "iPhone 14 Pro",
|
||||
"iPhone15,3": "iPhone 14 Pro Max",
|
||||
"iPhone14,7": "iPhone 14",
|
||||
"iPhone14,8": "iPhone 14 Plus",
|
||||
// iPhone 13 series
|
||||
"iPhone14,2": "iPhone 13 Pro",
|
||||
"iPhone14,3": "iPhone 13 Pro Max",
|
||||
"iPhone14,5": "iPhone 13",
|
||||
"iPhone14,4": "iPhone 13 mini",
|
||||
// Simulator
|
||||
"arm64": "iPhone (Simulator)",
|
||||
"x86_64": "iPhone (Simulator)",
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,123 @@
|
||||
import Foundation
|
||||
import Network
|
||||
import os
|
||||
|
||||
struct DiscoveredCluster: Identifiable, Equatable {
|
||||
let id: String
|
||||
let name: String
|
||||
let endpoint: NWEndpoint
|
||||
|
||||
static func == (lhs: DiscoveredCluster, rhs: DiscoveredCluster) -> Bool {
|
||||
lhs.id == rhs.id && lhs.name == rhs.name
|
||||
}
|
||||
}
|
||||
|
||||
@Observable
|
||||
@MainActor
|
||||
final class DiscoveryService {
|
||||
private(set) var discoveredClusters: [DiscoveredCluster] = []
|
||||
private(set) var isSearching = false
|
||||
|
||||
private var browser: NWBrowser?
|
||||
|
||||
func startBrowsing() {
|
||||
guard browser == nil else { return }
|
||||
|
||||
let browser = NWBrowser(for: .bonjour(type: "_exo._tcp", domain: nil), using: .tcp)
|
||||
|
||||
browser.stateUpdateHandler = { [weak self] state in
|
||||
guard let service = self else { return }
|
||||
Task { @MainActor in
|
||||
switch state {
|
||||
case .ready:
|
||||
service.isSearching = true
|
||||
case .failed, .cancelled:
|
||||
service.isSearching = false
|
||||
default:
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
browser.browseResultsChangedHandler = { [weak self] results, _ in
|
||||
guard let service = self else { return }
|
||||
Task { @MainActor in
|
||||
service.discoveredClusters = results.compactMap { result in
|
||||
guard case .service(let name, _, _, _) = result.endpoint else {
|
||||
return nil
|
||||
}
|
||||
return DiscoveredCluster(
|
||||
id: name,
|
||||
name: name,
|
||||
endpoint: result.endpoint
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
browser.start(queue: .main)
|
||||
self.browser = browser
|
||||
}
|
||||
|
||||
func stopBrowsing() {
|
||||
browser?.cancel()
|
||||
browser = nil
|
||||
isSearching = false
|
||||
discoveredClusters = []
|
||||
}
|
||||
|
||||
/// Resolve a discovered Bonjour endpoint to an IP address and port, then return a ConnectionInfo.
|
||||
func resolve(_ cluster: DiscoveredCluster) async -> ConnectionInfo? {
|
||||
await withCheckedContinuation { continuation in
|
||||
let didResume = OSAllocatedUnfairLock(initialState: false)
|
||||
let connection = NWConnection(to: cluster.endpoint, using: .tcp)
|
||||
connection.stateUpdateHandler = { state in
|
||||
guard
|
||||
didResume.withLock({
|
||||
guard !$0 else { return false }
|
||||
$0 = true
|
||||
return true
|
||||
})
|
||||
else { return }
|
||||
switch state {
|
||||
case .ready:
|
||||
if let innerEndpoint = connection.currentPath?.remoteEndpoint,
|
||||
case .hostPort(let host, let port) = innerEndpoint
|
||||
{
|
||||
var hostString: String
|
||||
switch host {
|
||||
case .ipv4(let addr):
|
||||
hostString = "\(addr)"
|
||||
case .ipv6(let addr):
|
||||
hostString = "\(addr)"
|
||||
case .name(let name, _):
|
||||
hostString = name
|
||||
@unknown default:
|
||||
hostString = "\(host)"
|
||||
}
|
||||
// Strip interface scope suffix (e.g. "%en0")
|
||||
if let pct = hostString.firstIndex(of: "%") {
|
||||
hostString = String(hostString[..<pct])
|
||||
}
|
||||
let info = ConnectionInfo(
|
||||
host: hostString,
|
||||
port: Int(port.rawValue),
|
||||
nodeId: nil
|
||||
)
|
||||
connection.cancel()
|
||||
continuation.resume(returning: info)
|
||||
} else {
|
||||
connection.cancel()
|
||||
continuation.resume(returning: nil)
|
||||
}
|
||||
case .failed, .cancelled:
|
||||
continuation.resume(returning: nil)
|
||||
default:
|
||||
// Not a terminal state — allow future callbacks
|
||||
didResume.withLock { $0 = false }
|
||||
}
|
||||
}
|
||||
connection.start(queue: .global(qos: .userInitiated))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,201 @@
|
||||
import Foundation
|
||||
import MLXLLM
|
||||
import MLXLMCommon
|
||||
|
||||
enum LocalModelState: Equatable {
|
||||
case notDownloaded
|
||||
case downloading(progress: Double)
|
||||
case downloaded
|
||||
case loading
|
||||
case ready
|
||||
case generating
|
||||
case error(String)
|
||||
}
|
||||
|
||||
@Observable
|
||||
@MainActor
|
||||
final class LocalInferenceService {
|
||||
private(set) var modelState: LocalModelState = .notDownloaded
|
||||
private var modelContainer: ModelContainer?
|
||||
private var generationTask: Task<Void, Never>?
|
||||
|
||||
let defaultModelId = "mlx-community/Qwen3-0.6B-4bit"
|
||||
|
||||
private static let modelDownloadedKey = "exo_local_model_downloaded"
|
||||
|
||||
var isReady: Bool {
|
||||
modelState == .ready
|
||||
}
|
||||
|
||||
var isAvailable: Bool {
|
||||
modelState == .ready || modelState == .generating
|
||||
}
|
||||
|
||||
init() {
|
||||
if UserDefaults.standard.bool(forKey: Self.modelDownloadedKey) {
|
||||
modelState = .downloaded
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Model Lifecycle
|
||||
|
||||
func prepareModel() async {
|
||||
guard modelState == .notDownloaded || modelState == .downloaded else { return }
|
||||
|
||||
let wasDownloaded = modelState == .downloaded
|
||||
|
||||
if !wasDownloaded {
|
||||
modelState = .downloading(progress: 0)
|
||||
} else {
|
||||
modelState = .loading
|
||||
}
|
||||
|
||||
do {
|
||||
let container = try await loadModelContainer(
|
||||
id: defaultModelId
|
||||
) { [weak self] progress in
|
||||
guard let self else { return }
|
||||
Task { @MainActor in
|
||||
if case .downloading = self.modelState {
|
||||
self.modelState = .downloading(progress: progress.fractionCompleted)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.modelContainer = container
|
||||
UserDefaults.standard.set(true, forKey: Self.modelDownloadedKey)
|
||||
modelState = .ready
|
||||
} catch {
|
||||
modelState = .error(error.localizedDescription)
|
||||
}
|
||||
}
|
||||
|
||||
func unloadModel() {
|
||||
cancelGeneration()
|
||||
modelContainer = nil
|
||||
modelState = .downloaded
|
||||
}
|
||||
|
||||
// MARK: - Generation
|
||||
|
||||
func streamChatCompletion(request: ChatCompletionRequest) -> AsyncThrowingStream<
|
||||
ChatCompletionChunk, Error
|
||||
> {
|
||||
AsyncThrowingStream { continuation in
|
||||
let task = Task { [weak self] in
|
||||
guard let self else {
|
||||
continuation.finish(throwing: LocalInferenceError.serviceUnavailable)
|
||||
return
|
||||
}
|
||||
|
||||
guard let container = self.modelContainer else {
|
||||
continuation.finish(throwing: LocalInferenceError.modelNotLoaded)
|
||||
return
|
||||
}
|
||||
|
||||
await MainActor.run {
|
||||
self.modelState = .generating
|
||||
}
|
||||
|
||||
defer {
|
||||
Task { @MainActor [weak self] in
|
||||
if self?.modelState == .generating {
|
||||
self?.modelState = .ready
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let chunkId = "local-\(UUID().uuidString)"
|
||||
|
||||
do {
|
||||
// Build Chat.Message array from the request
|
||||
var chatMessages: [Chat.Message] = []
|
||||
for msg in request.messages {
|
||||
switch msg.role {
|
||||
case "system":
|
||||
chatMessages.append(.system(msg.content))
|
||||
case "assistant":
|
||||
chatMessages.append(.assistant(msg.content))
|
||||
default:
|
||||
chatMessages.append(.user(msg.content))
|
||||
}
|
||||
}
|
||||
|
||||
// Use ChatSession for streaming generation
|
||||
let session = ChatSession(
|
||||
container,
|
||||
history: chatMessages,
|
||||
generateParameters: GenerateParameters(
|
||||
maxTokens: request.maxTokens ?? 4096,
|
||||
temperature: Float(request.temperature ?? 0.7)
|
||||
)
|
||||
)
|
||||
|
||||
// Stream with an empty prompt since history already contains the conversation
|
||||
let stream = session.streamResponse(to: "")
|
||||
for try await text in stream {
|
||||
if Task.isCancelled { break }
|
||||
|
||||
let chunk = ChatCompletionChunk(
|
||||
id: chunkId,
|
||||
model: request.model,
|
||||
choices: [
|
||||
StreamingChoice(
|
||||
index: 0,
|
||||
delta: Delta(role: nil, content: text),
|
||||
finishReason: nil
|
||||
)
|
||||
],
|
||||
usage: nil
|
||||
)
|
||||
continuation.yield(chunk)
|
||||
}
|
||||
|
||||
// Send final chunk with finish reason
|
||||
let finalChunk = ChatCompletionChunk(
|
||||
id: chunkId,
|
||||
model: request.model,
|
||||
choices: [
|
||||
StreamingChoice(
|
||||
index: 0,
|
||||
delta: Delta(role: nil, content: nil),
|
||||
finishReason: "stop"
|
||||
)
|
||||
],
|
||||
usage: nil
|
||||
)
|
||||
continuation.yield(finalChunk)
|
||||
continuation.finish()
|
||||
} catch {
|
||||
continuation.finish(throwing: error)
|
||||
}
|
||||
}
|
||||
|
||||
self.generationTask = task
|
||||
|
||||
continuation.onTermination = { _ in
|
||||
task.cancel()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func cancelGeneration() {
|
||||
generationTask?.cancel()
|
||||
generationTask = nil
|
||||
if modelState == .generating {
|
||||
modelState = .ready
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum LocalInferenceError: LocalizedError {
|
||||
case serviceUnavailable
|
||||
case modelNotLoaded
|
||||
|
||||
var errorDescription: String? {
|
||||
switch self {
|
||||
case .serviceUnavailable: "Local inference service is unavailable"
|
||||
case .modelNotLoaded: "Local model is not loaded"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
import Foundation
|
||||
|
||||
struct SSEStreamParser<T: Decodable>: AsyncSequence {
|
||||
typealias Element = T
|
||||
|
||||
let bytes: URLSession.AsyncBytes
|
||||
let decoder: JSONDecoder
|
||||
|
||||
init(bytes: URLSession.AsyncBytes, decoder: JSONDecoder = JSONDecoder()) {
|
||||
self.bytes = bytes
|
||||
self.decoder = decoder
|
||||
}
|
||||
|
||||
func makeAsyncIterator() -> AsyncIterator {
|
||||
AsyncIterator(lines: bytes.lines, decoder: decoder)
|
||||
}
|
||||
|
||||
struct AsyncIterator: AsyncIteratorProtocol {
|
||||
var lines: AsyncLineSequence<URLSession.AsyncBytes>.AsyncIterator
|
||||
let decoder: JSONDecoder
|
||||
|
||||
init(lines: AsyncLineSequence<URLSession.AsyncBytes>, decoder: JSONDecoder) {
|
||||
self.lines = lines.makeAsyncIterator()
|
||||
self.decoder = decoder
|
||||
}
|
||||
|
||||
mutating func next() async throws -> T? {
|
||||
while let line = try await lines.next() {
|
||||
let trimmed = line.trimmingCharacters(in: .whitespaces)
|
||||
|
||||
guard trimmed.hasPrefix("data: ") else { continue }
|
||||
|
||||
let payload = String(trimmed.dropFirst(6))
|
||||
|
||||
if payload == "[DONE]" {
|
||||
return nil
|
||||
}
|
||||
|
||||
guard let data = payload.data(using: .utf8) else { continue }
|
||||
|
||||
do {
|
||||
return try decoder.decode(T.self, from: data)
|
||||
} catch {
|
||||
continue
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
import SwiftUI
|
||||
|
||||
struct ChatView: View {
|
||||
@Environment(ClusterService.self) private var clusterService
|
||||
@Environment(ChatService.self) private var chatService
|
||||
@Environment(LocalInferenceService.self) private var localInferenceService
|
||||
@State private var inputText = ""
|
||||
@State private var showModelSelector = false
|
||||
|
||||
var body: some View {
|
||||
VStack(spacing: 0) {
|
||||
modelBar
|
||||
|
||||
GradientDivider()
|
||||
|
||||
messageList
|
||||
|
||||
GradientDivider()
|
||||
|
||||
inputBar
|
||||
}
|
||||
.background(Color.exoBlack)
|
||||
.sheet(isPresented: $showModelSelector) {
|
||||
ModelSelectorView(
|
||||
models: clusterService.availableModels,
|
||||
selectedModelId: chatService.activeConversation?.modelId
|
||||
) { modelId in
|
||||
chatService.setModelForActiveConversation(modelId)
|
||||
}
|
||||
.presentationBackground(Color.exoDarkGray)
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Model Bar
|
||||
|
||||
private var useLocalModel: Bool {
|
||||
!clusterService.isConnected && localInferenceService.isAvailable
|
||||
}
|
||||
|
||||
private var modelBar: some View {
|
||||
Button {
|
||||
if !useLocalModel {
|
||||
showModelSelector = true
|
||||
}
|
||||
} label: {
|
||||
HStack {
|
||||
Image(systemName: useLocalModel ? "iphone" : "cpu")
|
||||
.font(.exoCaption)
|
||||
.foregroundStyle(useLocalModel ? Color.exoYellow : Color.exoLightGray)
|
||||
|
||||
if useLocalModel {
|
||||
Text(localInferenceService.defaultModelId)
|
||||
.font(.exoSubheadline)
|
||||
.foregroundStyle(Color.exoForeground)
|
||||
.lineLimit(1)
|
||||
} else if let modelId = chatService.activeConversation?.modelId {
|
||||
Text(modelId)
|
||||
.font(.exoSubheadline)
|
||||
.foregroundStyle(Color.exoForeground)
|
||||
.lineLimit(1)
|
||||
} else {
|
||||
Text("SELECT MODEL")
|
||||
.font(.exoSubheadline)
|
||||
.tracking(1.5)
|
||||
.foregroundStyle(Color.exoLightGray)
|
||||
}
|
||||
|
||||
Spacer()
|
||||
|
||||
if useLocalModel {
|
||||
Text("ON-DEVICE")
|
||||
.font(.exoCaption)
|
||||
.tracking(1)
|
||||
.foregroundStyle(Color.exoYellow)
|
||||
.padding(.horizontal, 6)
|
||||
.padding(.vertical, 2)
|
||||
.background(Color.exoYellow.opacity(0.15))
|
||||
.clipShape(Capsule())
|
||||
} else {
|
||||
Image(systemName: "chevron.right")
|
||||
.font(.caption)
|
||||
.foregroundStyle(Color.exoLightGray)
|
||||
}
|
||||
}
|
||||
.padding(.horizontal)
|
||||
.padding(.vertical, 10)
|
||||
.background(Color.exoDarkGray)
|
||||
}
|
||||
.tint(.primary)
|
||||
.disabled(useLocalModel)
|
||||
}
|
||||
|
||||
// MARK: - Messages
|
||||
|
||||
private var messageList: some View {
|
||||
ScrollViewReader { proxy in
|
||||
ScrollView {
|
||||
LazyVStack(spacing: 12) {
|
||||
if chatService.activeMessages.isEmpty {
|
||||
emptyState
|
||||
} else {
|
||||
ForEach(chatService.activeMessages) { message in
|
||||
MessageBubbleView(message: message)
|
||||
.id(message.id)
|
||||
}
|
||||
}
|
||||
}
|
||||
.padding()
|
||||
}
|
||||
.background(Color.exoBlack)
|
||||
.onChange(of: chatService.activeMessages.last?.content) {
|
||||
if let lastId = chatService.activeMessages.last?.id {
|
||||
withAnimation(.easeOut(duration: 0.2)) {
|
||||
proxy.scrollTo(lastId, anchor: .bottom)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private var emptyState: some View {
|
||||
VStack(spacing: 16) {
|
||||
Spacer(minLength: 80)
|
||||
|
||||
ZStack {
|
||||
Circle()
|
||||
.stroke(Color.exoYellow.opacity(0.15), lineWidth: 1)
|
||||
.frame(width: 80, height: 80)
|
||||
Circle()
|
||||
.stroke(Color.exoYellow.opacity(0.3), lineWidth: 1)
|
||||
.frame(width: 56, height: 56)
|
||||
Circle()
|
||||
.fill(Color.exoYellow.opacity(0.15))
|
||||
.frame(width: 32, height: 32)
|
||||
Circle()
|
||||
.fill(Color.exoYellow)
|
||||
.frame(width: 8, height: 8)
|
||||
.shadow(color: Color.exoYellow.opacity(0.6), radius: 6)
|
||||
}
|
||||
|
||||
Text("AWAITING INPUT")
|
||||
.font(.exoSubheadline)
|
||||
.tracking(3)
|
||||
.foregroundStyle(Color.exoLightGray)
|
||||
|
||||
Text("Send a message to begin.")
|
||||
.font(.exoCaption)
|
||||
.foregroundStyle(Color.exoLightGray.opacity(0.6))
|
||||
|
||||
Spacer(minLength: 80)
|
||||
}
|
||||
.padding()
|
||||
}
|
||||
|
||||
// MARK: - Input
|
||||
|
||||
private var inputBar: some View {
|
||||
HStack(alignment: .bottom, spacing: 8) {
|
||||
TextField("Message...", text: $inputText, axis: .vertical)
|
||||
.font(.exoBody)
|
||||
.lineLimit(1...6)
|
||||
.textFieldStyle(.plain)
|
||||
.padding(10)
|
||||
.background(Color.exoMediumGray)
|
||||
.foregroundStyle(Color.exoForeground)
|
||||
.clipShape(RoundedRectangle(cornerRadius: 8))
|
||||
|
||||
if chatService.isGenerating {
|
||||
Button {
|
||||
chatService.cancelGeneration()
|
||||
} label: {
|
||||
Image(systemName: "stop.circle.fill")
|
||||
.font(.title2)
|
||||
.foregroundStyle(Color.exoDestructive)
|
||||
}
|
||||
} else {
|
||||
Button {
|
||||
let text = inputText
|
||||
inputText = ""
|
||||
chatService.sendMessage(text)
|
||||
} label: {
|
||||
Text("SEND")
|
||||
.font(.exoMono(12, weight: .bold))
|
||||
.tracking(1)
|
||||
.foregroundStyle(canSend ? Color.exoBlack : Color.exoLightGray)
|
||||
.padding(.horizontal, 14)
|
||||
.padding(.vertical, 8)
|
||||
.background(canSend ? Color.exoYellow : Color.exoMediumGray)
|
||||
.clipShape(RoundedRectangle(cornerRadius: 8))
|
||||
}
|
||||
.disabled(!canSend)
|
||||
}
|
||||
}
|
||||
.padding(.horizontal)
|
||||
.padding(.vertical, 8)
|
||||
.background(Color.exoDarkGray)
|
||||
}
|
||||
|
||||
private var canSend: Bool {
|
||||
!inputText.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
|
||||
&& (clusterService.isConnected || localInferenceService.isAvailable)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
import SwiftUI
|
||||
|
||||
struct MessageBubbleView: View {
|
||||
let message: ChatMessage
|
||||
|
||||
private var isAssistant: Bool { message.role == .assistant }
|
||||
|
||||
var body: some View {
|
||||
HStack {
|
||||
if message.role == .user { Spacer(minLength: 48) }
|
||||
|
||||
VStack(alignment: isAssistant ? .leading : .trailing, spacing: 6) {
|
||||
// Header
|
||||
HStack(spacing: 4) {
|
||||
if isAssistant {
|
||||
Circle()
|
||||
.fill(Color.exoYellow)
|
||||
.frame(width: 6, height: 6)
|
||||
.shadow(color: Color.exoYellow.opacity(0.6), radius: 4)
|
||||
Text("EXO")
|
||||
.font(.exoMono(10, weight: .bold))
|
||||
.tracking(1.5)
|
||||
.foregroundStyle(Color.exoYellow)
|
||||
} else {
|
||||
Text("QUERY")
|
||||
.font(.exoMono(10, weight: .medium))
|
||||
.tracking(1.5)
|
||||
.foregroundStyle(Color.exoLightGray)
|
||||
}
|
||||
}
|
||||
|
||||
// Bubble
|
||||
HStack(spacing: 0) {
|
||||
if isAssistant {
|
||||
RoundedRectangle(cornerRadius: 1)
|
||||
.fill(Color.exoYellow.opacity(0.5))
|
||||
.frame(width: 2)
|
||||
}
|
||||
|
||||
Text(message.content + (message.isStreaming ? " \u{258C}" : ""))
|
||||
.font(.exoBody)
|
||||
.textSelection(.enabled)
|
||||
.foregroundStyle(Color.exoForeground)
|
||||
.padding(.horizontal, 14)
|
||||
.padding(.vertical, 10)
|
||||
}
|
||||
.background(Color.exoDarkGray)
|
||||
.clipShape(RoundedRectangle(cornerRadius: 8))
|
||||
}
|
||||
|
||||
if isAssistant { Spacer(minLength: 48) }
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
import SwiftUI
|
||||
|
||||
struct ModelSelectorView: View {
|
||||
let models: [ModelOption]
|
||||
let selectedModelId: String?
|
||||
let onSelect: (String) -> Void
|
||||
@Environment(\.dismiss) private var dismiss
|
||||
|
||||
var body: some View {
|
||||
NavigationStack {
|
||||
List {
|
||||
if models.isEmpty {
|
||||
emptyContent
|
||||
} else {
|
||||
modelsList
|
||||
}
|
||||
}
|
||||
.scrollContentBackground(.hidden)
|
||||
.background(Color.exoBlack)
|
||||
.navigationTitle("SELECT MODEL")
|
||||
.navigationBarTitleDisplayMode(.inline)
|
||||
.toolbar {
|
||||
ToolbarItem(placement: .cancellationAction) {
|
||||
Button("Cancel") { dismiss() }
|
||||
.font(.exoSubheadline)
|
||||
.foregroundStyle(Color.exoYellow)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private var emptyContent: some View {
|
||||
ContentUnavailableView(
|
||||
"No Models Available",
|
||||
systemImage: "cpu",
|
||||
description: Text("Connect to an EXO cluster to see available models.")
|
||||
)
|
||||
.foregroundStyle(Color.exoLightGray)
|
||||
.listRowBackground(Color.exoBlack)
|
||||
}
|
||||
|
||||
private var modelsList: some View {
|
||||
ForEach(models) { model in
|
||||
Button {
|
||||
onSelect(model.id)
|
||||
dismiss()
|
||||
} label: {
|
||||
modelRow(model)
|
||||
}
|
||||
.tint(.primary)
|
||||
.listRowBackground(Color.exoDarkGray)
|
||||
}
|
||||
}
|
||||
|
||||
private func modelRow(_ model: ModelOption) -> some View {
|
||||
HStack {
|
||||
VStack(alignment: .leading, spacing: 2) {
|
||||
Text(model.displayName)
|
||||
.font(.exoSubheadline)
|
||||
.foregroundStyle(Color.exoForeground)
|
||||
Text(model.id)
|
||||
.font(.exoCaption)
|
||||
.foregroundStyle(Color.exoLightGray)
|
||||
}
|
||||
|
||||
Spacer()
|
||||
|
||||
if model.id == selectedModelId {
|
||||
Image(systemName: "checkmark")
|
||||
.font(.exoSubheadline)
|
||||
.foregroundStyle(Color.exoYellow)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
import SwiftUI
|
||||
|
||||
struct ConnectionStatusBadge: View {
|
||||
let connectionState: ConnectionState
|
||||
var localModelState: LocalModelState = .notDownloaded
|
||||
|
||||
private var isLocalReady: Bool {
|
||||
if case .disconnected = connectionState {
|
||||
return localModelState == .ready || localModelState == .generating
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
var body: some View {
|
||||
HStack(spacing: 6) {
|
||||
Circle()
|
||||
.fill(dotColor)
|
||||
.frame(width: 8, height: 8)
|
||||
.shadow(color: dotColor.opacity(0.6), radius: 4)
|
||||
|
||||
Text(label.uppercased())
|
||||
.font(.exoMono(10, weight: .medium))
|
||||
.tracking(1)
|
||||
.foregroundStyle(Color.exoForeground)
|
||||
}
|
||||
.padding(.horizontal, 10)
|
||||
.padding(.vertical, 5)
|
||||
.background(backgroundColor)
|
||||
.clipShape(Capsule())
|
||||
.overlay(
|
||||
Capsule()
|
||||
.stroke(dotColor.opacity(0.3), lineWidth: 1)
|
||||
)
|
||||
}
|
||||
|
||||
private var dotColor: Color {
|
||||
if isLocalReady {
|
||||
return .exoYellow
|
||||
}
|
||||
switch connectionState {
|
||||
case .connected: return .green
|
||||
case .connecting: return .orange
|
||||
case .disconnected: return .exoLightGray
|
||||
}
|
||||
}
|
||||
|
||||
private var label: String {
|
||||
if isLocalReady {
|
||||
return "Local"
|
||||
}
|
||||
switch connectionState {
|
||||
case .connected: return "Connected"
|
||||
case .connecting: return "Connecting"
|
||||
case .disconnected: return "Disconnected"
|
||||
}
|
||||
}
|
||||
|
||||
private var backgroundColor: Color {
|
||||
if isLocalReady {
|
||||
return Color.exoYellow.opacity(0.1)
|
||||
}
|
||||
switch connectionState {
|
||||
case .connected: return .green.opacity(0.1)
|
||||
case .connecting: return .orange.opacity(0.1)
|
||||
case .disconnected: return Color.exoMediumGray.opacity(0.5)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,136 @@
|
||||
import SwiftUI
|
||||
|
||||
struct RootView: View {
|
||||
@Environment(ClusterService.self) private var clusterService
|
||||
@Environment(DiscoveryService.self) private var discoveryService
|
||||
@Environment(ChatService.self) private var chatService
|
||||
@Environment(LocalInferenceService.self) private var localInferenceService
|
||||
@State private var showSettings = false
|
||||
@State private var showConversations = false
|
||||
|
||||
var body: some View {
|
||||
NavigationStack {
|
||||
ChatView()
|
||||
.navigationTitle("EXO")
|
||||
.navigationBarTitleDisplayMode(.inline)
|
||||
.toolbar {
|
||||
ToolbarItem(placement: .topBarLeading) {
|
||||
conversationMenuButton
|
||||
}
|
||||
|
||||
ToolbarItem(placement: .principal) {
|
||||
ConnectionStatusBadge(
|
||||
connectionState: clusterService.connectionState,
|
||||
localModelState: localInferenceService.modelState
|
||||
)
|
||||
}
|
||||
|
||||
ToolbarItem(placement: .topBarTrailing) {
|
||||
Button {
|
||||
showSettings = true
|
||||
} label: {
|
||||
Image(systemName: "gear")
|
||||
.foregroundStyle(Color.exoYellow)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
.tint(Color.exoYellow)
|
||||
.sheet(isPresented: $showSettings) {
|
||||
SettingsView()
|
||||
.environment(discoveryService)
|
||||
.presentationBackground(Color.exoDarkGray)
|
||||
}
|
||||
.sheet(isPresented: $showConversations) {
|
||||
conversationList
|
||||
.presentationBackground(Color.exoDarkGray)
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Conversations
|
||||
|
||||
private var conversationMenuButton: some View {
|
||||
HStack(spacing: 12) {
|
||||
Button {
|
||||
showConversations = true
|
||||
} label: {
|
||||
Image(systemName: "sidebar.left")
|
||||
.foregroundStyle(Color.exoYellow)
|
||||
}
|
||||
|
||||
Button {
|
||||
chatService.createConversation()
|
||||
} label: {
|
||||
Image(systemName: "square.and.pencil")
|
||||
.foregroundStyle(Color.exoYellow)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private var conversationList: some View {
|
||||
NavigationStack {
|
||||
List {
|
||||
if chatService.conversations.isEmpty {
|
||||
Text("No conversations yet")
|
||||
.font(.exoBody)
|
||||
.foregroundStyle(Color.exoLightGray)
|
||||
.listRowBackground(Color.exoDarkGray)
|
||||
} else {
|
||||
ForEach(chatService.conversations) { conversation in
|
||||
let isActive = conversation.id == chatService.activeConversationId
|
||||
Button {
|
||||
chatService.setActiveConversation(id: conversation.id)
|
||||
showConversations = false
|
||||
} label: {
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
Text(conversation.title)
|
||||
.font(.exoSubheadline)
|
||||
.fontWeight(isActive ? .semibold : .regular)
|
||||
.foregroundStyle(
|
||||
isActive ? Color.exoYellow : Color.exoForeground
|
||||
)
|
||||
.lineLimit(1)
|
||||
|
||||
if let modelId = conversation.modelId {
|
||||
Text(modelId)
|
||||
.font(.exoCaption)
|
||||
.foregroundStyle(Color.exoLightGray)
|
||||
.lineLimit(1)
|
||||
}
|
||||
}
|
||||
}
|
||||
.listRowBackground(
|
||||
isActive
|
||||
? Color.exoYellow.opacity(0.1)
|
||||
: Color.exoDarkGray
|
||||
)
|
||||
}
|
||||
.onDelete { indexSet in
|
||||
for index in indexSet {
|
||||
chatService.deleteConversation(id: chatService.conversations[index].id)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
.scrollContentBackground(.hidden)
|
||||
.background(Color.exoBlack)
|
||||
.navigationTitle("Conversations")
|
||||
.navigationBarTitleDisplayMode(.inline)
|
||||
.toolbar {
|
||||
ToolbarItem(placement: .confirmationAction) {
|
||||
Button("Done") { showConversations = false }
|
||||
.font(.exoSubheadline)
|
||||
.foregroundStyle(Color.exoYellow)
|
||||
}
|
||||
ToolbarItem(placement: .topBarLeading) {
|
||||
Button {
|
||||
chatService.createConversation()
|
||||
} label: {
|
||||
Image(systemName: "plus")
|
||||
.foregroundStyle(Color.exoYellow)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,314 @@
|
||||
import SwiftUI
|
||||
|
||||
struct SettingsView: View {
|
||||
@Environment(ClusterService.self) private var clusterService
|
||||
@Environment(DiscoveryService.self) private var discoveryService
|
||||
@Environment(LocalInferenceService.self) private var localInferenceService
|
||||
@Environment(\.dismiss) private var dismiss
|
||||
@State private var host: String = ""
|
||||
@State private var port: String = "52415"
|
||||
|
||||
var body: some View {
|
||||
NavigationStack {
|
||||
Form {
|
||||
localModelSection
|
||||
nearbyClustersSection
|
||||
connectionSection
|
||||
statusSection
|
||||
}
|
||||
.scrollContentBackground(.hidden)
|
||||
.background(Color.exoBlack)
|
||||
.navigationTitle("Settings")
|
||||
.navigationBarTitleDisplayMode(.inline)
|
||||
.toolbar {
|
||||
ToolbarItem(placement: .confirmationAction) {
|
||||
Button("Done") { dismiss() }
|
||||
.font(.exoSubheadline)
|
||||
.foregroundStyle(Color.exoYellow)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Section Headers
|
||||
|
||||
private func sectionHeader(_ title: String) -> some View {
|
||||
Text(title.uppercased())
|
||||
.font(.exoMono(10, weight: .semibold))
|
||||
.tracking(2)
|
||||
.foregroundStyle(Color.exoYellow)
|
||||
}
|
||||
|
||||
// MARK: - Local Model
|
||||
|
||||
private var localModelSection: some View {
|
||||
Section {
|
||||
HStack {
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
Text(localInferenceService.defaultModelId)
|
||||
.font(.exoSubheadline)
|
||||
.foregroundStyle(Color.exoForeground)
|
||||
|
||||
Text(localModelStatusText)
|
||||
.font(.exoCaption)
|
||||
.foregroundStyle(Color.exoLightGray)
|
||||
}
|
||||
|
||||
Spacer()
|
||||
|
||||
localModelActionButton
|
||||
}
|
||||
.listRowBackground(Color.exoDarkGray)
|
||||
|
||||
if case .downloading(let progress) = localInferenceService.modelState {
|
||||
ProgressView(value: progress)
|
||||
.tint(Color.exoYellow)
|
||||
.listRowBackground(Color.exoDarkGray)
|
||||
}
|
||||
} header: {
|
||||
sectionHeader("Local Model")
|
||||
} footer: {
|
||||
Text(
|
||||
"When disconnected from a cluster, messages are processed on-device using this model."
|
||||
)
|
||||
.font(.exoCaption)
|
||||
.foregroundStyle(Color.exoLightGray.opacity(0.7))
|
||||
}
|
||||
}
|
||||
|
||||
private var localModelStatusText: String {
|
||||
switch localInferenceService.modelState {
|
||||
case .notDownloaded: "Not downloaded"
|
||||
case .downloading(let progress): "Downloading \(Int(progress * 100))%..."
|
||||
case .downloaded: "Downloaded — not loaded"
|
||||
case .loading: "Loading into memory..."
|
||||
case .ready: "Ready"
|
||||
case .generating: "Generating..."
|
||||
case .error(let message): "Error: \(message)"
|
||||
}
|
||||
}
|
||||
|
||||
@ViewBuilder
|
||||
private var localModelActionButton: some View {
|
||||
switch localInferenceService.modelState {
|
||||
case .notDownloaded:
|
||||
exoButton("Download") {
|
||||
Task { await localInferenceService.prepareModel() }
|
||||
}
|
||||
case .downloading:
|
||||
ProgressView()
|
||||
.controlSize(.small)
|
||||
.tint(Color.exoYellow)
|
||||
case .downloaded:
|
||||
exoButton("Load") {
|
||||
Task { await localInferenceService.prepareModel() }
|
||||
}
|
||||
case .loading:
|
||||
ProgressView()
|
||||
.controlSize(.small)
|
||||
.tint(Color.exoYellow)
|
||||
case .ready, .generating:
|
||||
exoButton("Unload") {
|
||||
localInferenceService.unloadModel()
|
||||
}
|
||||
case .error:
|
||||
exoButton("Retry", destructive: true) {
|
||||
Task { await localInferenceService.prepareModel() }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func exoButton(_ title: String, destructive: Bool = false, action: @escaping () -> Void)
|
||||
-> some View
|
||||
{
|
||||
let borderColor = destructive ? Color.exoDestructive : Color.exoYellow
|
||||
return Button(action: action) {
|
||||
Text(title.uppercased())
|
||||
.font(.exoMono(11, weight: .semibold))
|
||||
.tracking(1)
|
||||
.foregroundStyle(borderColor)
|
||||
.padding(.horizontal, 10)
|
||||
.padding(.vertical, 5)
|
||||
.overlay(
|
||||
RoundedRectangle(cornerRadius: 6)
|
||||
.stroke(borderColor, lineWidth: 1)
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Nearby Clusters
|
||||
|
||||
private var nearbyClustersSection: some View {
|
||||
Section {
|
||||
if discoveryService.discoveredClusters.isEmpty {
|
||||
if discoveryService.isSearching {
|
||||
HStack {
|
||||
ProgressView()
|
||||
.tint(Color.exoYellow)
|
||||
.padding(.trailing, 8)
|
||||
Text("Searching for clusters...")
|
||||
.font(.exoBody)
|
||||
.foregroundStyle(Color.exoLightGray)
|
||||
}
|
||||
.listRowBackground(Color.exoDarkGray)
|
||||
} else {
|
||||
Text("No clusters found")
|
||||
.font(.exoBody)
|
||||
.foregroundStyle(Color.exoLightGray)
|
||||
.listRowBackground(Color.exoDarkGray)
|
||||
}
|
||||
} else {
|
||||
ForEach(discoveryService.discoveredClusters) { cluster in
|
||||
HStack {
|
||||
VStack(alignment: .leading) {
|
||||
Text(cluster.name)
|
||||
.font(.exoBody)
|
||||
.foregroundStyle(Color.exoForeground)
|
||||
}
|
||||
Spacer()
|
||||
exoButton("Connect") {
|
||||
Task {
|
||||
await clusterService.connectToDiscoveredCluster(
|
||||
cluster, using: discoveryService
|
||||
)
|
||||
if clusterService.isConnected {
|
||||
dismiss()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
.listRowBackground(Color.exoDarkGray)
|
||||
}
|
||||
}
|
||||
} header: {
|
||||
sectionHeader("Nearby Clusters")
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Manual Connection
|
||||
|
||||
private var connectionSection: some View {
|
||||
Section {
|
||||
TextField("IP Address (e.g. 192.168.1.42)", text: $host)
|
||||
.font(.exoBody)
|
||||
.keyboardType(.decimalPad)
|
||||
.textContentType(.URL)
|
||||
.autocorrectionDisabled()
|
||||
.foregroundStyle(Color.exoForeground)
|
||||
.listRowBackground(Color.exoDarkGray)
|
||||
|
||||
TextField("Port", text: $port)
|
||||
.font(.exoBody)
|
||||
.keyboardType(.numberPad)
|
||||
.foregroundStyle(Color.exoForeground)
|
||||
.listRowBackground(Color.exoDarkGray)
|
||||
|
||||
Button {
|
||||
Task {
|
||||
let portNum = Int(port) ?? ConnectionInfo.defaultPort
|
||||
let info = ConnectionInfo(host: host, port: portNum, nodeId: nil)
|
||||
await clusterService.connect(to: info)
|
||||
if clusterService.isConnected {
|
||||
dismiss()
|
||||
}
|
||||
}
|
||||
} label: {
|
||||
Text(clusterService.isConnected ? "RECONNECT" : "CONNECT")
|
||||
.font(.exoMono(13, weight: .semibold))
|
||||
.tracking(1.5)
|
||||
.foregroundStyle(
|
||||
host.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
|
||||
? Color.exoLightGray : Color.exoYellow
|
||||
)
|
||||
.frame(maxWidth: .infinity)
|
||||
}
|
||||
.disabled(host.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty)
|
||||
.listRowBackground(Color.exoDarkGray)
|
||||
} header: {
|
||||
sectionHeader("Manual Connection")
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Status
|
||||
|
||||
private var statusSection: some View {
|
||||
Section {
|
||||
if let connection = clusterService.currentConnection {
|
||||
LabeledContent {
|
||||
Text(connection.host)
|
||||
.font(.exoCaption)
|
||||
.foregroundStyle(Color.exoForeground)
|
||||
} label: {
|
||||
Text("Host")
|
||||
.font(.exoCaption)
|
||||
.foregroundStyle(Color.exoLightGray)
|
||||
}
|
||||
.listRowBackground(Color.exoDarkGray)
|
||||
|
||||
LabeledContent {
|
||||
Text("\(connection.port)")
|
||||
.font(.exoCaption)
|
||||
.foregroundStyle(Color.exoForeground)
|
||||
} label: {
|
||||
Text("Port")
|
||||
.font(.exoCaption)
|
||||
.foregroundStyle(Color.exoLightGray)
|
||||
}
|
||||
.listRowBackground(Color.exoDarkGray)
|
||||
|
||||
if let nodeId = connection.nodeId {
|
||||
LabeledContent {
|
||||
Text(String(nodeId.prefix(12)) + "...")
|
||||
.font(.exoCaption)
|
||||
.foregroundStyle(Color.exoForeground)
|
||||
} label: {
|
||||
Text("Node ID")
|
||||
.font(.exoCaption)
|
||||
.foregroundStyle(Color.exoLightGray)
|
||||
}
|
||||
.listRowBackground(Color.exoDarkGray)
|
||||
}
|
||||
|
||||
LabeledContent {
|
||||
Text("\(clusterService.availableModels.count)")
|
||||
.font(.exoCaption)
|
||||
.foregroundStyle(Color.exoForeground)
|
||||
} label: {
|
||||
Text("Models")
|
||||
.font(.exoCaption)
|
||||
.foregroundStyle(Color.exoLightGray)
|
||||
}
|
||||
.listRowBackground(Color.exoDarkGray)
|
||||
|
||||
Button(role: .destructive) {
|
||||
clusterService.disconnect()
|
||||
} label: {
|
||||
Text("DISCONNECT")
|
||||
.font(.exoMono(13, weight: .semibold))
|
||||
.tracking(1.5)
|
||||
.foregroundStyle(Color.exoDestructive)
|
||||
.frame(maxWidth: .infinity)
|
||||
}
|
||||
.listRowBackground(Color.exoDarkGray)
|
||||
} else {
|
||||
if let error = clusterService.lastError {
|
||||
Label {
|
||||
Text(error)
|
||||
.font(.exoCaption)
|
||||
} icon: {
|
||||
Image(systemName: "exclamationmark.triangle")
|
||||
}
|
||||
.foregroundStyle(Color.exoDestructive)
|
||||
.listRowBackground(Color.exoDarkGray)
|
||||
} else {
|
||||
Text("Not connected")
|
||||
.font(.exoBody)
|
||||
.foregroundStyle(Color.exoLightGray)
|
||||
.listRowBackground(Color.exoDarkGray)
|
||||
}
|
||||
}
|
||||
} header: {
|
||||
sectionHeader("Status")
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
import SwiftUI
|
||||
|
||||
// MARK: - EXO Color Palette
|
||||
|
||||
extension Color {
|
||||
/// Primary background — near-black (#121212)
|
||||
static let exoBlack = Color(red: 0x12 / 255.0, green: 0x12 / 255.0, blue: 0x12 / 255.0)
|
||||
/// Card / surface background (#1F1F1F)
|
||||
static let exoDarkGray = Color(red: 0x1F / 255.0, green: 0x1F / 255.0, blue: 0x1F / 255.0)
|
||||
/// Input field / elevated surface (#353535)
|
||||
static let exoMediumGray = Color(red: 0x35 / 255.0, green: 0x35 / 255.0, blue: 0x35 / 255.0)
|
||||
/// Secondary text (#999999)
|
||||
static let exoLightGray = Color(red: 0x99 / 255.0, green: 0x99 / 255.0, blue: 0x99 / 255.0)
|
||||
/// Accent yellow — matches dashboard (#FFD700)
|
||||
static let exoYellow = Color(red: 0xFF / 255.0, green: 0xD7 / 255.0, blue: 0x00 / 255.0)
|
||||
/// Primary foreground text (#E5E5E5)
|
||||
static let exoForeground = Color(red: 0xE5 / 255.0, green: 0xE5 / 255.0, blue: 0xE5 / 255.0)
|
||||
/// Destructive / error (#E74C3C)
|
||||
static let exoDestructive = Color(red: 0xE7 / 255.0, green: 0x4C / 255.0, blue: 0x3C / 255.0)
|
||||
}
|
||||
|
||||
// MARK: - EXO Typography (SF Mono via .monospaced design)
|
||||
|
||||
extension Font {
|
||||
/// Monospaced font at a given size and weight.
|
||||
static func exoMono(_ size: CGFloat, weight: Font.Weight = .regular) -> Font {
|
||||
.system(size: size, weight: weight, design: .monospaced)
|
||||
}
|
||||
|
||||
/// Body text — 15pt monospaced
|
||||
static let exoBody: Font = .system(size: 15, weight: .regular, design: .monospaced)
|
||||
/// Caption — 11pt monospaced
|
||||
static let exoCaption: Font = .system(size: 11, weight: .regular, design: .monospaced)
|
||||
/// Subheadline — 13pt monospaced medium
|
||||
static let exoSubheadline: Font = .system(size: 13, weight: .medium, design: .monospaced)
|
||||
/// Headline — 17pt monospaced semibold
|
||||
static let exoHeadline: Font = .system(size: 17, weight: .semibold, design: .monospaced)
|
||||
}
|
||||
|
||||
// MARK: - Reusable Gradient Divider
|
||||
|
||||
struct GradientDivider: View {
|
||||
var body: some View {
|
||||
LinearGradient(
|
||||
colors: [.clear, Color.exoYellow.opacity(0.3), .clear],
|
||||
startPoint: .leading,
|
||||
endPoint: .trailing
|
||||
)
|
||||
.frame(height: 1)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
//
|
||||
// EXO_iOSTests.swift
|
||||
// EXO-iOSTests
|
||||
//
|
||||
// Created by Sami Khan on 2026-02-17.
|
||||
//
|
||||
|
||||
import Testing
|
||||
|
||||
@testable import EXO_iOS
|
||||
|
||||
struct EXO_iOSTests {
|
||||
|
||||
@Test func example() async throws {
|
||||
// Write your test here and use APIs like `#expect(...)` to check expected conditions.
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
//
|
||||
// EXO_iOSUITests.swift
|
||||
// EXO-iOSUITests
|
||||
//
|
||||
// Created by Sami Khan on 2026-02-17.
|
||||
//
|
||||
|
||||
import XCTest
|
||||
|
||||
final class EXO_iOSUITests: XCTestCase {
|
||||
|
||||
override func setUpWithError() throws {
|
||||
// Put setup code here. This method is called before the invocation of each test method in the class.
|
||||
|
||||
// In UI tests it is usually best to stop immediately when a failure occurs.
|
||||
continueAfterFailure = false
|
||||
|
||||
// In UI tests it’s important to set the initial state - such as interface orientation - required for your tests before they run. The setUp method is a good place to do this.
|
||||
}
|
||||
|
||||
override func tearDownWithError() throws {
|
||||
// Put teardown code here. This method is called after the invocation of each test method in the class.
|
||||
}
|
||||
|
||||
@MainActor
|
||||
func testExample() throws {
|
||||
// UI tests must launch the application that they test.
|
||||
let app = XCUIApplication()
|
||||
app.launch()
|
||||
|
||||
// Use XCTAssert and related functions to verify your tests produce the correct results.
|
||||
}
|
||||
|
||||
@MainActor
|
||||
func testLaunchPerformance() throws {
|
||||
// This measures how long it takes to launch your application.
|
||||
measure(metrics: [XCTApplicationLaunchMetric()]) {
|
||||
XCUIApplication().launch()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
//
|
||||
// EXO_iOSUITestsLaunchTests.swift
|
||||
// EXO-iOSUITests
|
||||
//
|
||||
// Created by Sami Khan on 2026-02-17.
|
||||
//
|
||||
|
||||
import XCTest
|
||||
|
||||
final class EXO_iOSUITestsLaunchTests: XCTestCase {
|
||||
|
||||
override class var runsForEachTargetApplicationUIConfiguration: Bool {
|
||||
true
|
||||
}
|
||||
|
||||
override func setUpWithError() throws {
|
||||
continueAfterFailure = false
|
||||
}
|
||||
|
||||
@MainActor
|
||||
func testLaunch() throws {
|
||||
let app = XCUIApplication()
|
||||
app.launch()
|
||||
|
||||
// Insert steps here to perform after app launch but before taking a screenshot,
|
||||
// such as logging into a test account or navigating somewhere in the app
|
||||
|
||||
let attachment = XCTAttachment(screenshot: app.screenshot())
|
||||
attachment.name = "Launch Screen"
|
||||
attachment.lifetime = .keepAlways
|
||||
add(attachment)
|
||||
}
|
||||
}
|
||||
+208
-65
@@ -15,14 +15,14 @@ struct ContentView: View {
|
||||
@EnvironmentObject private var localNetworkChecker: LocalNetworkChecker
|
||||
@EnvironmentObject private var updater: SparkleUpdater
|
||||
@EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService
|
||||
@EnvironmentObject private var settingsWindowController: SettingsWindowController
|
||||
@EnvironmentObject private var bugReportWindowController: BugReportWindowController
|
||||
@State private var focusedNode: NodeViewModel?
|
||||
@State private var deletingInstanceIDs: Set<String> = []
|
||||
@State private var showAllNodes = false
|
||||
@State private var showAllInstances = false
|
||||
@State private var baseURLCopied = false
|
||||
@State private var showAdvanced = false
|
||||
@State private var showDebugInfo = false
|
||||
@State private var bugReportInFlight = false
|
||||
@State private var bugReportMessage: String?
|
||||
@State private var uninstallInProgress = false
|
||||
@State private var pendingNamespace: String = ""
|
||||
@State private var pendingHFToken: String = ""
|
||||
@@ -258,86 +258,139 @@ struct ContentView: View {
|
||||
VStack(alignment: .leading, spacing: 0) {
|
||||
if controller.status != .stopped {
|
||||
dashboardButton
|
||||
baseURLRow
|
||||
Divider()
|
||||
.padding(.vertical, 8)
|
||||
} else {
|
||||
Divider()
|
||||
.padding(.vertical, 4)
|
||||
}
|
||||
HoverButton(
|
||||
title: "Settings",
|
||||
tint: .primary,
|
||||
trailingSystemImage: "gear"
|
||||
) {
|
||||
settingsWindowController.open(
|
||||
controller: controller,
|
||||
updater: updater,
|
||||
networkStatusService: networkStatusService,
|
||||
thunderboltBridgeService: thunderboltBridgeService,
|
||||
stateService: stateService
|
||||
)
|
||||
}
|
||||
HoverButton(
|
||||
title: "Check for Updates",
|
||||
tint: .primary,
|
||||
trailingSystemImage: "arrow.triangle.2.circlepath"
|
||||
) {
|
||||
updater.checkForUpdates()
|
||||
}
|
||||
HoverButton(
|
||||
title: "Share Bug Report…",
|
||||
tint: .primary,
|
||||
trailingSystemImage: "ladybug"
|
||||
) {
|
||||
bugReportWindowController.open()
|
||||
}
|
||||
.padding(.bottom, 8)
|
||||
HoverButton(title: "Quit", tint: .secondary) {
|
||||
advancedSection
|
||||
.padding(.bottom, 8)
|
||||
controlButton(title: "Quit", tint: .secondary) {
|
||||
controller.stop()
|
||||
NSApplication.shared.terminate(nil)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private var dashboardButton: some View {
|
||||
HoverButton(
|
||||
title: "Web Dashboard",
|
||||
tint: .primary,
|
||||
trailingSystemImage: "arrow.up.right"
|
||||
) {
|
||||
guard let url = URL(string: "http://localhost:52415/") else { return }
|
||||
NSWorkspace.shared.open(url)
|
||||
private var advancedSection: some View {
|
||||
VStack(alignment: .leading, spacing: 6) {
|
||||
HStack {
|
||||
Text("Advanced")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
Spacer()
|
||||
collapseButton(isExpanded: $showAdvanced)
|
||||
}
|
||||
.animation(nil, value: showAdvanced)
|
||||
if showAdvanced {
|
||||
VStack(alignment: .leading, spacing: 8) {
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
Text("Cluster Namespace")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
HStack {
|
||||
TextField("optional", text: $pendingNamespace)
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.font(.caption2)
|
||||
.onAppear {
|
||||
pendingNamespace = controller.customNamespace
|
||||
}
|
||||
Button("Save & Restart") {
|
||||
controller.customNamespace = pendingNamespace
|
||||
if controller.status == .running || controller.status == .starting {
|
||||
controller.restart()
|
||||
}
|
||||
}
|
||||
.font(.caption2)
|
||||
.disabled(pendingNamespace == controller.customNamespace)
|
||||
}
|
||||
}
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
Text("HuggingFace Token")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
HStack {
|
||||
SecureField("optional", text: $pendingHFToken)
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.font(.caption2)
|
||||
.onAppear {
|
||||
pendingHFToken = controller.hfToken
|
||||
}
|
||||
Button("Save & Restart") {
|
||||
controller.hfToken = pendingHFToken
|
||||
if controller.status == .running || controller.status == .starting {
|
||||
controller.restart()
|
||||
}
|
||||
}
|
||||
.font(.caption2)
|
||||
.disabled(pendingHFToken == controller.hfToken)
|
||||
}
|
||||
}
|
||||
Divider()
|
||||
HStack {
|
||||
Toggle(
|
||||
"Enable Image Models (experimental)", isOn: $pendingEnableImageModels
|
||||
)
|
||||
.toggleStyle(.switch)
|
||||
.font(.caption2)
|
||||
.onAppear {
|
||||
pendingEnableImageModels = controller.enableImageModels
|
||||
}
|
||||
|
||||
Spacer()
|
||||
|
||||
Button("Save & Restart") {
|
||||
controller.enableImageModels = pendingEnableImageModels
|
||||
if controller.status == .running || controller.status == .starting {
|
||||
controller.restart()
|
||||
}
|
||||
}
|
||||
.font(.caption2)
|
||||
.disabled(pendingEnableImageModels == controller.enableImageModels)
|
||||
}
|
||||
HoverButton(title: "Check for Updates", small: true) {
|
||||
updater.checkForUpdates()
|
||||
}
|
||||
debugSection
|
||||
HoverButton(title: "Uninstall", tint: .red, small: true) {
|
||||
showUninstallConfirmationAlert()
|
||||
}
|
||||
.disabled(uninstallInProgress)
|
||||
}
|
||||
.transition(.opacity)
|
||||
}
|
||||
}
|
||||
.animation(.easeInOut(duration: 0.25), value: showAdvanced)
|
||||
}
|
||||
|
||||
private var baseURLRow: some View {
|
||||
HStack(spacing: 6) {
|
||||
Image(systemName: "link")
|
||||
.imageScale(.small)
|
||||
.foregroundColor(.secondary)
|
||||
Text("localhost:52415/v1")
|
||||
.font(.system(.caption, design: .monospaced))
|
||||
.foregroundColor(.primary)
|
||||
Spacer()
|
||||
Button {
|
||||
NSPasteboard.general.clearContents()
|
||||
NSPasteboard.general.setString("http://localhost:52415/v1", forType: .string)
|
||||
baseURLCopied = true
|
||||
DispatchQueue.main.asyncAfter(deadline: .now() + 2) {
|
||||
baseURLCopied = false
|
||||
}
|
||||
} label: {
|
||||
Image(systemName: baseURLCopied ? "checkmark" : "doc.on.doc")
|
||||
private func controlButton(title: String, tint: Color = .primary, action: @escaping () -> Void)
|
||||
-> some View
|
||||
{
|
||||
HoverButton(title: title, tint: tint, trailingSystemImage: nil, action: action)
|
||||
}
|
||||
|
||||
private var dashboardButton: some View {
|
||||
Button {
|
||||
guard let url = URL(string: "http://localhost:52415/") else { return }
|
||||
NSWorkspace.shared.open(url)
|
||||
} label: {
|
||||
HStack {
|
||||
Image(systemName: "arrow.up.right.square")
|
||||
.imageScale(.small)
|
||||
.foregroundColor(baseURLCopied ? .green : .secondary)
|
||||
.contentTransition(.symbolEffect(.replace))
|
||||
Text("Dashboard")
|
||||
.fontWeight(.medium)
|
||||
Spacer()
|
||||
}
|
||||
.buttonStyle(.plain)
|
||||
.help("Copy API base URL")
|
||||
.padding(.vertical, 8)
|
||||
.padding(.horizontal, 10)
|
||||
.background(
|
||||
RoundedRectangle(cornerRadius: 8, style: .continuous)
|
||||
.fill(Color(red: 1.0, green: 0.87, blue: 0.0).opacity(0.2))
|
||||
)
|
||||
}
|
||||
.padding(.vertical, 4)
|
||||
.padding(.horizontal, 8)
|
||||
.buttonStyle(.plain)
|
||||
.padding(.bottom, 4)
|
||||
}
|
||||
|
||||
private func collapseButton(isExpanded: Binding<Bool>) -> some View {
|
||||
@@ -475,6 +528,40 @@ struct ContentView: View {
|
||||
}
|
||||
}
|
||||
|
||||
private var debugSection: some View {
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
HoverButton(
|
||||
title: "Debug Info",
|
||||
tint: .primary,
|
||||
trailingSystemImage: showDebugInfo ? "chevron.up" : "chevron.down",
|
||||
small: true
|
||||
) {
|
||||
showDebugInfo.toggle()
|
||||
}
|
||||
if showDebugInfo {
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
Text("Version: \(buildTag)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
Text("Commit: \(buildCommit)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
Text(thunderboltStatusText)
|
||||
.font(.caption2)
|
||||
.foregroundColor(thunderboltStatusColor)
|
||||
clusterThunderboltBridgeView
|
||||
interfaceIpList
|
||||
rdmaStatusView
|
||||
sendBugReportButton
|
||||
.padding(.top, 6)
|
||||
}
|
||||
.padding(.leading, 8)
|
||||
.transition(.opacity)
|
||||
}
|
||||
}
|
||||
.animation(.easeInOut(duration: 0.25), value: showDebugInfo)
|
||||
}
|
||||
|
||||
private var rdmaStatusView: some View {
|
||||
let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:]
|
||||
let localNodeId = stateService.localNodeId
|
||||
@@ -523,6 +610,42 @@ struct ContentView: View {
|
||||
}
|
||||
}
|
||||
|
||||
private var sendBugReportButton: some View {
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
Button {
|
||||
Task {
|
||||
await sendBugReport()
|
||||
}
|
||||
} label: {
|
||||
HStack {
|
||||
if bugReportInFlight {
|
||||
ProgressView()
|
||||
.scaleEffect(0.6)
|
||||
}
|
||||
Text("Send Bug Report")
|
||||
.font(.caption)
|
||||
.fontWeight(.semibold)
|
||||
Spacer()
|
||||
}
|
||||
.padding(.vertical, 6)
|
||||
.padding(.horizontal, 8)
|
||||
.background(
|
||||
RoundedRectangle(cornerRadius: 6)
|
||||
.fill(Color.accentColor.opacity(0.12))
|
||||
)
|
||||
}
|
||||
.buttonStyle(.plain)
|
||||
.disabled(bugReportInFlight)
|
||||
|
||||
if let message = bugReportMessage {
|
||||
Text(message)
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
.fixedSize(horizontal: false, vertical: true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private var processToggleBinding: Binding<Bool> {
|
||||
Binding(
|
||||
get: {
|
||||
@@ -563,6 +686,19 @@ struct ContentView: View {
|
||||
)
|
||||
}
|
||||
|
||||
private func sendBugReport() async {
|
||||
bugReportInFlight = true
|
||||
bugReportMessage = "Collecting logs..."
|
||||
let service = BugReportService()
|
||||
do {
|
||||
let outcome = try await service.sendReport(isManual: true)
|
||||
bugReportMessage = outcome.message
|
||||
} catch {
|
||||
bugReportMessage = error.localizedDescription
|
||||
}
|
||||
bugReportInFlight = false
|
||||
}
|
||||
|
||||
private func showUninstallConfirmationAlert() {
|
||||
let alert = NSAlert()
|
||||
alert.messageText = "Uninstall EXO"
|
||||
@@ -645,6 +781,13 @@ struct ContentView: View {
|
||||
}
|
||||
}
|
||||
|
||||
private var buildTag: String {
|
||||
Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
|
||||
}
|
||||
|
||||
private var buildCommit: String {
|
||||
Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown"
|
||||
}
|
||||
}
|
||||
|
||||
private struct HoverButton: View {
|
||||
|
||||
@@ -21,10 +21,7 @@ struct EXOApp: App {
|
||||
@StateObject private var localNetworkChecker: LocalNetworkChecker
|
||||
@StateObject private var updater: SparkleUpdater
|
||||
@StateObject private var thunderboltBridgeService: ThunderboltBridgeService
|
||||
@StateObject private var settingsWindowController: SettingsWindowController
|
||||
@StateObject private var bugReportWindowController: BugReportWindowController
|
||||
private let terminationObserver: TerminationObserver
|
||||
private let firstLaunchPopout = FirstLaunchPopout()
|
||||
private let ciContext = CIContext(options: nil)
|
||||
|
||||
init() {
|
||||
@@ -46,14 +43,12 @@ struct EXOApp: App {
|
||||
_updater = StateObject(wrappedValue: updater)
|
||||
let thunderboltBridge = ThunderboltBridgeService(clusterStateService: service)
|
||||
_thunderboltBridgeService = StateObject(wrappedValue: thunderboltBridge)
|
||||
_settingsWindowController = StateObject(wrappedValue: SettingsWindowController())
|
||||
_bugReportWindowController = StateObject(wrappedValue: BugReportWindowController())
|
||||
enableLaunchAtLoginIfNeeded()
|
||||
// Install LaunchDaemon to disable Thunderbolt Bridge on startup (prevents network loops)
|
||||
NetworkSetupHelper.promptAndInstallIfNeeded()
|
||||
// Check local network access periodically (warning disappears when user grants permission)
|
||||
localNetwork.startPeriodicChecking(interval: 10)
|
||||
controller.scheduleLaunch(after: 5)
|
||||
controller.scheduleLaunch(after: 15)
|
||||
service.startPolling()
|
||||
networkStatus.startPolling()
|
||||
}
|
||||
@@ -67,20 +62,8 @@ struct EXOApp: App {
|
||||
.environmentObject(localNetworkChecker)
|
||||
.environmentObject(updater)
|
||||
.environmentObject(thunderboltBridgeService)
|
||||
.environmentObject(settingsWindowController)
|
||||
.environmentObject(bugReportWindowController)
|
||||
} label: {
|
||||
menuBarIcon
|
||||
.onReceive(controller.$isFirstLaunchReady) { ready in
|
||||
if ready {
|
||||
DispatchQueue.main.asyncAfter(deadline: .now() + 3.0) {
|
||||
self.firstLaunchPopout.onComplete = { [weak controller] in
|
||||
controller?.markOnboardingCompleted()
|
||||
}
|
||||
self.firstLaunchPopout.show()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
.menuBarExtraStyle(.window)
|
||||
}
|
||||
|
||||
@@ -4,30 +4,7 @@ import Foundation
|
||||
|
||||
private let customNamespaceKey = "EXOCustomNamespace"
|
||||
private let hfTokenKey = "EXOHFToken"
|
||||
private let hfEndpointKey = "EXOHFEndpoint"
|
||||
private let enableImageModelsKey = "EXOEnableImageModels"
|
||||
private let offlineModeKey = "EXOOfflineMode"
|
||||
private let fastSynchEnabledKey = "EXOFastSynchEnabled"
|
||||
private let onboardingCompletedKey = "EXOOnboardingCompleted"
|
||||
private let defaultModelsDirKey = "EXODefaultModelsDir"
|
||||
private let additionalModelsDirsKey = "EXOAdditionalModelsDirs"
|
||||
private let readOnlyModelsDirsKey = "EXOReadOnlyModelsDirs"
|
||||
private let customEnvironmentVariablesKey = "EXOCustomEnvironmentVariables"
|
||||
|
||||
/// A user-defined environment variable that is injected into the exo child
|
||||
/// process at launch. Used as an escape hatch for env vars that don't have
|
||||
/// first-class typed UI in Settings.
|
||||
struct CustomEnvironmentVariable: Codable, Identifiable, Equatable {
|
||||
var id: UUID
|
||||
var key: String
|
||||
var value: String
|
||||
|
||||
init(id: UUID = UUID(), key: String = "", value: String = "") {
|
||||
self.id = id
|
||||
self.key = key
|
||||
self.value = value
|
||||
}
|
||||
}
|
||||
|
||||
@MainActor
|
||||
final class ExoProcessController: ObservableObject {
|
||||
@@ -74,14 +51,6 @@ final class ExoProcessController: ObservableObject {
|
||||
UserDefaults.standard.set(hfToken, forKey: hfTokenKey)
|
||||
}
|
||||
}
|
||||
@Published var hfEndpoint: String = {
|
||||
return UserDefaults.standard.string(forKey: hfEndpointKey) ?? ""
|
||||
}()
|
||||
{
|
||||
didSet {
|
||||
UserDefaults.standard.set(hfEndpoint, forKey: hfEndpointKey)
|
||||
}
|
||||
}
|
||||
@Published var enableImageModels: Bool = {
|
||||
return UserDefaults.standard.bool(forKey: enableImageModelsKey)
|
||||
}()
|
||||
@@ -90,71 +59,6 @@ final class ExoProcessController: ObservableObject {
|
||||
UserDefaults.standard.set(enableImageModels, forKey: enableImageModelsKey)
|
||||
}
|
||||
}
|
||||
@Published var offlineMode: Bool = {
|
||||
return UserDefaults.standard.bool(forKey: offlineModeKey)
|
||||
}()
|
||||
{
|
||||
didSet {
|
||||
UserDefaults.standard.set(offlineMode, forKey: offlineModeKey)
|
||||
}
|
||||
}
|
||||
@Published var fastSynchEnabled: Bool = {
|
||||
if UserDefaults.standard.object(forKey: fastSynchEnabledKey) == nil {
|
||||
return true
|
||||
}
|
||||
return UserDefaults.standard.bool(forKey: fastSynchEnabledKey)
|
||||
}()
|
||||
{
|
||||
didSet {
|
||||
UserDefaults.standard.set(fastSynchEnabled, forKey: fastSynchEnabledKey)
|
||||
}
|
||||
}
|
||||
@Published var defaultModelsDir: String = {
|
||||
return UserDefaults.standard.string(forKey: defaultModelsDirKey) ?? ""
|
||||
}()
|
||||
{
|
||||
didSet {
|
||||
UserDefaults.standard.set(defaultModelsDir, forKey: defaultModelsDirKey)
|
||||
}
|
||||
}
|
||||
@Published var additionalModelsDirs: String = {
|
||||
return UserDefaults.standard.string(forKey: additionalModelsDirsKey) ?? ""
|
||||
}()
|
||||
{
|
||||
didSet {
|
||||
UserDefaults.standard.set(additionalModelsDirs, forKey: additionalModelsDirsKey)
|
||||
}
|
||||
}
|
||||
@Published var readOnlyModelsDirs: String = {
|
||||
return UserDefaults.standard.string(forKey: readOnlyModelsDirsKey) ?? ""
|
||||
}()
|
||||
{
|
||||
didSet {
|
||||
UserDefaults.standard.set(readOnlyModelsDirs, forKey: readOnlyModelsDirsKey)
|
||||
}
|
||||
}
|
||||
@Published var customEnvironmentVariables: [CustomEnvironmentVariable] = {
|
||||
guard
|
||||
let data = UserDefaults.standard.data(forKey: customEnvironmentVariablesKey),
|
||||
let decoded = try? JSONDecoder().decode(
|
||||
[CustomEnvironmentVariable].self, from: data
|
||||
)
|
||||
else {
|
||||
return []
|
||||
}
|
||||
return decoded
|
||||
}()
|
||||
{
|
||||
didSet {
|
||||
guard let data = try? JSONEncoder().encode(customEnvironmentVariables) else {
|
||||
return
|
||||
}
|
||||
UserDefaults.standard.set(data, forKey: customEnvironmentVariablesKey)
|
||||
}
|
||||
}
|
||||
|
||||
/// Fires once when EXO transitions to `.running` for the very first time (fresh install).
|
||||
@Published private(set) var isFirstLaunchReady = false
|
||||
|
||||
private var process: Process?
|
||||
private var runtimeDirectoryURL: URL?
|
||||
@@ -209,9 +113,6 @@ final class ExoProcessController: ObservableObject {
|
||||
try child.run()
|
||||
process = child
|
||||
status = .running
|
||||
|
||||
// Show welcome popout on every launch
|
||||
isFirstLaunchReady = true
|
||||
} catch {
|
||||
process = nil
|
||||
status = .failed(message: "Launch error")
|
||||
@@ -225,37 +126,11 @@ final class ExoProcessController: ObservableObject {
|
||||
return
|
||||
}
|
||||
process.terminationHandler = nil
|
||||
status = .stopped
|
||||
|
||||
guard process.isRunning else {
|
||||
self.process = nil
|
||||
return
|
||||
if process.isRunning {
|
||||
process.terminate()
|
||||
}
|
||||
|
||||
let proc = process
|
||||
self.process = nil
|
||||
|
||||
Task.detached {
|
||||
proc.interrupt()
|
||||
|
||||
for _ in 0..<50 {
|
||||
if !proc.isRunning { return }
|
||||
try? await Task.sleep(nanoseconds: 100_000_000)
|
||||
}
|
||||
|
||||
if proc.isRunning {
|
||||
proc.terminate()
|
||||
}
|
||||
|
||||
for _ in 0..<30 {
|
||||
if !proc.isRunning { return }
|
||||
try? await Task.sleep(nanoseconds: 100_000_000)
|
||||
}
|
||||
|
||||
if proc.isRunning {
|
||||
kill(proc.processIdentifier, SIGKILL)
|
||||
}
|
||||
}
|
||||
status = .stopped
|
||||
}
|
||||
|
||||
func restart() {
|
||||
@@ -263,17 +138,6 @@ final class ExoProcessController: ObservableObject {
|
||||
launch()
|
||||
}
|
||||
|
||||
/// Mark onboarding as completed (user interacted with the welcome popout).
|
||||
func markOnboardingCompleted() {
|
||||
UserDefaults.standard.set(true, forKey: onboardingCompletedKey)
|
||||
}
|
||||
|
||||
/// Reset onboarding so the welcome popout appears on next launch.
|
||||
func resetOnboarding() {
|
||||
UserDefaults.standard.removeObject(forKey: onboardingCompletedKey)
|
||||
isFirstLaunchReady = false
|
||||
}
|
||||
|
||||
func scheduleLaunch(after seconds: TimeInterval) {
|
||||
cancelPendingLaunch()
|
||||
let start = max(1, Int(ceil(seconds)))
|
||||
@@ -356,16 +220,9 @@ final class ExoProcessController: ObservableObject {
|
||||
if !hfToken.isEmpty {
|
||||
environment["HF_TOKEN"] = hfToken
|
||||
}
|
||||
if !hfEndpoint.isEmpty {
|
||||
environment["HF_ENDPOINT"] = hfEndpoint
|
||||
}
|
||||
if enableImageModels {
|
||||
environment["EXO_ENABLE_IMAGE_MODELS"] = "true"
|
||||
}
|
||||
if offlineMode {
|
||||
environment["EXO_OFFLINE"] = "true"
|
||||
}
|
||||
environment["EXO_FAST_SYNCH"] = fastSynchEnabled ? "true" : "false"
|
||||
|
||||
var paths: [String] = []
|
||||
if let existing = environment["PATH"], !existing.isEmpty {
|
||||
@@ -390,29 +247,6 @@ final class ExoProcessController: ObservableObject {
|
||||
}
|
||||
|
||||
environment["PATH"] = paths.joined(separator: ":")
|
||||
|
||||
let trimmedDefaultModelsDir = defaultModelsDir.trimmingCharacters(in: .whitespaces)
|
||||
if !trimmedDefaultModelsDir.isEmpty {
|
||||
environment["EXO_DEFAULT_MODELS_DIR"] = trimmedDefaultModelsDir
|
||||
}
|
||||
let trimmedAdditionalModelsDirs = additionalModelsDirs.trimmingCharacters(in: .whitespaces)
|
||||
if !trimmedAdditionalModelsDirs.isEmpty {
|
||||
environment["EXO_MODELS_DIRS"] = trimmedAdditionalModelsDirs
|
||||
}
|
||||
let trimmedReadOnlyModelsDirs = readOnlyModelsDirs.trimmingCharacters(in: .whitespaces)
|
||||
if !trimmedReadOnlyModelsDirs.isEmpty {
|
||||
environment["EXO_MODELS_READ_ONLY_DIRS"] = trimmedReadOnlyModelsDirs
|
||||
}
|
||||
|
||||
// Apply user-defined arbitrary environment variables last so that
|
||||
// power users can override any of the typed fields above when
|
||||
// necessary. Empty keys are ignored.
|
||||
for variable in customEnvironmentVariables {
|
||||
let trimmedKey = variable.key.trimmingCharacters(in: .whitespaces)
|
||||
guard !trimmedKey.isEmpty else { continue }
|
||||
environment[trimmedKey] = variable.value
|
||||
}
|
||||
|
||||
return environment
|
||||
}
|
||||
|
||||
|
||||
@@ -9,7 +9,7 @@
|
||||
<key>EXOBuildCommit</key>
|
||||
<string>$(EXO_BUILD_COMMIT)</string>
|
||||
<key>EXOBugReportPresignedUrlEndpoint</key>
|
||||
<string>https://reports.exolabs.net/presigned-urls</string>
|
||||
<string>$(EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT)</string>
|
||||
<key>NSLocalNetworkUsageDescription</key>
|
||||
<string>EXO needs local network access to discover and connect to other devices in your cluster for distributed AI inference.</string>
|
||||
<key>NSBonjourServices</key>
|
||||
|
||||
@@ -38,8 +38,7 @@ struct BugReportService {
|
||||
func sendReport(
|
||||
baseURL: URL = URL(string: "http://127.0.0.1:52415")!,
|
||||
now: Date = Date(),
|
||||
isManual: Bool = false,
|
||||
userDescription: String? = nil
|
||||
isManual: Bool = false
|
||||
) async throws -> BugReportOutcome {
|
||||
let timestamp = Self.runTimestampString(now)
|
||||
let dayPrefix = Self.dayPrefixString(now)
|
||||
@@ -61,8 +60,7 @@ struct BugReportService {
|
||||
ifconfig: ifconfigText,
|
||||
debugInfo: debugInfo,
|
||||
isManual: isManual,
|
||||
clusterTbBridgeStatus: clusterTbBridgeStatus,
|
||||
userDescription: userDescription
|
||||
clusterTbBridgeStatus: clusterTbBridgeStatus
|
||||
)
|
||||
|
||||
let eventLogFiles = readAllEventLogs()
|
||||
@@ -308,8 +306,7 @@ struct BugReportService {
|
||||
ifconfig: String,
|
||||
debugInfo: DebugInfo,
|
||||
isManual: Bool,
|
||||
clusterTbBridgeStatus: [[String: Any]]?,
|
||||
userDescription: String? = nil
|
||||
clusterTbBridgeStatus: [[String: Any]]?
|
||||
) -> Data? {
|
||||
let system = readSystemMetadata()
|
||||
let exo = readExoMetadata()
|
||||
@@ -326,9 +323,6 @@ struct BugReportService {
|
||||
if let tbStatus = clusterTbBridgeStatus {
|
||||
payload["cluster_thunderbolt_bridge"] = tbStatus
|
||||
}
|
||||
if let desc = userDescription, !desc.isEmpty {
|
||||
payload["user_description"] = desc
|
||||
}
|
||||
return try? JSONSerialization.data(withJSONObject: payload, options: [.prettyPrinted])
|
||||
}
|
||||
|
||||
|
||||
@@ -17,7 +17,7 @@ final class ClusterStateService: ObservableObject {
|
||||
|
||||
init(
|
||||
baseURL: URL = URL(string: "http://127.0.0.1:52415")!,
|
||||
session: URLSession = ClusterStateService.makeNonCachingSession()
|
||||
session: URLSession = .shared
|
||||
) {
|
||||
self.baseURL = baseURL
|
||||
self.endpoint = baseURL.appendingPathComponent("state")
|
||||
@@ -27,23 +27,6 @@ final class ClusterStateService: ObservableObject {
|
||||
self.decoder = decoder
|
||||
}
|
||||
|
||||
/// `URLSession.shared` carries an on-disk `URLCache` that persists every
|
||||
/// response body under `~/Library/Caches/exolabs.EXO/`. We poll `/state`
|
||||
/// at 2 Hz from `startPolling`, so leaving the shared cache attached
|
||||
/// dirties ~500–620 KB/sec of file-backed memory and trips macOS's
|
||||
/// per-process `disk writes` resource limit (microstackshot reports
|
||||
/// observed on M3 Ultra producing GBs of cached responses per hour).
|
||||
/// Cluster-state polling responses are time-sensitive and small; they
|
||||
/// gain nothing from being cached on disk. Use an ephemeral session
|
||||
/// with `urlCache = nil` so neither response bodies nor metadata
|
||||
/// touch disk.
|
||||
private static func makeNonCachingSession() -> URLSession {
|
||||
let config = URLSessionConfiguration.ephemeral
|
||||
config.urlCache = nil
|
||||
config.requestCachePolicy = .reloadIgnoringLocalCacheData
|
||||
return URLSession(configuration: config)
|
||||
}
|
||||
|
||||
func startPolling(interval: TimeInterval = 0.5) {
|
||||
stopPolling()
|
||||
Task {
|
||||
|
||||
@@ -1,242 +0,0 @@
|
||||
import AppKit
|
||||
import SwiftUI
|
||||
|
||||
/// Manages a standalone window for the bug-report flow.
|
||||
/// Ensures only one instance exists and brings it to front on repeated opens.
|
||||
@MainActor
|
||||
final class BugReportWindowController: ObservableObject {
|
||||
private var window: NSWindow?
|
||||
|
||||
func open() {
|
||||
if let existing = window, existing.isVisible {
|
||||
existing.makeKeyAndOrderFront(nil)
|
||||
NSApp.activate()
|
||||
return
|
||||
}
|
||||
|
||||
let view = BugReportView(onDismiss: { [weak self] in
|
||||
self?.window?.close()
|
||||
})
|
||||
|
||||
let hostingController = NSHostingController(rootView: view)
|
||||
hostingController.sizingOptions = [.preferredContentSize, .minSize]
|
||||
|
||||
let newWindow = NSWindow(contentViewController: hostingController)
|
||||
newWindow.styleMask = [.titled, .closable, .resizable]
|
||||
newWindow.title = "Send a Bug Report"
|
||||
newWindow.center()
|
||||
newWindow.setFrameAutosaveName("ExoBugReportWindow")
|
||||
newWindow.isReleasedWhenClosed = false
|
||||
newWindow.makeKeyAndOrderFront(nil)
|
||||
NSApp.activate()
|
||||
|
||||
window = newWindow
|
||||
}
|
||||
}
|
||||
|
||||
private struct BugReportView: View {
|
||||
fileprivate enum Phase: Equatable {
|
||||
case prompting
|
||||
case sending(String)
|
||||
case success(String)
|
||||
case failure(String)
|
||||
}
|
||||
|
||||
let onDismiss: () -> Void
|
||||
|
||||
@State private var phase: Phase = .prompting
|
||||
@State private var userDescription: String = ""
|
||||
@FocusState private var descriptionFocused: Bool
|
||||
|
||||
var body: some View {
|
||||
VStack(alignment: .leading, spacing: 12) {
|
||||
switch phase {
|
||||
case .prompting:
|
||||
promptingView
|
||||
case .sending(let message):
|
||||
sendingView(message: message)
|
||||
case .success(let message):
|
||||
successView(message: message)
|
||||
case .failure(let message):
|
||||
failureView(message: message)
|
||||
}
|
||||
}
|
||||
.padding(16)
|
||||
.frame(minWidth: 380)
|
||||
.animation(.easeInOut(duration: 0.2), value: phase)
|
||||
.onAppear { descriptionFocused = true }
|
||||
}
|
||||
|
||||
private var promptingView: some View {
|
||||
VStack(alignment: .leading, spacing: 8) {
|
||||
Text("Description (optional)")
|
||||
.font(.subheadline)
|
||||
.foregroundColor(.secondary)
|
||||
ZStack(alignment: .topLeading) {
|
||||
if userDescription.isEmpty {
|
||||
Text("What were you doing when it broke?")
|
||||
.font(.body)
|
||||
.foregroundColor(Color(nsColor: .placeholderTextColor))
|
||||
.padding(.horizontal, 10)
|
||||
.padding(.vertical, 8)
|
||||
.allowsHitTesting(false)
|
||||
}
|
||||
TextEditor(text: $userDescription)
|
||||
.font(.body)
|
||||
.scrollContentBackground(.hidden)
|
||||
.padding(4)
|
||||
.frame(height: 72)
|
||||
.focused($descriptionFocused)
|
||||
}
|
||||
.background(
|
||||
RoundedRectangle(cornerRadius: 6)
|
||||
.fill(Color(nsColor: .textBackgroundColor))
|
||||
)
|
||||
.overlay(
|
||||
RoundedRectangle(cornerRadius: 6)
|
||||
.strokeBorder(Color(nsColor: .separatorColor), lineWidth: 1)
|
||||
)
|
||||
|
||||
Text("Diagnostic logs will be uploaded with your report.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
|
||||
HStack {
|
||||
Spacer()
|
||||
Button("Cancel") { onDismiss() }
|
||||
.keyboardShortcut(.cancelAction)
|
||||
Button("Send") {
|
||||
Task { await send() }
|
||||
}
|
||||
.keyboardShortcut(.defaultAction)
|
||||
}
|
||||
.padding(.top, 4)
|
||||
}
|
||||
}
|
||||
|
||||
private func sendingView(message: String) -> some View {
|
||||
VStack(alignment: .leading, spacing: 12) {
|
||||
HStack(spacing: 10) {
|
||||
ProgressView().controlSize(.small)
|
||||
Text(message)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
HStack {
|
||||
Spacer()
|
||||
Button("Cancel") { onDismiss() }
|
||||
.keyboardShortcut(.cancelAction)
|
||||
.disabled(true)
|
||||
Button("Send") {}
|
||||
.disabled(true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func successView(message: String) -> some View {
|
||||
VStack(alignment: .leading, spacing: 12) {
|
||||
HStack(alignment: .top, spacing: 10) {
|
||||
Image(systemName: "checkmark.circle.fill")
|
||||
.foregroundColor(.green)
|
||||
.font(.title2)
|
||||
Text(message)
|
||||
.fixedSize(horizontal: false, vertical: true)
|
||||
}
|
||||
HStack {
|
||||
Button {
|
||||
openGitHubIssue()
|
||||
} label: {
|
||||
HStack(spacing: 4) {
|
||||
Image(systemName: "arrow.up.right.square")
|
||||
Text("Open GitHub Issue")
|
||||
}
|
||||
}
|
||||
Spacer()
|
||||
Button("Done") { onDismiss() }
|
||||
.keyboardShortcut(.defaultAction)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func failureView(message: String) -> some View {
|
||||
VStack(alignment: .leading, spacing: 12) {
|
||||
HStack(alignment: .top, spacing: 10) {
|
||||
Image(systemName: "exclamationmark.triangle.fill")
|
||||
.foregroundColor(.orange)
|
||||
.font(.title2)
|
||||
Text(message)
|
||||
.fixedSize(horizontal: false, vertical: true)
|
||||
}
|
||||
HStack {
|
||||
Spacer()
|
||||
Button("Try Again") {
|
||||
phase = .prompting
|
||||
}
|
||||
Button("Close") { onDismiss() }
|
||||
.keyboardShortcut(.defaultAction)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func send() async {
|
||||
phase = .sending("Collecting logs and uploading…")
|
||||
let service = BugReportService()
|
||||
let description = userDescription.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
do {
|
||||
let outcome = try await service.sendReport(
|
||||
isManual: true,
|
||||
userDescription: description.isEmpty ? nil : description
|
||||
)
|
||||
if outcome.success {
|
||||
phase = .success(outcome.message)
|
||||
} else {
|
||||
phase = .failure(outcome.message)
|
||||
}
|
||||
} catch {
|
||||
phase = .failure(error.localizedDescription)
|
||||
}
|
||||
}
|
||||
|
||||
private func openGitHubIssue() {
|
||||
let description = userDescription.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
|
||||
var bodyParts: [String] = []
|
||||
bodyParts.append("## Describe the bug")
|
||||
bodyParts.append("")
|
||||
if !description.isEmpty {
|
||||
bodyParts.append(description)
|
||||
} else {
|
||||
bodyParts.append("A clear and concise description of what the bug is.")
|
||||
}
|
||||
bodyParts.append("")
|
||||
bodyParts.append("## Environment")
|
||||
bodyParts.append("")
|
||||
bodyParts.append("- macOS Version: \(ProcessInfo.processInfo.operatingSystemVersionString)")
|
||||
bodyParts.append("- EXO Version: \(buildTag) (\(buildCommit))")
|
||||
bodyParts.append("")
|
||||
bodyParts.append("## Additional context")
|
||||
bodyParts.append("")
|
||||
bodyParts.append("A bug report with diagnostic logs was submitted via the app.")
|
||||
|
||||
let body = bodyParts.joined(separator: "\n")
|
||||
|
||||
var components = URLComponents(string: "https://github.com/exo-explore/exo/issues/new")!
|
||||
components.queryItems = [
|
||||
URLQueryItem(name: "template", value: "bug_report.md"),
|
||||
URLQueryItem(name: "title", value: "[BUG] "),
|
||||
URLQueryItem(name: "body", value: body),
|
||||
URLQueryItem(name: "labels", value: "bug"),
|
||||
]
|
||||
|
||||
if let url = components.url {
|
||||
NSWorkspace.shared.open(url)
|
||||
}
|
||||
}
|
||||
|
||||
private var buildTag: String {
|
||||
Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
|
||||
}
|
||||
|
||||
private var buildCommit: String {
|
||||
Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown"
|
||||
}
|
||||
}
|
||||
@@ -1,200 +0,0 @@
|
||||
import AppKit
|
||||
import SwiftUI
|
||||
|
||||
/// A popover callout anchored to the menu bar icon on every launch,
|
||||
/// pointing the user to the web dashboard with an arrow connecting to the icon.
|
||||
@MainActor
|
||||
final class FirstLaunchPopout {
|
||||
private var popover: NSPopover?
|
||||
private var countdownTask: Task<Void, Never>?
|
||||
private static let dashboardURL = "http://localhost:52415/"
|
||||
|
||||
/// Called when the user completes onboarding (clicks Open Dashboard or dismisses).
|
||||
var onComplete: (() -> Void)?
|
||||
|
||||
func show() {
|
||||
guard popover == nil else { return }
|
||||
|
||||
// The status bar button may not exist yet on first launch; retry generously.
|
||||
showWithRetry(attemptsRemaining: 15)
|
||||
}
|
||||
|
||||
private func showWithRetry(attemptsRemaining: Int) {
|
||||
guard attemptsRemaining > 0 else {
|
||||
// Exhausted retries — fall back to just opening the dashboard directly.
|
||||
openDashboard()
|
||||
onComplete?()
|
||||
return
|
||||
}
|
||||
|
||||
guard let button = Self.findStatusItemButton() else {
|
||||
DispatchQueue.main.asyncAfter(deadline: .now() + 1.0) { [weak self] in
|
||||
self?.showWithRetry(attemptsRemaining: attemptsRemaining - 1)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
let pop = NSPopover()
|
||||
pop.behavior = .applicationDefined
|
||||
pop.animates = true
|
||||
pop.contentSize = NSSize(width: 280, height: 120)
|
||||
pop.contentViewController = NSHostingController(
|
||||
rootView: WelcomeCalloutView(
|
||||
countdownDuration: 10,
|
||||
onDismiss: { [weak self] in
|
||||
self?.onComplete?()
|
||||
self?.dismiss()
|
||||
},
|
||||
onOpen: { [weak self] in
|
||||
self?.openDashboard()
|
||||
self?.onComplete?()
|
||||
self?.dismiss()
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
self.popover = pop
|
||||
pop.show(relativeTo: button.bounds, of: button, preferredEdge: .minY)
|
||||
|
||||
// Auto-open dashboard after 10s then dismiss
|
||||
countdownTask = Task {
|
||||
try? await Task.sleep(nanoseconds: 10_000_000_000)
|
||||
if !Task.isCancelled {
|
||||
openDashboard()
|
||||
onComplete?()
|
||||
dismiss()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func dismiss() {
|
||||
countdownTask?.cancel()
|
||||
countdownTask = nil
|
||||
guard let pop = popover else { return }
|
||||
popover = nil
|
||||
pop.performClose(nil)
|
||||
}
|
||||
|
||||
private func openDashboard() {
|
||||
guard let url = URL(string: Self.dashboardURL) else { return }
|
||||
NSWorkspace.shared.open(url)
|
||||
}
|
||||
|
||||
/// Finds the NSStatusBarButton created by SwiftUI's MenuBarExtra.
|
||||
/// Walks the view hierarchy to find the actual button rather than the content view.
|
||||
private static func findStatusItemButton() -> NSView? {
|
||||
for window in NSApp.windows {
|
||||
let className = NSStringFromClass(type(of: window))
|
||||
// Match NSStatusBarWindow or any internal SwiftUI status bar window
|
||||
guard className.contains("StatusBar") || className.contains("MenuBarExtra") else {
|
||||
continue
|
||||
}
|
||||
if let content = window.contentView {
|
||||
if let button = findButton(in: content) {
|
||||
return button
|
||||
}
|
||||
// Fall back to the content view itself if it has a non-zero frame
|
||||
if content.frame.width > 0 {
|
||||
return content
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
/// Recursively searches the view hierarchy for an NSStatusBarButton.
|
||||
private static func findButton(in view: NSView) -> NSView? {
|
||||
let className = NSStringFromClass(type(of: view))
|
||||
if className.contains("StatusBarButton") || className.contains("StatusItem") {
|
||||
return view
|
||||
}
|
||||
for subview in view.subviews {
|
||||
if let found = findButton(in: subview) {
|
||||
return found
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
/// Minimal welcome callout — friendly pointer, not a wall of text.
|
||||
/// Rendered inside the NSPopover which provides its own chrome and arrow.
|
||||
private struct WelcomeCalloutView: View {
|
||||
let countdownDuration: Int
|
||||
let onDismiss: () -> Void
|
||||
let onOpen: () -> Void
|
||||
@State private var countdown: Int
|
||||
@State private var timerTask: Task<Void, Never>?
|
||||
|
||||
init(countdownDuration: Int, onDismiss: @escaping () -> Void, onOpen: @escaping () -> Void) {
|
||||
self.countdownDuration = countdownDuration
|
||||
self.onDismiss = onDismiss
|
||||
self.onOpen = onOpen
|
||||
self._countdown = State(initialValue: countdownDuration)
|
||||
}
|
||||
|
||||
var body: some View {
|
||||
VStack(alignment: .leading, spacing: 10) {
|
||||
HStack(alignment: .top) {
|
||||
Text("EXO is running")
|
||||
.font(.system(.headline, design: .rounded))
|
||||
.fontWeight(.semibold)
|
||||
.foregroundColor(.primary)
|
||||
Spacer()
|
||||
Button {
|
||||
onDismiss()
|
||||
} label: {
|
||||
Image(systemName: "xmark.circle.fill")
|
||||
.font(.system(size: 14))
|
||||
.foregroundStyle(.tertiary)
|
||||
}
|
||||
.buttonStyle(.plain)
|
||||
}
|
||||
|
||||
Text("Run your first model here:")
|
||||
.font(.system(.subheadline, design: .default))
|
||||
.foregroundColor(.secondary)
|
||||
|
||||
HStack {
|
||||
Button {
|
||||
onOpen()
|
||||
} label: {
|
||||
Label("Open Dashboard", systemImage: "arrow.up.right.square")
|
||||
.font(.system(.caption, design: .default))
|
||||
.fontWeight(.medium)
|
||||
}
|
||||
.buttonStyle(.borderedProminent)
|
||||
.tint(.accentColor)
|
||||
.controlSize(.small)
|
||||
|
||||
Spacer()
|
||||
|
||||
if countdown > 0 {
|
||||
Text("Launching in \(countdown) secs...")
|
||||
.font(.system(.caption2, design: .default))
|
||||
.foregroundColor(.secondary.opacity(0.6))
|
||||
.monospacedDigit()
|
||||
}
|
||||
}
|
||||
}
|
||||
.padding(14)
|
||||
.onAppear {
|
||||
startCountdown()
|
||||
}
|
||||
.onDisappear {
|
||||
timerTask?.cancel()
|
||||
timerTask = nil
|
||||
}
|
||||
}
|
||||
|
||||
private func startCountdown() {
|
||||
timerTask = Task {
|
||||
while countdown > 0 {
|
||||
try? await Task.sleep(nanoseconds: 1_000_000_000)
|
||||
if !Task.isCancelled {
|
||||
countdown -= 1
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,718 +0,0 @@
|
||||
import AppKit
|
||||
import SwiftUI
|
||||
|
||||
/// Native macOS Settings window following Apple HIG.
|
||||
/// Organized into General, Model, Advanced, and About sections.
|
||||
struct SettingsView: View {
|
||||
@EnvironmentObject private var controller: ExoProcessController
|
||||
@EnvironmentObject private var updater: SparkleUpdater
|
||||
@EnvironmentObject private var networkStatusService: NetworkStatusService
|
||||
@EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService
|
||||
@EnvironmentObject private var stateService: ClusterStateService
|
||||
|
||||
@State private var pendingNamespace: String = ""
|
||||
@State private var pendingHFToken: String = ""
|
||||
@State private var pendingHFEndpoint: String = ""
|
||||
@State private var pendingEnableImageModels = false
|
||||
@State private var pendingOfflineMode = false
|
||||
@State private var pendingFastSynchEnabled = false
|
||||
@State private var pendingDefaultModelsDir: String = ""
|
||||
@State private var pendingAdditionalModelsDirs: String = ""
|
||||
@State private var pendingReadOnlyModelsDirs: String = ""
|
||||
@State private var pendingCustomEnvironmentVariables: [CustomEnvironmentVariable] = []
|
||||
@State private var needsRestart = false
|
||||
@State private var uninstallInProgress = false
|
||||
|
||||
var body: some View {
|
||||
TabView {
|
||||
generalTab
|
||||
.tabItem {
|
||||
Label("General", systemImage: "gear")
|
||||
}
|
||||
modelTab
|
||||
.tabItem {
|
||||
Label("Model", systemImage: "cube")
|
||||
}
|
||||
advancedTab
|
||||
.tabItem {
|
||||
Label("Advanced", systemImage: "wrench.and.screwdriver")
|
||||
}
|
||||
environmentTab
|
||||
.tabItem {
|
||||
Label("Environment", systemImage: "terminal")
|
||||
}
|
||||
aboutTab
|
||||
.tabItem {
|
||||
Label("About", systemImage: "info.circle")
|
||||
}
|
||||
}
|
||||
.frame(width: 640, height: 560)
|
||||
.onAppear {
|
||||
pendingNamespace = controller.customNamespace
|
||||
pendingHFToken = controller.hfToken
|
||||
pendingHFEndpoint = controller.hfEndpoint
|
||||
pendingEnableImageModels = controller.enableImageModels
|
||||
pendingOfflineMode = controller.offlineMode
|
||||
pendingFastSynchEnabled = controller.fastSynchEnabled
|
||||
pendingDefaultModelsDir = controller.defaultModelsDir
|
||||
pendingAdditionalModelsDirs = controller.additionalModelsDirs
|
||||
pendingReadOnlyModelsDirs = controller.readOnlyModelsDirs
|
||||
pendingCustomEnvironmentVariables = controller.customEnvironmentVariables
|
||||
needsRestart = false
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - General Tab
|
||||
|
||||
private var generalTab: some View {
|
||||
Form {
|
||||
Section {
|
||||
LabeledContent("Cluster Namespace") {
|
||||
TextField("", text: $pendingNamespace, prompt: Text("default"))
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.frame(width: 260)
|
||||
}
|
||||
Text("Nodes with the same namespace form a cluster. Leave empty for default.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
|
||||
Section {
|
||||
LabeledContent("HuggingFace Token") {
|
||||
SecureField("", text: $pendingHFToken, prompt: Text("optional"))
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.frame(width: 260)
|
||||
}
|
||||
Text("Required for gated models. Get yours at huggingface.co/settings/tokens")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
|
||||
Section {
|
||||
LabeledContent("HuggingFace Endpoint") {
|
||||
TextField("", text: $pendingHFEndpoint, prompt: Text("default"))
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.frame(width: 260)
|
||||
}
|
||||
Text("Defaults to huggingface.co. Use a mirror (e.g. hf-mirror.com) for China.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
|
||||
Section {
|
||||
Toggle("Offline Mode", isOn: $pendingOfflineMode)
|
||||
Text("Skip internet checks and use only locally available models.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
|
||||
Section {
|
||||
HStack {
|
||||
Spacer()
|
||||
Button("Save & Restart") {
|
||||
applyGeneralSettings()
|
||||
}
|
||||
.disabled(!hasGeneralChanges)
|
||||
}
|
||||
}
|
||||
}
|
||||
.formStyle(.grouped)
|
||||
.padding()
|
||||
}
|
||||
|
||||
// MARK: - Model Tab
|
||||
|
||||
private var modelTab: some View {
|
||||
Form {
|
||||
Section {
|
||||
Toggle("Enable Image Models (experimental)", isOn: $pendingEnableImageModels)
|
||||
Text("Allow text-to-image and image-to-image models in the model picker.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
|
||||
Section {
|
||||
HStack {
|
||||
Spacer()
|
||||
Button("Save & Restart") {
|
||||
applyModelSettings()
|
||||
}
|
||||
.disabled(!hasModelChanges)
|
||||
}
|
||||
}
|
||||
}
|
||||
.formStyle(.grouped)
|
||||
.padding()
|
||||
}
|
||||
|
||||
// MARK: - Advanced Tab
|
||||
|
||||
private var advancedTab: some View {
|
||||
Form {
|
||||
Section("Performance") {
|
||||
Toggle("Fast Synch Enabled", isOn: $pendingFastSynchEnabled)
|
||||
Text(
|
||||
"Experimental: enables fast CPU to GPU synchronization. Can sometimes cause a \"GPU lock\" where inference hangs for ~10 seconds before starting. Necessary for low latency with RDMA and Tensor Parallelism."
|
||||
)
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
|
||||
HStack {
|
||||
Spacer()
|
||||
Button("Save & Restart") {
|
||||
applyAdvancedSettings()
|
||||
}
|
||||
.disabled(!hasAdvancedChanges)
|
||||
}
|
||||
}
|
||||
|
||||
Section("Onboarding") {
|
||||
HStack {
|
||||
VStack(alignment: .leading) {
|
||||
Text("Reset Onboarding")
|
||||
Text("Opens the dashboard and resets the onboarding wizard.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
Spacer()
|
||||
Button("Reset") {
|
||||
guard let url = URL(string: "http://localhost:52415/?reset-onboarding")
|
||||
else { return }
|
||||
NSWorkspace.shared.open(url)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Section("Debug Info") {
|
||||
LabeledContent("Thunderbolt Bridge") {
|
||||
Text(thunderboltStatusText)
|
||||
.foregroundColor(thunderboltStatusColor)
|
||||
}
|
||||
|
||||
VStack(alignment: .leading, spacing: 2) {
|
||||
clusterThunderboltBridgeView
|
||||
}
|
||||
|
||||
VStack(alignment: .leading, spacing: 2) {
|
||||
interfaceIpList
|
||||
}
|
||||
|
||||
VStack(alignment: .leading, spacing: 2) {
|
||||
rdmaStatusView
|
||||
}
|
||||
}
|
||||
|
||||
Section("Danger Zone") {
|
||||
Button(role: .destructive) {
|
||||
showUninstallConfirmationAlert()
|
||||
} label: {
|
||||
HStack {
|
||||
Text("Uninstall EXO")
|
||||
Spacer()
|
||||
Image(systemName: "trash")
|
||||
.imageScale(.small)
|
||||
}
|
||||
}
|
||||
.disabled(uninstallInProgress)
|
||||
}
|
||||
}
|
||||
.formStyle(.grouped)
|
||||
.padding()
|
||||
}
|
||||
|
||||
// MARK: - Environment Tab
|
||||
|
||||
private var environmentTab: some View {
|
||||
Form {
|
||||
Section("Models Directories") {
|
||||
LabeledContent("Default Models Directory") {
|
||||
TextField(
|
||||
"",
|
||||
text: $pendingDefaultModelsDir,
|
||||
prompt: Text("~/.exo/models")
|
||||
)
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.font(.system(.body, design: .monospaced))
|
||||
.frame(width: 260)
|
||||
}
|
||||
Text("Sets EXO_DEFAULT_MODELS_DIR. Where models are downloaded.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
|
||||
LabeledContent("Additional Directories") {
|
||||
TextField(
|
||||
"",
|
||||
text: $pendingAdditionalModelsDirs,
|
||||
prompt: Text("optional, colon-separated")
|
||||
)
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.font(.system(.body, design: .monospaced))
|
||||
.frame(width: 260)
|
||||
}
|
||||
Text("Sets EXO_MODELS_DIRS. Extra writable model directories.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
|
||||
LabeledContent("Read-Only Directories") {
|
||||
TextField(
|
||||
"",
|
||||
text: $pendingReadOnlyModelsDirs,
|
||||
prompt: Text("optional, colon-separated")
|
||||
)
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.font(.system(.body, design: .monospaced))
|
||||
.frame(width: 260)
|
||||
}
|
||||
Text("Sets EXO_MODELS_READ_ONLY_DIRS. Never written to.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
|
||||
Section("Custom Environment Variables") {
|
||||
Text(
|
||||
"Escape hatch for env vars that don't have typed fields above. "
|
||||
+ "Values here override the typed fields on conflict."
|
||||
)
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
|
||||
if pendingCustomEnvironmentVariables.isEmpty {
|
||||
Text("No custom variables.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
} else {
|
||||
ForEach($pendingCustomEnvironmentVariables) { $variable in
|
||||
HStack(alignment: .center, spacing: 8) {
|
||||
VStack(spacing: 4) {
|
||||
TextField("key", text: $variable.key)
|
||||
.labelsHidden()
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.font(.system(.body, design: .monospaced))
|
||||
TextField("value", text: $variable.value)
|
||||
.labelsHidden()
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.font(.system(.body, design: .monospaced))
|
||||
}
|
||||
VStack(spacing: 4) {
|
||||
Button {
|
||||
pendingCustomEnvironmentVariables.removeAll {
|
||||
$0.id == variable.id
|
||||
}
|
||||
} label: {
|
||||
Image(systemName: "minus.circle")
|
||||
}
|
||||
.buttonStyle(.borderless)
|
||||
.help("Remove variable")
|
||||
if !isValidEnvironmentVariableName(variable.key) {
|
||||
Image(systemName: "exclamationmark.triangle.fill")
|
||||
.foregroundColor(.orange)
|
||||
.help(
|
||||
"Invalid environment variable name. "
|
||||
+ "Must match [A-Za-z_][A-Za-z0-9_]*."
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
HStack {
|
||||
Button {
|
||||
pendingCustomEnvironmentVariables.append(
|
||||
CustomEnvironmentVariable()
|
||||
)
|
||||
} label: {
|
||||
Label("Add Variable", systemImage: "plus")
|
||||
}
|
||||
Spacer()
|
||||
}
|
||||
}
|
||||
|
||||
Section {
|
||||
HStack {
|
||||
Spacer()
|
||||
Button("Save & Restart") {
|
||||
applyEnvironmentSettings()
|
||||
}
|
||||
.disabled(!hasEnvironmentChanges)
|
||||
}
|
||||
}
|
||||
}
|
||||
.formStyle(.grouped)
|
||||
.padding()
|
||||
}
|
||||
|
||||
// MARK: - About Tab
|
||||
|
||||
private var aboutTab: some View {
|
||||
Form {
|
||||
Section {
|
||||
LabeledContent("Version") {
|
||||
Text(buildTag)
|
||||
.textSelection(.enabled)
|
||||
}
|
||||
LabeledContent("Commit") {
|
||||
Text(buildCommit)
|
||||
.font(.system(.body, design: .monospaced))
|
||||
.textSelection(.enabled)
|
||||
}
|
||||
}
|
||||
|
||||
Section {
|
||||
Button("Check for Updates") {
|
||||
updater.checkForUpdates()
|
||||
}
|
||||
}
|
||||
}
|
||||
.formStyle(.grouped)
|
||||
.padding()
|
||||
}
|
||||
|
||||
// MARK: - Debug Info Views (moved from ContentView)
|
||||
|
||||
private var thunderboltStatusText: String {
|
||||
switch networkStatusService.status.thunderboltBridgeState {
|
||||
case .some(.disabled):
|
||||
return "Disabled"
|
||||
case .some(.deleted):
|
||||
return "Deleted"
|
||||
case .some(.enabled):
|
||||
return "Enabled"
|
||||
case nil:
|
||||
return "Unknown"
|
||||
}
|
||||
}
|
||||
|
||||
private var thunderboltStatusColor: Color {
|
||||
switch networkStatusService.status.thunderboltBridgeState {
|
||||
case .some(.disabled), .some(.deleted):
|
||||
return .green
|
||||
case .some(.enabled):
|
||||
return .red
|
||||
case nil:
|
||||
return .secondary
|
||||
}
|
||||
}
|
||||
|
||||
private var clusterThunderboltBridgeView: some View {
|
||||
let bridgeStatuses = stateService.latestSnapshot?.nodeThunderboltBridge ?? [:]
|
||||
let localNodeId = stateService.localNodeId
|
||||
let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:]
|
||||
|
||||
return VStack(alignment: .leading, spacing: 1) {
|
||||
if bridgeStatuses.isEmpty {
|
||||
Text("Cluster TB Bridge: No data")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
} else {
|
||||
Text("Cluster TB Bridge Status:")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
ForEach(Array(bridgeStatuses.keys.sorted()), id: \.self) { nodeId in
|
||||
if let status = bridgeStatuses[nodeId] {
|
||||
let nodeName =
|
||||
nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8))
|
||||
let isLocal = nodeId == localNodeId
|
||||
let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):"
|
||||
let statusText =
|
||||
!status.exists
|
||||
? "N/A"
|
||||
: (status.enabled ? "Enabled" : "Disabled")
|
||||
let color: Color =
|
||||
!status.exists
|
||||
? .secondary
|
||||
: (status.enabled ? .red : .green)
|
||||
Text("\(prefix) \(statusText)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(color)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private var interfaceIpList: some View {
|
||||
let statuses = networkStatusService.status.interfaceStatuses
|
||||
return VStack(alignment: .leading, spacing: 1) {
|
||||
Text("Interfaces (en0–en7):")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
if statuses.isEmpty {
|
||||
Text(" Unknown")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
} else {
|
||||
ForEach(statuses, id: \.interfaceName) { status in
|
||||
let ipText = status.ipAddress ?? "No IP"
|
||||
Text(" \(status.interfaceName): \(ipText)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(status.ipAddress == nil ? .red : .green)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private var rdmaStatusView: some View {
|
||||
let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:]
|
||||
let localNodeId = stateService.localNodeId
|
||||
let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:]
|
||||
let localDevices = networkStatusService.status.localRdmaDevices
|
||||
let localPorts = networkStatusService.status.localRdmaActivePorts
|
||||
|
||||
return VStack(alignment: .leading, spacing: 1) {
|
||||
if rdmaStatuses.isEmpty {
|
||||
Text("Cluster RDMA: No data")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
} else {
|
||||
Text("Cluster RDMA Status:")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
ForEach(Array(rdmaStatuses.keys.sorted()), id: \.self) { nodeId in
|
||||
if let status = rdmaStatuses[nodeId] {
|
||||
let nodeName =
|
||||
nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8))
|
||||
let isLocal = nodeId == localNodeId
|
||||
let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):"
|
||||
let statusText = status.enabled ? "Enabled" : "Disabled"
|
||||
let color: Color = status.enabled ? .green : .orange
|
||||
Text("\(prefix) \(statusText)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(color)
|
||||
}
|
||||
}
|
||||
}
|
||||
if !localDevices.isEmpty {
|
||||
Text(" Local Devices: \(localDevices.joined(separator: ", "))")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
if !localPorts.isEmpty {
|
||||
Text(" Local Active Ports:")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
ForEach(localPorts, id: \.device) { port in
|
||||
Text(" \(port.device) port \(port.port): \(port.state)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.green)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Actions
|
||||
|
||||
private func showUninstallConfirmationAlert() {
|
||||
let alert = NSAlert()
|
||||
alert.messageText = "Uninstall EXO"
|
||||
alert.informativeText = """
|
||||
This will remove EXO and all its components:
|
||||
|
||||
• Network configuration daemon
|
||||
• Launch at login registration
|
||||
• EXO network location
|
||||
• EXO data directory (~/.exo)
|
||||
|
||||
The app will be moved to Trash.
|
||||
"""
|
||||
alert.alertStyle = .warning
|
||||
|
||||
let checkbox = NSButton(
|
||||
checkboxWithTitle: "Keep downloaded models (~/.exo/models)",
|
||||
target: nil, action: nil)
|
||||
checkbox.state = .off
|
||||
checkbox.sizeToFit()
|
||||
alert.accessoryView = checkbox
|
||||
|
||||
alert.addButton(withTitle: "Uninstall")
|
||||
alert.addButton(withTitle: "Cancel")
|
||||
|
||||
if let uninstallButton = alert.buttons.first {
|
||||
uninstallButton.hasDestructiveAction = true
|
||||
}
|
||||
|
||||
let response = alert.runModal()
|
||||
if response == .alertFirstButtonReturn {
|
||||
performUninstall(keepModels: checkbox.state == .on)
|
||||
}
|
||||
}
|
||||
|
||||
private func performUninstall(keepModels: Bool) {
|
||||
uninstallInProgress = true
|
||||
|
||||
controller.cancelPendingLaunch()
|
||||
controller.stop()
|
||||
stateService.stopPolling()
|
||||
|
||||
DispatchQueue.global(qos: .utility).async {
|
||||
do {
|
||||
try NetworkSetupHelper.uninstall()
|
||||
try Self.removeExoDirectory(keepModels: keepModels)
|
||||
|
||||
DispatchQueue.main.async {
|
||||
LaunchAtLoginHelper.disable()
|
||||
self.moveAppToTrash()
|
||||
|
||||
DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) {
|
||||
NSApplication.shared.terminate(nil)
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
DispatchQueue.main.async {
|
||||
let errorAlert = NSAlert()
|
||||
errorAlert.messageText = "Uninstall Failed"
|
||||
errorAlert.informativeText = error.localizedDescription
|
||||
errorAlert.alertStyle = .critical
|
||||
errorAlert.addButton(withTitle: "OK")
|
||||
errorAlert.runModal()
|
||||
self.uninstallInProgress = false
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private static func removeExoDirectory(keepModels: Bool) throws {
|
||||
let fm = FileManager.default
|
||||
let exoDir = ExoProcessController.exoDirectoryURL
|
||||
guard fm.fileExists(atPath: exoDir.path) else { return }
|
||||
|
||||
if !keepModels {
|
||||
try fm.removeItem(at: exoDir)
|
||||
return
|
||||
}
|
||||
|
||||
let contents = try fm.contentsOfDirectory(
|
||||
at: exoDir, includingPropertiesForKeys: nil, options: [])
|
||||
for entry in contents where entry.lastPathComponent != "models" {
|
||||
try? fm.removeItem(at: entry)
|
||||
}
|
||||
}
|
||||
|
||||
private func moveAppToTrash() {
|
||||
guard let appURL = Bundle.main.bundleURL as URL? else { return }
|
||||
do {
|
||||
try FileManager.default.trashItem(at: appURL, resultingItemURL: nil)
|
||||
} catch {
|
||||
// If we can't trash the app, that's OK - user can do it manually
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Helpers
|
||||
|
||||
private var hasGeneralChanges: Bool {
|
||||
pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken
|
||||
|| pendingHFEndpoint != controller.hfEndpoint
|
||||
|| pendingOfflineMode != controller.offlineMode
|
||||
}
|
||||
|
||||
private var hasModelChanges: Bool {
|
||||
pendingEnableImageModels != controller.enableImageModels
|
||||
}
|
||||
|
||||
private var hasAdvancedChanges: Bool {
|
||||
pendingFastSynchEnabled != controller.fastSynchEnabled
|
||||
}
|
||||
|
||||
private var hasEnvironmentChanges: Bool {
|
||||
pendingDefaultModelsDir != controller.defaultModelsDir
|
||||
|| pendingAdditionalModelsDirs != controller.additionalModelsDirs
|
||||
|| pendingReadOnlyModelsDirs != controller.readOnlyModelsDirs
|
||||
|| pendingCustomEnvironmentVariables != controller.customEnvironmentVariables
|
||||
}
|
||||
|
||||
private func applyGeneralSettings() {
|
||||
controller.customNamespace = pendingNamespace
|
||||
controller.hfToken = pendingHFToken
|
||||
controller.hfEndpoint = pendingHFEndpoint
|
||||
controller.offlineMode = pendingOfflineMode
|
||||
restartIfRunning()
|
||||
}
|
||||
|
||||
private func applyModelSettings() {
|
||||
controller.enableImageModels = pendingEnableImageModels
|
||||
restartIfRunning()
|
||||
}
|
||||
|
||||
private func applyAdvancedSettings() {
|
||||
controller.fastSynchEnabled = pendingFastSynchEnabled
|
||||
restartIfRunning()
|
||||
}
|
||||
|
||||
private func applyEnvironmentSettings() {
|
||||
controller.defaultModelsDir = pendingDefaultModelsDir.trimmingCharacters(
|
||||
in: .whitespaces)
|
||||
controller.additionalModelsDirs = pendingAdditionalModelsDirs.trimmingCharacters(
|
||||
in: .whitespaces)
|
||||
controller.readOnlyModelsDirs = pendingReadOnlyModelsDirs.trimmingCharacters(
|
||||
in: .whitespaces)
|
||||
|
||||
pendingDefaultModelsDir = controller.defaultModelsDir
|
||||
pendingAdditionalModelsDirs = controller.additionalModelsDirs
|
||||
pendingReadOnlyModelsDirs = controller.readOnlyModelsDirs
|
||||
|
||||
// Trim whitespace from keys and drop empty ones so that the stored
|
||||
// form matches what is actually injected into the child process and
|
||||
// hasEnvironmentChanges doesn't show a stale diff after save.
|
||||
let trimmed: [CustomEnvironmentVariable] =
|
||||
pendingCustomEnvironmentVariables.compactMap { variable in
|
||||
let key = variable.key.trimmingCharacters(in: .whitespaces)
|
||||
guard !key.isEmpty else { return nil }
|
||||
return CustomEnvironmentVariable(
|
||||
id: variable.id, key: key, value: variable.value
|
||||
)
|
||||
}
|
||||
|
||||
// De-duplicate keys, keeping the last occurrence. This matches the
|
||||
// effective semantics of the dictionary assignment in
|
||||
// ExoProcessController.makeEnvironment and avoids silently losing
|
||||
// visible rows after save.
|
||||
var seenKeys = Set<String>()
|
||||
var deduplicatedReversed: [CustomEnvironmentVariable] = []
|
||||
for variable in trimmed.reversed() {
|
||||
if seenKeys.insert(variable.key).inserted {
|
||||
deduplicatedReversed.append(variable)
|
||||
}
|
||||
}
|
||||
let sanitized = Array(deduplicatedReversed.reversed())
|
||||
|
||||
pendingCustomEnvironmentVariables = sanitized
|
||||
controller.customEnvironmentVariables = sanitized
|
||||
|
||||
restartIfRunning()
|
||||
}
|
||||
|
||||
/// Validates a POSIX-style environment variable name:
|
||||
/// `[A-Za-z_][A-Za-z0-9_]*`. Uses an ASCII-only charset so that
|
||||
/// Unicode letters (e.g. `ñ`, Cyrillic) are rejected in line with what
|
||||
/// the help tooltip advertises. Empty strings are treated as valid
|
||||
/// here so that a freshly added blank row does not immediately look
|
||||
/// broken; the save step filters empty keys out instead.
|
||||
private func isValidEnvironmentVariableName(_ key: String) -> Bool {
|
||||
if key.isEmpty { return true }
|
||||
let headAllowed = CharacterSet(
|
||||
charactersIn: "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz_"
|
||||
)
|
||||
let tailAllowed = headAllowed.union(CharacterSet(charactersIn: "0123456789"))
|
||||
guard let first = key.unicodeScalars.first, headAllowed.contains(first) else {
|
||||
return false
|
||||
}
|
||||
for scalar in key.unicodeScalars.dropFirst() {
|
||||
if !tailAllowed.contains(scalar) { return false }
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
private func restartIfRunning() {
|
||||
if controller.status == .running || controller.status == .starting {
|
||||
controller.restart()
|
||||
}
|
||||
}
|
||||
|
||||
private var buildTag: String {
|
||||
Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
|
||||
}
|
||||
|
||||
private var buildCommit: String {
|
||||
Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown"
|
||||
}
|
||||
}
|
||||
@@ -1,47 +0,0 @@
|
||||
import AppKit
|
||||
import SwiftUI
|
||||
|
||||
/// Manages a standalone native macOS Settings window.
|
||||
/// Ensures only one instance exists and brings it to front on repeated opens.
|
||||
@MainActor
|
||||
final class SettingsWindowController: ObservableObject {
|
||||
private var window: NSWindow?
|
||||
|
||||
func open(
|
||||
controller: ExoProcessController,
|
||||
updater: SparkleUpdater,
|
||||
networkStatusService: NetworkStatusService,
|
||||
thunderboltBridgeService: ThunderboltBridgeService,
|
||||
stateService: ClusterStateService
|
||||
) {
|
||||
if let existing = window, existing.isVisible {
|
||||
existing.makeKeyAndOrderFront(nil)
|
||||
NSApp.activate(ignoringOtherApps: true)
|
||||
return
|
||||
}
|
||||
|
||||
let settingsView = SettingsView()
|
||||
.environmentObject(controller)
|
||||
.environmentObject(updater)
|
||||
.environmentObject(networkStatusService)
|
||||
.environmentObject(thunderboltBridgeService)
|
||||
.environmentObject(stateService)
|
||||
|
||||
let hostingView = NSHostingView(rootView: settingsView)
|
||||
|
||||
let newWindow = NSWindow(
|
||||
contentRect: NSRect(x: 0, y: 0, width: 640, height: 560),
|
||||
styleMask: [.titled, .closable],
|
||||
backing: .buffered,
|
||||
defer: false
|
||||
)
|
||||
newWindow.title = "EXO Settings"
|
||||
newWindow.contentView = hostingView
|
||||
newWindow.center()
|
||||
newWindow.isReleasedWhenClosed = false
|
||||
newWindow.makeKeyAndOrderFront(nil)
|
||||
NSApp.activate(ignoringOtherApps: true)
|
||||
|
||||
window = newWindow
|
||||
}
|
||||
}
|
||||
@@ -3,55 +3,25 @@
|
||||
# EXO Uninstaller Script
|
||||
#
|
||||
# This script removes all EXO system components that persist after deleting the app.
|
||||
# Run with: sudo ./uninstall-exo.sh [--keep-models]
|
||||
#
|
||||
# Options:
|
||||
# --keep-models Preserve ~/.exo/models when removing the EXO data directory.
|
||||
# Run with: sudo ./uninstall-exo.sh
|
||||
#
|
||||
# Components removed:
|
||||
# - LaunchDaemon: /Library/LaunchDaemons/io.exo.networksetup.plist
|
||||
# - Network script: /Library/Application Support/EXO/
|
||||
# - Log files: /var/log/io.exo.networksetup.*
|
||||
# - Network location: "exo"
|
||||
# - EXO data directory: ~/.exo (or all of ~/.exo except models/ when --keep-models is set)
|
||||
# - Launch at login registration
|
||||
#
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
KEEP_MODELS=0
|
||||
for arg in "$@"; do
|
||||
case "$arg" in
|
||||
--keep-models)
|
||||
KEEP_MODELS=1
|
||||
;;
|
||||
-h | --help)
|
||||
echo "Usage: sudo ./uninstall-exo.sh [--keep-models]"
|
||||
echo " --keep-models Preserve ~/.exo/models when removing the EXO data directory."
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "Unknown argument: $arg" >&2
|
||||
echo "Usage: sudo ./uninstall-exo.sh [--keep-models]" >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
LABEL="io.exo.networksetup"
|
||||
# Current script path. Older installs used a different filename; keep the
|
||||
# legacy path here so a fresh uninstall still cleans up upgraded machines.
|
||||
CURRENT_SCRIPT_DEST="/Library/Application Support/EXO/disable_bridge.sh"
|
||||
LEGACY_SCRIPT_DEST="/Library/Application Support/EXO/disable_bridge_enable_dhcp.sh"
|
||||
SCRIPT_DEST="/Library/Application Support/EXO/disable_bridge_enable_dhcp.sh"
|
||||
PLIST_DEST="/Library/LaunchDaemons/io.exo.networksetup.plist"
|
||||
LOG_OUT="/var/log/${LABEL}.log"
|
||||
LOG_ERR="/var/log/${LABEL}.err.log"
|
||||
APP_BUNDLE_ID="io.exo.EXO"
|
||||
|
||||
# Resolve the invoking user's home, even when run via sudo.
|
||||
USER_HOME="$(eval echo "~${SUDO_USER:-$USER}")"
|
||||
EXO_DIR="$USER_HOME/.exo"
|
||||
|
||||
# Colors for output
|
||||
RED='\033[0;31m'
|
||||
GREEN='\033[0;32m'
|
||||
@@ -99,17 +69,11 @@ else
|
||||
echo_warn "LaunchDaemon plist not found (already removed?)"
|
||||
fi
|
||||
|
||||
# Remove the script (current and legacy filenames) — backwards-compatible:
|
||||
# tolerate either, both, or neither being present.
|
||||
removed_any_script=0
|
||||
for script in "$CURRENT_SCRIPT_DEST" "$LEGACY_SCRIPT_DEST"; do
|
||||
if [[ -f $script ]]; then
|
||||
rm -f "$script"
|
||||
echo_info "Removed network setup script: $script"
|
||||
removed_any_script=1
|
||||
fi
|
||||
done
|
||||
if [[ $removed_any_script -eq 0 ]]; then
|
||||
# Remove the script and parent directory
|
||||
if [[ -f $SCRIPT_DEST ]]; then
|
||||
rm -f "$SCRIPT_DEST"
|
||||
echo_info "Removed network setup script"
|
||||
else
|
||||
echo_warn "Network setup script not found (already removed?)"
|
||||
fi
|
||||
|
||||
@@ -151,22 +115,6 @@ if networksetup -listnetworkservices 2>/dev/null | grep -q "Thunderbolt Bridge";
|
||||
echo_info "Re-enabled Thunderbolt Bridge"
|
||||
fi
|
||||
|
||||
# Remove EXO data directory (~/.exo)
|
||||
EXO_DIR_REMOVED=""
|
||||
if [[ -d $EXO_DIR ]]; then
|
||||
if [[ $KEEP_MODELS == "1" && -d "$EXO_DIR/models" ]]; then
|
||||
find "$EXO_DIR" -mindepth 1 -maxdepth 1 ! -name models -exec rm -rf {} +
|
||||
EXO_DIR_REMOVED="kept_models"
|
||||
echo_info "Removed ~/.exo (preserved models/)"
|
||||
else
|
||||
rm -rf "$EXO_DIR"
|
||||
EXO_DIR_REMOVED="full"
|
||||
echo_info "Removed ~/.exo"
|
||||
fi
|
||||
else
|
||||
echo_warn "~/.exo not found (already removed?)"
|
||||
fi
|
||||
|
||||
# Note about launch at login registration
|
||||
# SMAppService-based login items cannot be removed from a shell script.
|
||||
# They can only be unregistered from within the app itself or manually via System Settings.
|
||||
@@ -196,10 +144,6 @@ echo " • Network setup LaunchDaemon"
|
||||
echo " • Network configuration script"
|
||||
echo " • Log files"
|
||||
echo " • 'exo' network location"
|
||||
case "$EXO_DIR_REMOVED" in
|
||||
full) echo " • EXO data directory (~/.exo)" ;;
|
||||
kept_models) echo " • EXO data directory (~/.exo, models preserved)" ;;
|
||||
esac
|
||||
echo ""
|
||||
echo "Your network has been restored to use the 'Automatic' location."
|
||||
echo "Thunderbolt Bridge has been re-enabled (if present)."
|
||||
|
||||
@@ -1,160 +0,0 @@
|
||||
# Exo-Bench — Methodology
|
||||
|
||||
exo bench measures inference throughput and resource consumption of an exo cluster under controlled conditions. It sends prompts to the `/bench/chat/completions` endpoint, collects server-reported timing statistics, and records system-level metrics (power, GPU utilisation, temperature) throughout each run.
|
||||
|
||||
The goal is to have accurate, transparent and reproducible numbers to compare speed and scaling across different models and different setups, and to be able to track these results as optimizations and features are added to EXO.
|
||||
|
||||
Below is a technical summary of how Exo-Bench works. While the methodology and benchmark may change over time, this document will be kept up to date whenever this happens. If you find an issue with the methodology, or would like a feature to be added, please open a GitHub issue!
|
||||
|
||||
---
|
||||
|
||||
## Prompt Construction
|
||||
|
||||
Benchmarks need prompts of an exact token length. Unfortunately, we do not have direct access to the model but just the chat completion endpoint. To get around this fact, we create a request that will tokenise to a certain prompt length.
|
||||
|
||||
This is achieved by:
|
||||
|
||||
1. Tokenising a sample message through the model's `apply_chat_template()` to measure overhead (system tokens, special tokens, chat formatting).
|
||||
2. Binary-searching over a repeated atom string (default `"a "`) to find the content length that produces exactly the target number of tokens after template expansion.
|
||||
3. Returning both the content string and the verified token count.
|
||||
|
||||
The actual token count is recorded in every result row as `pp_tokens`, so downstream analysis can confirm the prompt hit its target.
|
||||
|
||||
Chat template formatting means that it may be impossible to attain very small pp benchmarks. e.g. pp=32 may not work. This tradeoff was made because the result of such a small prompt does not seem very interesting or useful for any real-world use cases.
|
||||
|
||||
---
|
||||
|
||||
## Bench Endpoint
|
||||
|
||||
When a request reaches the server via the `/bench/chat/completions` endpoint, three things change compared to a normal chat completion:
|
||||
|
||||
- **KV prefix cache is disabled by default**. Every request starts from a cold cache, ensuring prefill timing is not affected by prior requests. See [Prefix Cache Mode](#prefix-cache-mode) for the `--use-prefix-cache` option.
|
||||
- **EOS tokens are banned**. A logits processor suppresses all end-of-sequence tokens, forcing the model to generate exactly `max_tokens` tokens. This guarantees consistent generation length for fair TPS comparison — the model cannot short-circuit a run by stopping early.
|
||||
- **No model output parsing**. The bench collection path concatenates raw token text without any model-specific post-processing (thinking tag extraction, structured output handling, etc.). This is to avoid model outputs such as tool parsing or any structural mistakes from breaking the benchmark - we are testing for speed; see Exo-Eval for performance metrics.
|
||||
|
||||
---
|
||||
|
||||
## Timing
|
||||
|
||||
### Prefill TPS
|
||||
|
||||
Measured server-side per task.
|
||||
|
||||
```
|
||||
prefill_tps = num_prompt_tokens / prefill_wall_seconds
|
||||
```
|
||||
|
||||
### Generation TPS
|
||||
|
||||
Measured server-side per task. Each task records wall-clock timestamps as tokens arrive:
|
||||
|
||||
- First generated token: timestamp recorded
|
||||
- Every subsequent token: timestamp updated
|
||||
|
||||
When generation completes:
|
||||
|
||||
```
|
||||
gen_span = last_token_time - first_token_time
|
||||
generation_tps = (completion_tokens - 1) / gen_span
|
||||
```
|
||||
|
||||
The first token is excluded from the numerator because the rate measures inter-token throughput — the time between the first and last token divided by the number of intervals.
|
||||
|
||||
This does mean that tg=1 will not work.
|
||||
|
||||
---
|
||||
|
||||
## Concurrency
|
||||
|
||||
### Single Request
|
||||
|
||||
The client records wall-clock `elapsed_s` around the HTTP round-trip (network latency + server prefill + generation + response serialisation). This is a convenience metric for end-to-end latency. The authoritative TPS numbers come from the server-side per-task timing in the `generation_stats` response.
|
||||
|
||||
### Concurrent Requests
|
||||
|
||||
When `--concurrency N` is set with N > 1, all N requests must hit the server at the same instant. The mechanism:
|
||||
|
||||
1. The prompt is built once and shared across all threads.
|
||||
2. Each thread gets its own HTTP connection.
|
||||
3. A thread barrier blocks all threads until every thread is ready.
|
||||
4. The first thread past the barrier records the batch start time and signals the others.
|
||||
5. All threads use the same start time as their reference, then fire their HTTP request.
|
||||
6. Each thread's `elapsed_s` is measured from the shared start time to its own response completion.
|
||||
|
||||
**Batch wall time** is the maximum `elapsed_s` across all N requests — the time until the last request finishes.
|
||||
|
||||
### Aggregate TPS
|
||||
|
||||
```
|
||||
per_req_tps = max(generation_tps across N concurrent requests)
|
||||
agg_gen_tps = per_req_tps * concurrency
|
||||
```
|
||||
|
||||
`max` is used instead of `mean` because all requests run in parallel against the same model. The fastest request's generation rate represents the system's per-stream throughput capacity; multiplying by concurrency gives aggregate throughput.
|
||||
|
||||
---
|
||||
|
||||
## Prefix Cache Mode
|
||||
|
||||
When `--use-prefix-cache` is passed, the KV prefix cache remains active during benchmarking. This speeds up repeated runs by skipping redundant prefill work, which is useful when prompt processing is not the focus of the benchmark (e.g. when measuring generation throughput or power consumption across many configurations).
|
||||
|
||||
Each response includes a `prefix_cache_hit` field (`"none"`, `"partial"`, or `"exact"`):
|
||||
|
||||
- **none**: Cold prefill — no cached KV state was available. The reported `prompt_tps` is the real prefill throughput.
|
||||
- **partial**: A prefix of the prompt was found in cache. Only the remaining tokens were prefilled. The reported `prompt_tps` reflects the real throughput on the uncached portion. This occurs when multiple ascending `--pp` values share a common prefix (e.g. `--pp 1000,5000` — the 5000-token prompt reuses the 1000-token cache entry and prefills the remaining 4000 tokens).
|
||||
- **exact**: The entire prompt was found in cache (e.g. same `--pp` value on a `--repeat`). No prefill work was done. The reported `prompt_tps` is the TPS from when the cache entry was originally created, not a new measurement.
|
||||
|
||||
**Prompt TPS is approximate in this mode.** Exact-hit runs report the stored TPS from the original cold/partial prefill rather than a freshly measured value. For accurate cold prefill numbers, run without `--use-prefix-cache`.
|
||||
|
||||
Ascending `--pp` order (e.g. `--pp 1000,5000,10000`) gives the most useful data: each size gets a meaningful partial hit except the first which is cold. Descending order produces exact hits with approximate TPS from a longer prompt's original run.
|
||||
|
||||
---
|
||||
|
||||
## Warmup
|
||||
|
||||
Before measurement begins, `--warmup N` (default: 0) discarded requests are sent using the first pp/tg pair. Warmup results are not included in the output.
|
||||
|
||||
---
|
||||
|
||||
## System Metrics
|
||||
|
||||
A background thread polls each node at 1 Hz, collecting:
|
||||
|
||||
- GPU utilisation (%)
|
||||
- Temperature (C)
|
||||
- System power draw (W)
|
||||
- CPU cluster usage (performance and efficiency cores)
|
||||
|
||||
**Energy** is computed via trapezoidal integration of the power samples over each inference window (the wall-clock span of each benchmark request or concurrent batch). Average power is `total_joules / total_inference_seconds`.
|
||||
|
||||
---
|
||||
|
||||
## Output Format
|
||||
|
||||
Results are written as JSON with three top-level keys:
|
||||
|
||||
- **`runs`**: Array of per-request result objects, each containing:
|
||||
- `elapsed_s`, `output_text_preview` (first 200 chars)
|
||||
- `stats`: `{ prompt_tps, generation_tps, prompt_tokens, generation_tokens, peak_memory_usage }`
|
||||
- Placement metadata: `model_id`, `placement_sharding`, `placement_instance_meta`, `placement_nodes`
|
||||
- Run metadata: `pp_tokens`, `tg`, `repeat_index`, `concurrency`, `concurrent_index`
|
||||
- `download_duration_s` (if model was freshly downloaded)
|
||||
- **`cluster`**: Cluster state snapshot at time of benchmarking.
|
||||
- **`system_metrics`**: Per-node time-series samples (GPU, power, temperature).
|
||||
|
||||
---
|
||||
|
||||
## Reproducing Results
|
||||
|
||||
```bash
|
||||
cd bench && uv run python exo_bench.py \
|
||||
--model "mlx-community/Qwen3.5-27B-4bit" \
|
||||
--instance-meta jaccl \
|
||||
--sharding tensor \
|
||||
--min-nodes 2 --max-nodes 2 \
|
||||
--pp 512 4096 --tg 128 \
|
||||
--repeat 3 \
|
||||
--warmup 1
|
||||
```
|
||||
|
||||
Run --help for all the available flags.
|
||||
+3
-1
@@ -2,4 +2,6 @@
|
||||
#
|
||||
# Lists the suite files to include. Each file defines benchmarks
|
||||
# with shared constraints, topology, and default args.
|
||||
include = ["single-m3-ultra.toml"]
|
||||
include = [
|
||||
"single-m3-ultra.toml",
|
||||
]
|
||||
@@ -1,347 +0,0 @@
|
||||
# Model evaluation configurations for exo_eval.
|
||||
#
|
||||
# Each [[model]] entry uses `patterns` — a list of substrings matched
|
||||
# against the model_id. First matching entry wins.
|
||||
#
|
||||
# Required fields:
|
||||
# name, patterns, reasoning
|
||||
#
|
||||
# Optional per-model overrides (CLI flags take priority over these):
|
||||
# temperature, top_p, max_tokens, reasoning_effort, enable_thinking
|
||||
#
|
||||
# Fallback defaults (when no per-model config):
|
||||
# reasoning: temperature=1.0, max_tokens=131072, reasoning_effort="high"
|
||||
# non-reasoning: temperature=0.0, max_tokens=16384
|
||||
#
|
||||
# All per-model values below are sourced from official model cards,
|
||||
# generation_config.json files, and vendor documentation.
|
||||
|
||||
# ─── Qwen3.5 (Feb 2026) ─────────────────────────────────────────────
|
||||
# Source: HuggingFace model cards (Qwen/Qwen3.5-*)
|
||||
# Model card recommends: temp=0.6, top_p=0.95, top_k=20
|
||||
# We omit top_k to match vllm eval (which doesn't set it).
|
||||
# max_tokens=121072 to match vllm eval (131072 context - 10000 safety margin).
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3.5 2B"
|
||||
patterns = ["Qwen3.5-2B"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
max_tokens = 121072
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3.5 9B"
|
||||
patterns = ["Qwen3.5-9B"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
max_tokens = 121072
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3.5 27B"
|
||||
patterns = ["Qwen3.5-27B"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
max_tokens = 121072
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3.5 35B A3B"
|
||||
patterns = ["Qwen3.5-35B-A3B"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
max_tokens = 121072
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3.5 122B A10B"
|
||||
patterns = ["Qwen3.5-122B-A10B"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
max_tokens = 121072
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3.5 397B A17B"
|
||||
patterns = ["Qwen3.5-397B-A17B"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
max_tokens = 121072
|
||||
|
||||
# ─── Qwen3 (Apr 2025) ───────────────────────────────────────────────
|
||||
# Source: HuggingFace model cards (Qwen/Qwen3-*)
|
||||
# Model card recommends: temp=0.6, top_p=0.95, top_k=20
|
||||
# We omit top_k to match vllm eval (which doesn't set it).
|
||||
# Non-thinking: temp=0.7, top_p=0.8
|
||||
# max_tokens: 32768 general, 38912 for complex math/code
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3 0.6B"
|
||||
patterns = ["Qwen3-0.6B"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
max_tokens = 38912
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3 30B A3B"
|
||||
patterns = ["Qwen3-30B-A3B"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
max_tokens = 38912
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3 235B A22B"
|
||||
patterns = ["Qwen3-235B-A22B"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
max_tokens = 38912
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3 Next 80B Thinking"
|
||||
patterns = ["Qwen3-Next-80B-A3B-Thinking"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
max_tokens = 38912
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3 Next 80B Instruct"
|
||||
patterns = ["Qwen3-Next-80B-A3B-Instruct"]
|
||||
reasoning = false
|
||||
temperature = 0.7
|
||||
top_p = 0.8
|
||||
max_tokens = 16384
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3 Coder 480B"
|
||||
patterns = ["Qwen3-Coder-480B"]
|
||||
reasoning = false
|
||||
temperature = 0.7
|
||||
top_p = 0.8
|
||||
max_tokens = 16384
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3 Coder Next"
|
||||
patterns = ["Qwen3-Coder-Next"]
|
||||
reasoning = false
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
max_tokens = 121072
|
||||
|
||||
# ─── GPT-OSS (OpenAI) ───────────────────────────────────────────────
|
||||
# Source: OpenAI GitHub README + HuggingFace discussion #21
|
||||
# temp=1.0, top_p=1.0, NO top_k, NO repetition_penalty
|
||||
# reasoning_effort supported: low/medium/high
|
||||
|
||||
[[model]]
|
||||
name = "GPT-OSS 20B"
|
||||
patterns = ["gpt-oss-20b"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 1.0
|
||||
|
||||
[[model]]
|
||||
name = "GPT-OSS 120B"
|
||||
patterns = ["gpt-oss-120b"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 1.0
|
||||
|
||||
# ─── DeepSeek ────────────────────────────────────────────────────────
|
||||
# Source: https://api-docs.deepseek.com/quick_start/parameter_settings
|
||||
# Coding/Math: temp=0.0, General: temp=1.3, Creative: temp=1.5
|
||||
# NOTE: DeepSeek API applies nonlinear temp mapping. These are API values.
|
||||
# When running model directly: API temp 1.0 = model temp 0.3
|
||||
# We use temp=0.0 for eval (coding/math focus).
|
||||
|
||||
[[model]]
|
||||
name = "DeepSeek V3.1"
|
||||
patterns = ["DeepSeek-V3.1"]
|
||||
reasoning = true
|
||||
temperature = 0.0
|
||||
|
||||
[[model]]
|
||||
name = "DeepSeek V3.2"
|
||||
patterns = ["DeepSeek-V3.2"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
|
||||
# ─── NVIDIA Nemotron ───────────────────────────────────────────────────
|
||||
# Source: HuggingFace model cards
|
||||
# All variants: temp=1.0, top_p=0.95, enable_thinking=true
|
||||
|
||||
[[model]]
|
||||
name = "Nemotron Cascade 2 30B A3B"
|
||||
patterns = ["Nemotron-Cascade-2-30B-A3B"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
|
||||
[[model]]
|
||||
name = "Nemotron 3 Super 120B A12B"
|
||||
patterns = ["Nemotron-3-Super-120B-A12B", "NVIDIA-Nemotron-3-Super-120B-A12B"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
|
||||
# ─── GLM (ZhipuAI / THUDM) ──────────────────────────────────────────
|
||||
# Source: HuggingFace model cards + generation_config.json + docs.z.ai
|
||||
# GLM 4.5+: temp=1.0, top_p=0.95
|
||||
# max_tokens=121072 to match vllm eval (131072 context - 10000 safety margin)
|
||||
|
||||
[[model]]
|
||||
name = "GLM-5"
|
||||
patterns = ["GLM-5"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
max_tokens = 121072
|
||||
|
||||
[[model]]
|
||||
name = "GLM 4.5 Air"
|
||||
patterns = ["GLM-4.5-Air"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
|
||||
[[model]]
|
||||
name = "GLM 4.7"
|
||||
patterns = ["GLM-4.7-"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
max_tokens = 121072
|
||||
# Note: matches both GLM-4.7 and GLM-4.7-Flash
|
||||
|
||||
# ─── Kimi (Moonshot AI) ─────────────────────────────────────────────
|
||||
# Source: HuggingFace model cards (moonshotai/Kimi-K2-*)
|
||||
# K2-Instruct: temp=0.6
|
||||
# K2-Thinking: temp=1.0, max_length=262144
|
||||
# K2.5: thinking temp=1.0, top_p=0.95; instant temp=0.6, top_p=0.95
|
||||
|
||||
[[model]]
|
||||
name = "Kimi K2 Thinking"
|
||||
patterns = ["Kimi-K2-Thinking"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
max_tokens = 131072
|
||||
|
||||
[[model]]
|
||||
name = "Kimi K2.5"
|
||||
patterns = ["Kimi-K2.5"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
max_tokens = 121072
|
||||
|
||||
[[model]]
|
||||
name = "Kimi K2 Instruct"
|
||||
patterns = ["Kimi-K2-Instruct"]
|
||||
reasoning = false
|
||||
temperature = 0.6
|
||||
|
||||
# ─── MiniMax ─────────────────────────────────────────────────────────
|
||||
# Source: HuggingFace model cards + generation_config.json
|
||||
# All models: temp=1.0, top_p=0.95
|
||||
# max_tokens=90000 to match vllm eval (100000 context - 10000 safety margin)
|
||||
|
||||
[[model]]
|
||||
name = "MiniMax M2.7"
|
||||
patterns = ["MiniMax-M2.7"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
max_tokens = 90000
|
||||
|
||||
[[model]]
|
||||
name = "MiniMax M2.5"
|
||||
patterns = ["MiniMax-M2.5"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
max_tokens = 90000
|
||||
|
||||
[[model]]
|
||||
name = "MiniMax M2.1"
|
||||
patterns = ["MiniMax-M2.1"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
|
||||
# ─── Step (StepFun) ─────────────────────────────────────────────────
|
||||
# Source: HuggingFace model card (stepfun-ai/Step-3.5-Flash)
|
||||
# Reasoning: temp=1.0, top_p=0.95
|
||||
# General chat: temp=0.6, top_p=0.95
|
||||
# We use reasoning settings for eval.
|
||||
|
||||
[[model]]
|
||||
name = "Step 3.5 Flash"
|
||||
patterns = ["Step-3.5-Flash"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
enable_thinking = true
|
||||
max_tokens = 121072
|
||||
|
||||
# ─── Llama (Meta) ───────────────────────────────────────────────────
|
||||
# Source: generation_config.json + meta-llama/llama-models generation.py
|
||||
# All variants: temp=0.6, top_p=0.9
|
||||
|
||||
[[model]]
|
||||
name = "Llama 3.2 1B"
|
||||
patterns = ["Llama-3.2-1B"]
|
||||
reasoning = false
|
||||
temperature = 0.6
|
||||
top_p = 0.9
|
||||
|
||||
[[model]]
|
||||
name = "Llama 3.2 3B"
|
||||
patterns = ["Llama-3.2-3B"]
|
||||
reasoning = false
|
||||
temperature = 0.6
|
||||
top_p = 0.9
|
||||
|
||||
[[model]]
|
||||
name = "Llama 3.1 8B"
|
||||
patterns = ["Llama-3.1-8B", "Meta-Llama-3.1-8B"]
|
||||
reasoning = false
|
||||
temperature = 0.6
|
||||
top_p = 0.9
|
||||
|
||||
[[model]]
|
||||
name = "Llama 3.1 70B"
|
||||
patterns = ["Llama-3.1-70B", "Meta-Llama-3.1-70B"]
|
||||
reasoning = false
|
||||
temperature = 0.6
|
||||
top_p = 0.9
|
||||
|
||||
[[model]]
|
||||
name = "Llama 3.3 70B"
|
||||
patterns = ["Llama-3.3-70B", "llama-3.3-70b"]
|
||||
reasoning = false
|
||||
temperature = 0.6
|
||||
top_p = 0.9
|
||||
File diff suppressed because it is too large.
Load diff
+567
-601
File diff suppressed because it is too large.
Load diff
-1753
File diff suppressed because it is too large.
Load diff
@@ -1,625 +0,0 @@
|
||||
# type: ignore
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import http.client
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from typing import Any
|
||||
from urllib.parse import urlencode
|
||||
|
||||
from loguru import logger
|
||||
|
||||
_SETTLE_INITIAL_BACKOFF_S = 1.0
|
||||
_SETTLE_MAX_BACKOFF_S = 60.0
|
||||
_SETTLE_BACKOFF_MULTIPLIER = 2.0
|
||||
|
||||
|
||||
class ExoHttpError(RuntimeError):
|
||||
def __init__(self, status: int, reason: str, body_preview: str):
|
||||
super().__init__(f"HTTP {status} {reason}: {body_preview}")
|
||||
self.status = status
|
||||
|
||||
|
||||
class ExoClient:
|
||||
def __init__(self, host: str, port: int, timeout_s: float = 7200.0):
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.timeout_s = timeout_s
|
||||
|
||||
def request_json(
|
||||
self,
|
||||
method: str,
|
||||
path: str,
|
||||
params: dict[str, Any] | None = None,
|
||||
body: dict[str, Any] | None = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
) -> Any:
|
||||
if not path.startswith("/"):
|
||||
path = "/" + path
|
||||
if params:
|
||||
path = path + "?" + urlencode(params)
|
||||
|
||||
conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
|
||||
try:
|
||||
payload: bytes | None = None
|
||||
hdrs: dict[str, str] = {"Accept": "application/json"}
|
||||
|
||||
if body is not None:
|
||||
payload = json.dumps(body).encode("utf-8")
|
||||
hdrs["Content-Type"] = "application/json"
|
||||
if headers:
|
||||
hdrs.update(headers)
|
||||
|
||||
conn.request(method.upper(), path, body=payload, headers=hdrs)
|
||||
resp = conn.getresponse()
|
||||
raw = resp.read()
|
||||
text = raw.decode("utf-8", errors="replace") if raw else ""
|
||||
|
||||
if resp.status >= 400:
|
||||
raise ExoHttpError(resp.status, resp.reason, text[:300])
|
||||
|
||||
if not text:
|
||||
return None
|
||||
return json.loads(text)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
return self.request_json("POST", "/bench/chat/completions", body=payload)
|
||||
|
||||
def stream_bench_chat_completions(self, payload: dict[str, Any]) -> Iterator[str]:
|
||||
"""POST /bench/chat/completions with stream=True, yielding raw SSE lines."""
|
||||
payload = {**payload, "stream": True}
|
||||
data = json.dumps(payload).encode("utf-8")
|
||||
conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
|
||||
try:
|
||||
conn.request(
|
||||
"POST",
|
||||
"/bench/chat/completions",
|
||||
body=data,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "text/event-stream",
|
||||
},
|
||||
)
|
||||
resp = conn.getresponse()
|
||||
if resp.status >= 400:
|
||||
raw = resp.read().decode("utf-8", errors="replace")
|
||||
raise ExoHttpError(resp.status, resp.reason, raw[:300])
|
||||
for line in resp:
|
||||
yield line.decode("utf-8", errors="replace")
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def get_state_path(self, path: str) -> Any:
|
||||
try:
|
||||
return self.request_json("GET", f"/state/{path}")
|
||||
except ExoHttpError as e:
|
||||
if e.status == 404:
|
||||
return None
|
||||
raise
|
||||
|
||||
def get_instance(self, instance_id: str) -> dict[str, Any] | None:
|
||||
return self.get_state_path(f"instances/{instance_id}")
|
||||
|
||||
def get_runner(self, runner_id: str) -> dict[str, Any] | None:
|
||||
return self.get_state_path(f"runners/{runner_id}")
|
||||
|
||||
def get_node_downloads(self, node_id: str) -> list[dict[str, Any]] | None:
|
||||
return self.get_state_path(f"downloads/{node_id}")
|
||||
|
||||
def get_node_disk(self, node_id: str) -> dict[str, Any] | None:
|
||||
return self.get_state_path(f"nodeDisk/{node_id}")
|
||||
|
||||
def get_node_system(self, node_id: str) -> dict[str, Any] | None:
|
||||
return self.get_state_path(f"nodeSystem/{node_id}")
|
||||
|
||||
def get_node_identities(self) -> dict[str, Any] | None:
|
||||
return self.get_state_path("nodeIdentities")
|
||||
|
||||
def get_topology(self) -> dict[str, Any] | None:
|
||||
return self.get_state_path("topology")
|
||||
|
||||
|
||||
def unwrap_instance(instance: dict[str, Any]) -> dict[str, Any]:
|
||||
if len(instance) != 1:
|
||||
raise KeyError(f"Expected 1 key, got keys={list(instance.keys())}")
|
||||
|
||||
tag = next(iter(instance))
|
||||
inner = instance[tag]
|
||||
if not isinstance(inner, dict):
|
||||
raise TypeError(f"payload for {tag} must be dict, got {type(inner)}")
|
||||
return inner
|
||||
|
||||
|
||||
def instance_id_from_instance(instance: dict[str, Any]) -> str:
|
||||
inner = unwrap_instance(instance)
|
||||
return str(inner["instanceId"])
|
||||
|
||||
|
||||
def nodes_used_in_instance(instance: dict[str, Any]) -> int:
|
||||
inner = unwrap_instance(instance)
|
||||
return len(inner["shardAssignments"]["nodeToRunner"])
|
||||
|
||||
|
||||
def runner_ids_from_instance(instance: dict[str, Any]) -> list[str]:
|
||||
inner = unwrap_instance(instance)
|
||||
runner_to_shard = inner["shardAssignments"]["runnerToShard"]
|
||||
return list(runner_to_shard.keys())
|
||||
|
||||
|
||||
def node_ids_from_instance(instance: dict[str, Any]) -> list[str]:
|
||||
inner = unwrap_instance(instance)
|
||||
return list(inner["shardAssignments"]["nodeToRunner"].keys())
|
||||
|
||||
|
||||
def runner_ready(runner: dict[str, Any]) -> bool:
|
||||
return "RunnerReady" in runner
|
||||
|
||||
|
||||
def runner_failed(runner: dict[str, Any]) -> bool:
|
||||
return "RunnerFailed" in runner
|
||||
|
||||
|
||||
def get_runner_failed_message(runner: dict[str, Any]) -> str | None:
|
||||
if "RunnerFailed" in runner:
|
||||
return runner["RunnerFailed"].get("errorMessage")
|
||||
return None
|
||||
|
||||
|
||||
def wait_for_instance_ready(
|
||||
client: ExoClient, instance_id: str, timeout: float = 24000.0
|
||||
) -> None:
|
||||
start_time = time.time()
|
||||
instance_existed = False
|
||||
last_loaded: dict[str, int] = {}
|
||||
while time.time() - start_time < timeout:
|
||||
instance = client.get_instance(instance_id)
|
||||
|
||||
if instance is None:
|
||||
if instance_existed:
|
||||
raise RuntimeError(
|
||||
f"Instance {instance_id} was deleted (runner may have failed)"
|
||||
)
|
||||
time.sleep(0.1)
|
||||
continue
|
||||
|
||||
instance_existed = True
|
||||
rids = runner_ids_from_instance(instance)
|
||||
|
||||
all_ready = True
|
||||
for rid in rids:
|
||||
runner = client.get_runner(rid) or {}
|
||||
if runner_failed(runner):
|
||||
error_msg = get_runner_failed_message(runner) or "Unknown error"
|
||||
raise RuntimeError(f"Runner {rid} failed: {error_msg}")
|
||||
if "RunnerLoading" in runner:
|
||||
loading = runner["RunnerLoading"]
|
||||
loaded = loading.get("layersLoaded", 0)
|
||||
total = loading.get("totalLayers", 0)
|
||||
if total > 0 and last_loaded.get(rid) != loaded:
|
||||
last_loaded[rid] = loaded
|
||||
logger.debug(f"Runner {rid}: loading layers {loaded}/{total}")
|
||||
if not runner_ready(runner):
|
||||
all_ready = False
|
||||
|
||||
if all_ready:
|
||||
return
|
||||
|
||||
time.sleep(0.1)
|
||||
|
||||
raise TimeoutError(f"Instance {instance_id} did not become ready within {timeout=}")
|
||||
|
||||
|
||||
def wait_for_instance_gone(
|
||||
client: ExoClient, instance_id: str, timeout: float = 3.0
|
||||
) -> None:
|
||||
start_time = time.time()
|
||||
while time.time() - start_time < timeout:
|
||||
try:
|
||||
client.request_json("GET", f"/instance/{instance_id}")
|
||||
time.sleep(0.4)
|
||||
except ExoHttpError as e:
|
||||
if e.status == 404:
|
||||
return
|
||||
raise
|
||||
|
||||
raise TimeoutError(f"Instance {instance_id} did not get deleted within {timeout=}")
|
||||
|
||||
|
||||
def capture_cluster_snapshot(client: ExoClient) -> dict[str, Any]:
|
||||
snapshot: dict[str, Any] = {}
|
||||
identities = client.get_node_identities()
|
||||
if identities:
|
||||
snapshot["nodeIdentities"] = identities
|
||||
topology = client.get_topology()
|
||||
if topology:
|
||||
snapshot["topology"] = topology
|
||||
node_memory = client.get_state_path("nodeMemory")
|
||||
if node_memory:
|
||||
snapshot["nodeMemory"] = node_memory
|
||||
node_system = client.get_state_path("nodeSystem")
|
||||
if node_system:
|
||||
snapshot["nodeSystem"] = node_system
|
||||
return snapshot
|
||||
|
||||
|
||||
def resolve_model_short_id(
|
||||
client: ExoClient, model_arg: str, *, force_download: bool = False
|
||||
) -> tuple[str, str]:
|
||||
models = client.request_json("GET", "/models") or {}
|
||||
data = models.get("data") or []
|
||||
|
||||
for m in data:
|
||||
if (m.get("name") or "").lower() == model_arg.lower():
|
||||
short_id = str(m["name"])
|
||||
full_id = str(m.get("hugging_face_id") or m["name"])
|
||||
return short_id, full_id
|
||||
|
||||
for m in data:
|
||||
if m.get("hugging_face_id") == model_arg:
|
||||
short_id = str(m["name"])
|
||||
full_id = str(m["hugging_face_id"])
|
||||
return short_id, full_id
|
||||
|
||||
if force_download and "/" in model_arg:
|
||||
logger.info(f"Model not in /models, adding from HuggingFace: {model_arg}")
|
||||
result = client.request_json(
|
||||
"POST", "/models/add", body={"model_id": model_arg}
|
||||
)
|
||||
if result:
|
||||
short_id = str(result.get("name") or model_arg.rsplit("/", 1)[-1])
|
||||
full_id = str(result.get("hugging_face_id") or model_arg)
|
||||
return short_id, full_id
|
||||
|
||||
raise ValueError(f"Model not found in /models: {model_arg}")
|
||||
|
||||
|
||||
def placement_filter(instance_meta: str, wanted: str) -> bool:
|
||||
s = (instance_meta or "").lower()
|
||||
if wanted == "both":
|
||||
return ("ring" in s) or ("jaccl" in s)
|
||||
return wanted in s
|
||||
|
||||
|
||||
def sharding_filter(sharding: str, wanted: str) -> bool:
|
||||
s = (sharding or "").lower()
|
||||
if wanted == "both":
|
||||
return ("pipeline" in s) or ("tensor" in s)
|
||||
return wanted in s
|
||||
|
||||
|
||||
def fetch_and_filter_placements(
|
||||
client: ExoClient,
|
||||
full_model_id: str,
|
||||
args: argparse.Namespace,
|
||||
node_id: str | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
params: dict[str, str] = {"model_id": full_model_id}
|
||||
if node_id is not None:
|
||||
params["node_ids"] = node_id
|
||||
previews_resp = client.request_json("GET", "/instance/previews", params=params)
|
||||
previews = previews_resp.get("previews") or []
|
||||
|
||||
selected: list[dict[str, Any]] = []
|
||||
for p in previews:
|
||||
if p.get("error") is not None:
|
||||
continue
|
||||
if not placement_filter(str(p.get("instance_meta", "")), args.instance_meta):
|
||||
continue
|
||||
if not sharding_filter(str(p.get("sharding", "")), args.sharding):
|
||||
continue
|
||||
|
||||
instance = p.get("instance")
|
||||
if not isinstance(instance, dict):
|
||||
continue
|
||||
|
||||
n = nodes_used_in_instance(instance)
|
||||
# Skip tensor ring single node as it is pointless when pipeline ring
|
||||
if n == 1 and (
|
||||
(args.sharding == "both" and "tensor" in p.get("sharding", "").lower())
|
||||
or (
|
||||
args.instance_meta == "both"
|
||||
and "jaccl" in p.get("instance_meta", "").lower()
|
||||
)
|
||||
):
|
||||
continue
|
||||
|
||||
if (
|
||||
args.skip_pipeline_jaccl
|
||||
and (
|
||||
args.instance_meta == "both"
|
||||
and "jaccl" in p.get("instance_meta", "").lower()
|
||||
)
|
||||
and (
|
||||
args.sharding == "both" and "pipeline" in p.get("sharding", "").lower()
|
||||
)
|
||||
):
|
||||
continue
|
||||
|
||||
if (
|
||||
args.skip_tensor_ring
|
||||
and (
|
||||
args.instance_meta == "both"
|
||||
and "ring" in p.get("instance_meta", "").lower()
|
||||
)
|
||||
and (args.sharding == "both" and "tensor" in p.get("sharding", "").lower())
|
||||
):
|
||||
continue
|
||||
|
||||
if args.min_nodes <= n <= args.max_nodes:
|
||||
selected.append(p)
|
||||
|
||||
return selected
|
||||
|
||||
|
||||
def settle_and_fetch_placements(
|
||||
client: ExoClient,
|
||||
full_model_id: str,
|
||||
args: argparse.Namespace,
|
||||
settle_timeout: float = 0,
|
||||
node_id: str | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
selected = fetch_and_filter_placements(client, full_model_id, args, node_id=node_id)
|
||||
|
||||
if not selected and settle_timeout > 0:
|
||||
backoff = _SETTLE_INITIAL_BACKOFF_S
|
||||
deadline = time.monotonic() + settle_timeout
|
||||
while not selected and time.monotonic() < deadline:
|
||||
remaining = deadline - time.monotonic()
|
||||
logger.warning(
|
||||
f"No valid placements yet (cluster may still be settling). "
|
||||
f"Retrying in {backoff:.1f}s ({remaining:.0f}s remaining)..."
|
||||
)
|
||||
time.sleep(min(backoff, remaining))
|
||||
backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
|
||||
selected = fetch_and_filter_placements(
|
||||
client, full_model_id, args, node_id=node_id
|
||||
)
|
||||
|
||||
return selected
|
||||
|
||||
|
||||
def run_planning_phase(
|
||||
client: ExoClient,
|
||||
full_model_id: str,
|
||||
preview: dict[str, Any],
|
||||
danger_delete: bool,
|
||||
timeout: float,
|
||||
settle_deadline: float | None,
|
||||
) -> float | None:
|
||||
"""Check disk space and ensure model is downloaded before benchmarking.
|
||||
|
||||
Returns the wall-clock download duration in seconds if a fresh download
|
||||
was needed, or None if the model was already cached on all nodes.
|
||||
"""
|
||||
# Get model size from /models
|
||||
models = client.request_json("GET", "/models") or {}
|
||||
model_bytes = 0
|
||||
for m in models.get("data", []):
|
||||
if m.get("hugging_face_id") == full_model_id:
|
||||
model_bytes = m.get("storage_size_megabytes", 0) * 1024 * 1024
|
||||
break
|
||||
|
||||
if not model_bytes:
|
||||
logger.warning(
|
||||
f"Could not determine size for {full_model_id}, skipping disk check"
|
||||
)
|
||||
return None
|
||||
|
||||
# Get nodes from preview
|
||||
inner = unwrap_instance(preview["instance"])
|
||||
node_ids = list(inner["shardAssignments"]["nodeToRunner"].keys())
|
||||
runner_to_shard = inner["shardAssignments"]["runnerToShard"]
|
||||
|
||||
needs_download = False
|
||||
|
||||
for node_id in node_ids:
|
||||
node_downloads = client.get_node_downloads(node_id) or []
|
||||
|
||||
already_downloaded = any(
|
||||
"DownloadCompleted" in p
|
||||
and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
|
||||
"modelId"
|
||||
]
|
||||
== full_model_id
|
||||
for p in node_downloads
|
||||
)
|
||||
if already_downloaded:
|
||||
continue
|
||||
|
||||
needs_download = True
|
||||
|
||||
disk_info = client.get_node_disk(node_id) or {}
|
||||
backoff = _SETTLE_INITIAL_BACKOFF_S
|
||||
while not disk_info and settle_deadline and time.monotonic() < settle_deadline:
|
||||
remaining = settle_deadline - time.monotonic()
|
||||
logger.info(
|
||||
f"Waiting for disk info on {node_id} ({remaining:.0f}s remaining)..."
|
||||
)
|
||||
time.sleep(min(backoff, remaining))
|
||||
backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
|
||||
disk_info = client.get_node_disk(node_id) or {}
|
||||
|
||||
if not disk_info:
|
||||
logger.warning(f"No disk info for {node_id}, skipping space check")
|
||||
continue
|
||||
|
||||
avail = disk_info.get("available", {}).get("inBytes", 0)
|
||||
if avail >= model_bytes:
|
||||
continue
|
||||
|
||||
if not danger_delete:
|
||||
raise RuntimeError(
|
||||
f"Insufficient disk on {node_id}: need {model_bytes // (1024**3)}GB, "
|
||||
f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space."
|
||||
)
|
||||
|
||||
completed = [
|
||||
(
|
||||
unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
|
||||
"modelId"
|
||||
],
|
||||
p["DownloadCompleted"]["total"]["inBytes"],
|
||||
)
|
||||
for p in node_downloads
|
||||
if "DownloadCompleted" in p
|
||||
and not p["DownloadCompleted"].get("readOnly", False)
|
||||
]
|
||||
for del_model, size in sorted(completed, key=lambda x: x[1]):
|
||||
logger.info(f"Deleting {del_model} from {node_id} ({size // (1024**2)}MB)")
|
||||
client.request_json("DELETE", f"/download/{node_id}/{del_model}")
|
||||
avail += size
|
||||
if avail >= model_bytes:
|
||||
break
|
||||
|
||||
if avail < model_bytes:
|
||||
raise RuntimeError(f"Could not free enough space on {node_id}")
|
||||
|
||||
# Start downloads (idempotent)
|
||||
download_t0 = time.perf_counter() if needs_download else None
|
||||
for node_id in node_ids:
|
||||
runner_id = inner["shardAssignments"]["nodeToRunner"][node_id]
|
||||
shard = runner_to_shard[runner_id]
|
||||
client.request_json(
|
||||
"POST",
|
||||
"/download/start",
|
||||
body={
|
||||
"targetNodeId": node_id,
|
||||
"shardMetadata": shard,
|
||||
},
|
||||
)
|
||||
logger.info(f"Started download on {node_id}")
|
||||
|
||||
# Wait for downloads (no timeout — poll until complete or failed)
|
||||
while True:
|
||||
all_done = True
|
||||
for node_id in node_ids:
|
||||
node_downloads = client.get_node_downloads(node_id) or []
|
||||
done = any(
|
||||
"DownloadCompleted" in p
|
||||
and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])[
|
||||
"modelCard"
|
||||
]["modelId"]
|
||||
== full_model_id
|
||||
for p in node_downloads
|
||||
)
|
||||
failed = [
|
||||
p["DownloadFailed"]["errorMessage"]
|
||||
for p in node_downloads
|
||||
if "DownloadFailed" in p
|
||||
and unwrap_instance(p["DownloadFailed"]["shardMetadata"])["modelCard"][
|
||||
"modelId"
|
||||
]
|
||||
== full_model_id
|
||||
]
|
||||
if failed:
|
||||
raise RuntimeError(f"Download failed on {node_id}: {failed[0]}")
|
||||
if not done:
|
||||
all_done = False
|
||||
ongoing = [
|
||||
p
|
||||
for p in node_downloads
|
||||
if "DownloadOngoing" in p
|
||||
and unwrap_instance(p["DownloadOngoing"]["shardMetadata"])[
|
||||
"modelCard"
|
||||
]["modelId"]
|
||||
== full_model_id
|
||||
]
|
||||
if ongoing:
|
||||
prog = ongoing[0]["DownloadOngoing"]["downloadProgress"]
|
||||
speed_mb = prog.get("speed", 0) / (1024 * 1024)
|
||||
eta_s = prog.get("etaMs", 0) / 1000
|
||||
dl_bytes = prog.get("downloaded", {}).get("inBytes", 0)
|
||||
total_bytes = prog.get("total", {}).get("inBytes", 0)
|
||||
pct = (dl_bytes / total_bytes * 100) if total_bytes else 0
|
||||
logger.info(
|
||||
f"Downloading on {node_id}: {pct:.1f}% @ {speed_mb:.1f} MB/s, "
|
||||
f"ETA {eta_s:.0f}s "
|
||||
f"({prog.get('completedFiles', 0)}/{prog.get('totalFiles', 0)} files)"
|
||||
)
|
||||
if all_done:
|
||||
if download_t0 is not None:
|
||||
return time.perf_counter() - download_t0
|
||||
return None
|
||||
time.sleep(10)
|
||||
|
||||
|
||||
def find_existing_instance(client: ExoClient, model_id: str) -> str | None:
|
||||
"""Find an existing running instance for the given model."""
|
||||
try:
|
||||
state = client.request_json("GET", "/state")
|
||||
except Exception:
|
||||
return None
|
||||
for inst_id, inst in state.get("instances", {}).items():
|
||||
# Instance structure is nested: {"MlxJacclInstance": {"shardAssignments": {"modelId": ...}}}
|
||||
for _inst_type, inner in inst.items():
|
||||
if not isinstance(inner, dict):
|
||||
continue
|
||||
sa = inner.get("shardAssignments", {})
|
||||
if sa.get("modelId") == model_id:
|
||||
return inst_id
|
||||
return None
|
||||
|
||||
|
||||
def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
|
||||
ap.add_argument("--host", default=os.environ.get("EXO_HOST", "localhost"))
|
||||
ap.add_argument(
|
||||
"--port", type=int, default=int(os.environ.get("EXO_PORT", "52415"))
|
||||
)
|
||||
ap.add_argument("--model", required=True, help="Model short id or huggingface id")
|
||||
ap.add_argument(
|
||||
"--force-download",
|
||||
action="store_true",
|
||||
help="If model not in /models, add it from HuggingFace via exo and download.",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--max-nodes",
|
||||
type=int,
|
||||
default=4,
|
||||
help="Only consider placements using <= this many nodes.",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--min-nodes",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Only consider placements using >= this many nodes.",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--instance-meta", choices=["ring", "jaccl", "both"], default="both"
|
||||
)
|
||||
ap.add_argument(
|
||||
"--sharding", choices=["pipeline", "tensor", "both"], default="both"
|
||||
)
|
||||
ap.add_argument(
|
||||
"--skip-pipeline-jaccl",
|
||||
action="store_true",
|
||||
help="Skip pipeline+jaccl placements, as it's often pointless.",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--skip-tensor-ring",
|
||||
action="store_true",
|
||||
help="Skip tensor+ring placements, as it's so slow.",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--timeout", type=float, default=7200.0, help="HTTP timeout (seconds)."
|
||||
)
|
||||
ap.add_argument(
|
||||
"--settle-timeout",
|
||||
type=float,
|
||||
default=60.0,
|
||||
help="Max seconds to wait for the cluster to produce valid placements (0 = try once).",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--danger-delete-downloads",
|
||||
action="store_true",
|
||||
help="Delete existing models from smallest to largest to make room for benchmark model.",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--reuse-instance",
|
||||
action="store_true",
|
||||
help="Reuse an existing running instance for this model instead of creating a new one.",
|
||||
)
|
||||
Loaded 100 of 474 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user