github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add Mac OS support for Apple Silicon (#770)

Browse Source 
* Made an install script and auto updates env for mac

* GPU sensors and initial training working for MAC. Still WIP.

* Switch dataloader to single threaded until I can work around some mac pickeling issues.

* Get quantization working on mac

* Fix mac exclusive imports so they don't break other builds.

* Add mac instructions to the UI
This commit is contained in:
Jaret Burkett
2026-03-30 09:37:47 -06:00
committed by GitHub
parent bc47fd6755
commit 171535833a
25 changed files with 1072 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
ui/src/app/api/cpu/route.ts
View File

@@ -1,20 +1,57 @@
import { NextResponse } from 'next/server';
import si from 'systeminformation';
import { createRequire } from 'module';
import os from 'os';
import { CpuInfo } from '@/types';
const isMac = os.platform() === 'darwin';
export async function GET() {
try {
const cpuInfoRaw = await si.cpu();
const memoryData = await si.mem();
let cpuInfo: CpuInfo = {
name: `${cpuInfoRaw.manufacturer} ${cpuInfoRaw.brand}`,
cores: cpuInfoRaw.cores,
temperature: (await si.cpuTemperature()).main || 0,
totalMemory: memoryData.total / (1024 * 1024),
availableMemory: memoryData.available / (1024 * 1024),
freeMemory: memoryData.free / (1024 * 1024),
currentLoad: (await si.currentLoad()).currentLoad || 0,
};
let cpuInfo: CpuInfo;
if (isMac) {
try {
const nativeRequire = createRequire(import.meta.url);
const ms = nativeRequire('macstats') as any;
const ramData = ms.getRAMUsageSync();
const cpuData = ms.getCpuDataSync();
cpuInfo = {
name: `${cpuInfoRaw.manufacturer} ${cpuInfoRaw.brand}`,
cores: cpuInfoRaw.cores,
temperature: cpuData.temperature || 0,
totalMemory: ramData.total / (1024 * 1024),
availableMemory: ramData.free / (1024 * 1024),
freeMemory: ramData.free / (1024 * 1024),
currentLoad: (await si.currentLoad()).currentLoad || 0,
};
} catch {
// Fallback to systeminformation if macstats fails
const memoryData = await si.mem();
cpuInfo = {
name: `${cpuInfoRaw.manufacturer} ${cpuInfoRaw.brand}`,
cores: cpuInfoRaw.cores,
temperature: (await si.cpuTemperature()).main || 0,
totalMemory: memoryData.total / (1024 * 1024),
availableMemory: memoryData.available / (1024 * 1024),
freeMemory: memoryData.free / (1024 * 1024),
currentLoad: (await si.currentLoad()).currentLoad || 0,
};
}
} else {
const memoryData = await si.mem();
cpuInfo = {
name: `${cpuInfoRaw.manufacturer} ${cpuInfoRaw.brand}`,
cores: cpuInfoRaw.cores,
temperature: (await si.cpuTemperature()).main || 0,
totalMemory: memoryData.total / (1024 * 1024),
availableMemory: memoryData.available / (1024 * 1024),
freeMemory: memoryData.free / (1024 * 1024),
currentLoad: (await si.currentLoad()).currentLoad || 0,
};
}
return NextResponse.json(cpuInfo);
} catch (error) {

130
ui/src/app/api/gpu/route.ts
View File

@@ -1,15 +1,140 @@
import { NextResponse } from 'next/server';
import { exec } from 'child_process';
import { exec, execSync } from 'child_process';
import { promisify } from 'util';
import { createRequire } from 'module';
import os from 'os';
const execAsync = promisify(exec);
interface MacGpuResult {
name: string;
memUsed: number;
memTotal: number;
gpuLoad: number;
temperature: number;
fanSpeed: number;
powerDraw: number;
}
async function getMacGpuInfo(): Promise<MacGpuResult | null> {
try {
const memoryTotal = os.totalmem() / (1024 * 1024);
// Get GPU name and core count from system_profiler
let gpuName = 'Apple GPU';
try {
const spOut = execSync(
'system_profiler SPDisplaysDataType 2>/dev/null | grep -E "Chipset Model|Total Number of Cores"',
{ encoding: 'utf-8', timeout: 5000 },
);
const nameMatch = spOut.match(/Chipset Model:\s*(.+)/);
const coresMatch = spOut.match(/Total Number of Cores:\s*(\d+)/);
if (nameMatch) {
gpuName = nameMatch[1].trim();
if (coresMatch) {
gpuName += ` GPU (${coresMatch[1]} cores)`;
}
}
} catch {
// fallback to generic name
}
let temperature = 0;
let gpuLoad = 0;
let fanSpeed = 0;
let powerDraw = 0;
let memUsed = 0;
let memTotal = memoryTotal;
try {
// Use createRequire to hide from webpack static analysis so it doesn't fail on non-mac platforms
const nativeRequire = createRequire(import.meta.url);
const ms = nativeRequire('macstats') as any;
try {
const gpuData = ms.getGpuDataSync();
temperature = gpuData.temperature || 0;
gpuLoad = gpuData.usage || 0;
} catch {
// ignore
}
try {
const fanData = ms.getFanDataSync();
const fanKeys = Object.keys(fanData);
if (fanKeys.length > 0) {
fanSpeed = fanData[fanKeys[0]].rpm || 0;
}
} catch {
// ignore
}
try {
const powerData = ms.getPowerDataSync();
powerDraw = powerData.gpu || 0;
} catch {
// ignore
}
try {
const ramData = ms.getRAMUsageSync();
memUsed = ramData.used / (1024 * 1024);
memTotal = ramData.total / (1024 * 1024);
} catch {
// ignore
}
} catch (error) {
console.warn('macstats not available:', error);
}
return { name: gpuName, memUsed, memTotal, gpuLoad, temperature, fanSpeed, powerDraw };
} catch {
return null;
}
}
export async function GET() {
try {
// Get platform
const platform = os.platform();
const isWindows = platform === 'win32';
const isMac = platform === 'darwin';
if (isMac) {
const macGpu = await getMacGpuInfo();
if (macGpu) {
return NextResponse.json({
hasNvidiaSmi: false,
isMac: true,
gpus: [
{
index: 0,
name: macGpu.name,
driverVersion: 'macOS',
temperature: Math.round(macGpu.temperature),
utilization: {
gpu: macGpu.gpuLoad,
memory: macGpu.memTotal > 0 ? Math.round((macGpu.memUsed / macGpu.memTotal) * 100) : 0,
},
memory: {
total: Math.round(macGpu.memTotal),
free: Math.round(macGpu.memTotal - macGpu.memUsed),
used: Math.round(macGpu.memUsed),
},
power: { draw: macGpu.powerDraw, limit: 0 },
clocks: { graphics: 0, memory: 0 },
fan: { speed: macGpu.fanSpeed },
},
],
});
}
return NextResponse.json({
hasNvidiaSmi: false,
isMac: true,
gpus: [],
error: 'Could not read Mac GPU stats',
});
}
// Check if nvidia-smi is available
const hasNvidiaSmi = await checkNvidiaSmi(isWindows);
@@ -17,6 +142,7 @@ export async function GET() {
if (!hasNvidiaSmi) {
return NextResponse.json({
hasNvidiaSmi: false,
isMac: false,
gpus: [],
error: 'nvidia-smi not found or not accessible',
});
@@ -34,6 +160,7 @@ export async function GET() {
return NextResponse.json(
{
hasNvidiaSmi: false,
isMac: false,
gpus: [],
error: `Failed to fetch GPU stats: ${error instanceof Error ? error.message : String(error)}`,
},
@@ -121,3 +248,4 @@ async function getGpuStats(isWindows: boolean) {
return gpus;
}

8
ui/src/app/api/jobs/route.ts
View File

@@ -1,5 +1,6 @@
import { NextResponse } from 'next/server';
import { PrismaClient } from '@prisma/client';
import { isMac } from '@/helpers/basic';
const prisma = new PrismaClient();
@@ -28,7 +29,12 @@ export async function GET(request: Request) {
export async function POST(request: Request) {
try {
const body = await request.json();
const { id, name, job_config, gpu_ids } = body;
const { id, name, job_config } = body;
let gpu_ids: string = body.gpu_ids;
if (isMac()) {
gpu_ids = "mps";
}
if (id) {
// Update existing training

21
ui/src/app/jobs/new/SimpleJob.tsx
View File

@@ -19,6 +19,7 @@ import SampleControlImage from '@/components/SampleControlImage';
import { FlipHorizontal2, FlipVertical2 } from 'lucide-react';
import { handleModelArchChange } from './utils';
import { IoFlaskSharp } from 'react-icons/io5';
import { isMac } from '@/helpers/basic';
type Props = {
jobConfig: JobConfig;
@@ -146,6 +147,8 @@ export default function SimpleJob({
return newQuantizationOptions;
}, [modelArch]);
const showGPUSelect = !isMac();
return (
<>
<form
@@ -171,13 +174,15 @@ export default function SimpleJob({
disabled={runId !== null}
required
/>
<SelectInput
label="GPU ID"
value={`${gpuIDs}`}
docKey="gpuids"
onChange={value => setGpuIDs(value)}
options={gpuList.map((gpu: any) => ({ value: `${gpu.index}`, label: `GPU #${gpu.index}` }))}
/>
{showGPUSelect && (
<SelectInput
label="GPU ID"
value={`${gpuIDs}`}
docKey="gpuids"
onChange={value => setGpuIDs(value)}
options={gpuList.map((gpu: any) => ({ value: `${gpu.index}`, label: `GPU #${gpu.index}` }))}
/>
)}
{disableSections.includes('trigger_word') ? null : (
<TextInput
label="Trigger Word"
@@ -249,7 +254,7 @@ export default function SimpleJob({
onChange={value => setJobConfig(value, 'config.process[0].model.model_kwargs.match_target_res')}
/>
)}
{modelArch?.additionalSections?.includes('model.layer_offloading') && (
{modelArch?.additionalSections?.includes('model.layer_offloading') && !isMac() && (
<>
<Checkbox
label={

6
ui/src/app/jobs/new/jobConfig.ts
View File

@@ -1,3 +1,5 @@
'use client';
import { isMac } from '@/helpers/basic';
import { JobConfig, DatasetConfig, SliderConfig } from '@/types';
export const defaultDatasetConfig: DatasetConfig = {
@@ -199,5 +201,9 @@ export const migrateJobConfig = (jobConfig: JobConfig): JobConfig => {
use_ui_logger: true,
};
}
if (isMac()) {
jobConfig.config.process[0].device = 'mps';
}
return jobConfig;
};

2
ui/src/app/jobs/new/page.tsx
View File

@@ -35,7 +35,7 @@ export default function TrainingForm() {
const [datasetOptions, setDatasetOptions] = useState<{ value: string; label: string }[]>([]);
const [showAdvancedView, setShowAdvancedView] = useState(false);
const [jobConfig, setJobConfig] = useNestedState<JobConfig>(objectCopy(defaultJobConfig));
const [jobConfig, setJobConfig] = useNestedState<JobConfig>(objectCopy(migrateJobConfig(defaultJobConfig)));
const [status, setStatus] = useState<'idle' | 'saving' | 'success' | 'error'>('idle');
const fileInputRef = useRef<HTMLInputElement>(null);

4
ui/src/app/layout.tsx
View File

@@ -7,6 +7,7 @@ import ConfirmModal from '@/components/ConfirmModal';
import { Suspense } from 'react';
import AuthWrapper from '@/components/AuthWrapper';
import DocModal from '@/components/DocModal';
import os from 'os';
export const dynamic = 'force-dynamic';
@@ -21,12 +22,15 @@ export default function RootLayout({ children }: { children: React.ReactNode })
// Check if the AI_TOOLKIT_AUTH environment variable is set
const authRequired = process.env.AI_TOOLKIT_AUTH ? true : false;
const platform = os.platform();
return (
<html lang="en" className="dark">
<head>
<meta name="apple-mobile-web-app-title" content="AI-Toolkit" />
</head>
<body className={inter.className}>
<script dangerouslySetInnerHTML={{ __html: `window.server_platform = "${platform}";` }} />
<ThemeProvider>
<AuthWrapper authRequired={authRequired}>
<div className="flex h-screen bg-gray-950">

2
ui/src/components/GPUMonitor.tsx
View File

@@ -98,7 +98,7 @@ const GpuMonitor: React.FC = () => {
);
}
if (!gpuData.hasNvidiaSmi) {
if (!gpuData.hasNvidiaSmi && !gpuData.isMac) {
return (
<div className="bg-yellow-900 border border-yellow-700 text-yellow-300 px-4 py-3 rounded relative" role="alert">
<strong className="font-bold">No NVIDIA GPUs detected!</strong>

9
ui/src/components/JobOverview.tsx
View File

@@ -14,12 +14,17 @@ interface JobOverviewProps {
}
export default function JobOverview({ job }: JobOverviewProps) {
const gpuIds = useMemo(() => job.gpu_ids.split(',').map(id => parseInt(id)), [job.gpu_ids]);
const gpuIds = useMemo(() => {
if (job.gpu_ids === 'mps') {
return [0]; // For MPS, we can just return a single GPU ID since it's virtualized
}
return job.gpu_ids.split(',').map(id => parseInt(id));
}, [job.gpu_ids]);
const { log, setLog, status: statusLog, refresh: refreshLog } = useJobLog(job.id, 2000);
const logRef = useRef<HTMLDivElement>(null);
// Track whether we should auto-scroll to bottom
const [isScrolledToBottom, setIsScrolledToBottom] = useState(true);
console.log('job.gpu_ids', job.gpu_ids);
const { gpuList, isGPUInfoLoaded } = useGPUInfo(gpuIds, 5000);
const { cpuInfo, isCPUInfoLoaded } = useCPUInfo(5000);
const totalSteps = getTotalSteps(job);

1
ui/src/helpers/basic.ts Normal file
View File

@@ -0,0 +1 @@
export const isMac = () => (typeof window !== 'undefined' && (window as any).server_platform === 'darwin') || process.platform === 'darwin';

1
ui/src/types.ts
View File

@@ -51,6 +51,7 @@ export interface CpuInfo {
export interface GPUApiResponse {
hasNvidiaSmi: boolean;
isMac: boolean;
gpus: GpuInfo[];
error?: string;
}