Merge main

This commit is contained in:
Ryuichi Leo Takashige
2026-03-16 23:01:48 +00:00
5 changed files with 51 additions and 43 deletions
@@ -82,6 +82,12 @@
d="M12.025 1.13c-5.77 0-10.449 4.647-10.449 10.378 0 1.112.178 2.181.503 3.185.064-.222.203-.444.416-.577a.96.96 0 0 1 .524-.15c.293 0 .584.124.84.284.278.173.48.408.71.694.226.282.458.611.684.951v-.014c.017-.324.106-.622.264-.874s.403-.487.762-.543c.3-.047.596.06.787.203s.31.313.4.467c.15.257.212.468.233.542.01.026.653 1.552 1.657 2.54.616.605 1.01 1.223 1.082 1.912.055.537-.096 1.059-.38 1.572.637.121 1.294.187 1.967.187.657 0 1.298-.063 1.921-.178-.287-.517-.44-1.041-.384-1.581.07-.69.465-1.307 1.081-1.913 1.004-.987 1.647-2.513 1.657-2.539.021-.074.083-.285.233-.542.09-.154.208-.323.4-.467a1.08 1.08 0 0 1 .787-.203c.359.056.604.29.762.543s.247.55.265.874v.015c.225-.34.457-.67.683-.952.23-.286.432-.52.71-.694.257-.16.547-.284.84-.285a.97.97 0 0 1 .524.151c.228.143.373.388.43.625l.006.04a10.3 10.3 0 0 0 .534-3.273c0-5.731-4.678-10.378-10.449-10.378M8.327 6.583a1.5 1.5 0 0 1 .713.174 1.487 1.487 0 0 1 .617 2.013c-.183.343-.762-.214-1.102-.094-.38.134-.532.914-.917.71a1.487 1.487 0 0 1 .69-2.803m7.486 0a1.487 1.487 0 0 1 .689 2.803c-.385.204-.536-.576-.916-.71-.34-.12-.92.437-1.103.094a1.487 1.487 0 0 1 .617-2.013 1.5 1.5 0 0 1 .713-.174m-10.68 1.55a.96.96 0 1 1 0 1.921.96.96 0 0 1 0-1.92m13.838 0a.96.96 0 1 1 0 1.92.96.96 0 0 1 0-1.92M8.489 11.458c.588.01 1.965 1.157 3.572 1.164 1.607-.007 2.984-1.155 3.572-1.164.196-.003.305.12.305.454 0 .886-.424 2.328-1.563 3.202-.22-.756-1.396-1.366-1.63-1.32q-.011.001-.02.006l-.044.026-.01.008-.03.024q-.018.017-.035.036l-.032.04a1 1 0 0 0-.058.09l-.014.025q-.049.088-.11.19a1 1 0 0 1-.083.116 1.2 1.2 0 0 1-.173.18q-.035.029-.075.058a1.3 1.3 0 0 1-.251-.243 1 1 0 0 1-.076-.107c-.124-.193-.177-.363-.337-.444-.034-.016-.104-.008-.2.022q-.094.03-.216.087-.06.028-.125.063l-.13.074q-.067.04-.136.086a3 3 0 0 0-.135.096 3 3 0 0 0-.26.219 2 2 0 0 0-.12.121 2 2 0 0 0-.106.128l-.002.002a2 2 0 0 0-.09.132l-.001.001a1.2 1.2 0 0 0-.105.212q-.013.036-.024.073c-1.139-.875-1.563-2.317-1.563-3.203 0-.334.109-.457.305-.454m.836 10.354c.824-1.19.766-2.082-.365-3.194-1.13-1.112-1.789-2.738-1.789-2.738s-.246-.945-.806-.858-.97 1.499.202 2.362c1.173.864-.233 1.45-.685.64-.45-.812-1.683-2.896-2.322-3.295s-1.089-.175-.938.647 2.822 2.813 2.562 3.244-1.176-.506-1.176-.506-2.866-2.567-3.49-1.898.473 1.23 2.037 2.16c1.564.932 1.686 1.178 1.464 1.53s-3.675-2.511-4-1.297c-.323 1.214 3.524 1.567 3.287 2.405-.238.839-2.71-1.587-3.216-.642-.506.946 3.49 2.056 3.522 2.064 1.29.33 4.568 1.028 5.713-.624m5.349 0c-.824-1.19-.766-2.082.365-3.194 1.13-1.112 1.789-2.738 1.789-2.738s.246-.945.806-.858.97 1.499-.202 2.362c-1.173.864.233 1.45.685.64.451-.812 1.683-2.896 2.322-3.295s1.089-.175.938.647-2.822 2.813-2.562 3.244 1.176-.506 1.176-.506 2.866-2.567 3.49-1.898-.473 1.23-2.037 2.16c-1.564.932-1.686 1.178-1.464 1.53s3.675-2.511 4-1.297c.323 1.214-3.524 1.567-3.287 2.405.238.839 2.71-1.587 3.216-.642.506.946-3.49 2.056-3.522 2.064-1.29.33-4.568 1.028-5.713-.624"
/>
</svg>
{:else if family === "step"}
<svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
<path
d="M22.012 0h1.032v.927H24v.968h-.956V3.78h-1.032V1.896h-1.878v-.97h1.878V0zM2.6 12.371V1.87h.969v10.502h-.97zm10.423.66h10.95v.918h-6.208v9.579h-4.742V13.03zM5.629 3.333v12.356H0v4.51h10.386V8L20.859 8l-.003-4.668-15.227.001z"
/>
</svg>
{:else}
<svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
<path
+16 -17
View File
@@ -114,23 +114,22 @@
};
};
packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin
(
let
uvLock = builtins.fromTOML (builtins.readFile ./uv.lock);
mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx" && p.source ? git) uvLock.package);
uvLockMlxVersion = mlxPackage.version;
in
{
metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
mlx = pkgs.callPackage ./nix/mlx.nix {
inherit (self'.packages) metal-toolchain;
inherit uvLockMlxVersion;
};
default = self'.packages.exo;
}
) // lib.optionalAttrs (pkgsCuda != null) {
# CUDA-compiled PyTorch and vLLM (built from source by nixpkgs, aarch64-linux only)
packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin (
let
uvLock = builtins.fromTOML (builtins.readFile ./uv.lock);
mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx" && p.source ? git) uvLock.package);
uvLockMlxVersion = mlxPackage.version;
uvLockMlxRev = builtins.elemAt (builtins.split "#" mlxPackage.source.git) 2;
in
{
metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
mlx = pkgs.callPackage ./nix/mlx.nix {
inherit (self'.packages) metal-toolchain;
inherit uvLockMlxVersion uvLockMlxRev;
};
default = self'.packages.exo;
}
) // lib.optionalAttrs (pkgsCuda != null) {
torch-cuda = pkgsCuda.python313Packages.torch;
vllm-cuda = pkgsCuda.python313Packages.vllm;
+3 -4
View File
@@ -11,6 +11,7 @@
, fmt
, python313Packages
, uvLockMlxVersion
, uvLockMlxRev
}:
assert stdenv.isDarwin;
@@ -41,15 +42,13 @@ let
mlx = stdenv.mkDerivation rec {
pname = "mlx";
version = let v = "0.30.7.dev20260225+257d5692"; in
assert v == uvLockMlxVersion || throw "MLX version mismatch: nix/mlx.nix has ${v} but uv.lock has ${uvLockMlxVersion}. Update both the version and hash in nix/mlx.nix.";
v;
version = uvLockMlxVersion;
pyproject = true;
src = fetchFromGitHub {
owner = "rltakashige";
repo = "mlx-jaccl-fix-small-recv";
rev = "257d5692fc7af6bba3b8afaeb63c549b7d1e43d5";
rev = uvLockMlxRev;
hash = "sha256-GosFIWxIB48Egb1MqJrR3xhsUsQeWdRk5rV93USY6wQ=";
};
+15 -19
View File
@@ -1,4 +1,3 @@
import asyncio
from dataclasses import dataclass, field
import anyio
@@ -47,7 +46,7 @@ class DownloadCoordinator:
# Local state
download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict)
active_downloads: dict[ModelId, asyncio.Task[None]] = field(default_factory=dict)
active_downloads: dict[ModelId, anyio.CancelScope] = field(default_factory=dict)
_tg: TaskGroup = field(init=False, default_factory=TaskGroup)
@@ -77,8 +76,6 @@ class DownloadCoordinator:
await self.event_sender.send(
NodeDownloadProgress(download_progress=completed)
)
if model_id in self.active_downloads:
del self.active_downloads[model_id]
self._last_progress_time.pop(model_id, None)
elif (
progress.status == "in_progress"
@@ -103,13 +100,9 @@ class DownloadCoordinator:
logger.info(
f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}"
)
try:
async with self._tg as tg:
tg.start_soon(self._command_processor)
tg.start_soon(self._emit_existing_download_progress)
finally:
for task in self.active_downloads.values():
task.cancel()
async with self._tg as tg:
tg.start_soon(self._command_processor)
tg.start_soon(self._emit_existing_download_progress)
def shutdown(self) -> None:
self._tg.cancel_tasks()
@@ -132,7 +125,7 @@ class DownloadCoordinator:
async def _cancel_download(self, model_id: ModelId) -> None:
if model_id in self.active_downloads and model_id in self.download_status:
logger.info(f"Cancelling download for {model_id}")
self.active_downloads.pop(model_id).cancel()
self.active_downloads[model_id].cancel()
current_status = self.download_status[model_id]
pending = DownloadPending(
shard_metadata=current_status.shard_metadata,
@@ -236,9 +229,10 @@ class DownloadCoordinator:
self.download_status[model_id] = status
self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))
async def download_wrapper() -> None:
async def download_wrapper(cancel_scope: anyio.CancelScope) -> None:
try:
await self.shard_downloader.ensure_shard(shard)
with cancel_scope:
await self.shard_downloader.ensure_shard(shard)
except Exception as e:
logger.error(f"Download failed for {model_id}: {e}")
failed = DownloadFailed(
@@ -251,12 +245,15 @@ class DownloadCoordinator:
await self.event_sender.send(
NodeDownloadProgress(download_progress=failed)
)
except anyio.get_cancelled_exc_class():
# ignore cancellation - let cleanup do its thing
pass
finally:
if model_id in self.active_downloads:
del self.active_downloads[model_id]
self.active_downloads.pop(model_id, None)
task = asyncio.create_task(download_wrapper())
self.active_downloads[model_id] = task
scope = anyio.CancelScope()
self._tg.start_soon(download_wrapper, scope)
self.active_downloads[model_id] = scope
async def _delete_download(self, model_id: ModelId) -> None:
# Protect read-only models (from EXO_MODELS_PATH) from deletion
@@ -272,7 +269,6 @@ class DownloadCoordinator:
if model_id in self.active_downloads:
logger.info(f"Cancelling active download for {model_id} before deletion")
self.active_downloads[model_id].cancel()
del self.active_downloads[model_id]
# Delete from disk
logger.info(f"Deleting model files for {model_id}")
+11 -3
View File
@@ -529,6 +529,9 @@ class InfoGatherer:
if self.macmon_interval is None:
return
# macmon pipe --interval [interval in ms]
# Timeout: if macmon produces no output for this many seconds, restart it.
# macmon writes every macmon_interval seconds, so 10x that is generous.
read_timeout = max(self.macmon_interval * 10, 30)
while True:
try:
async with await open_process(
@@ -542,10 +545,15 @@ class InfoGatherer:
if not p.stdout:
logger.critical("MacMon closed stdout")
return
async for text in TextReceiveStream(
BufferedByteReceiveStream(p.stdout)
):
stream = TextReceiveStream(BufferedByteReceiveStream(p.stdout))
while True:
with fail_after(read_timeout):
text = await stream.receive()
await self.info_sender.send(MacmonMetrics.from_raw_json(text))
except TimeoutError:
logger.warning(
f"MacMon produced no output for {read_timeout}s, restarting"
)
except CalledProcessError as e:
stderr_msg = "no stderr"
stderr_output = cast(bytes | str | None, e.stderr)