feat: add gpu options

This commit is contained in:
BennyKok 2024-01-11 21:57:30 +08:00
parent 6e3b7f0435
commit 29bf7636b1
10 changed files with 792 additions and 5 deletions

View File

@ -0,0 +1,7 @@
DO $$ BEGIN
CREATE TYPE "machine_gpu" AS ENUM('T4', 'A10G', 'A100');
EXCEPTION
WHEN duplicate_object THEN null;
END $$;
--> statement-breakpoint
ALTER TABLE "comfyui_deploy"."machines" ADD COLUMN "gpu" "machine_gpu";

View File

@ -0,0 +1,736 @@
{
"id": "9d175095-5a2f-45b8-a38f-2a5f01f93a31",
"prevId": "56dcfeff-7e8c-46de-9a84-319444df65d8",
"version": "5",
"dialect": "pg",
"tables": {
"api_keys": {
"name": "api_keys",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"key": {
"name": "key",
"type": "text",
"primaryKey": false,
"notNull": true
},
"name": {
"name": "name",
"type": "text",
"primaryKey": false,
"notNull": true
},
"user_id": {
"name": "user_id",
"type": "text",
"primaryKey": false,
"notNull": true
},
"org_id": {
"name": "org_id",
"type": "text",
"primaryKey": false,
"notNull": false
},
"revoked": {
"name": "revoked",
"type": "boolean",
"primaryKey": false,
"notNull": true,
"default": false
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"api_keys_user_id_users_id_fk": {
"name": "api_keys_user_id_users_id_fk",
"tableFrom": "api_keys",
"tableTo": "users",
"columnsFrom": [
"user_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {
"api_keys_key_unique": {
"name": "api_keys_key_unique",
"nullsNotDistinct": false,
"columns": [
"key"
]
}
}
},
"deployments": {
"name": "deployments",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"user_id": {
"name": "user_id",
"type": "text",
"primaryKey": false,
"notNull": true
},
"workflow_version_id": {
"name": "workflow_version_id",
"type": "uuid",
"primaryKey": false,
"notNull": true
},
"workflow_id": {
"name": "workflow_id",
"type": "uuid",
"primaryKey": false,
"notNull": true
},
"machine_id": {
"name": "machine_id",
"type": "uuid",
"primaryKey": false,
"notNull": true
},
"environment": {
"name": "environment",
"type": "deployment_environment",
"primaryKey": false,
"notNull": true
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"deployments_user_id_users_id_fk": {
"name": "deployments_user_id_users_id_fk",
"tableFrom": "deployments",
"tableTo": "users",
"columnsFrom": [
"user_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
},
"deployments_workflow_version_id_workflow_versions_id_fk": {
"name": "deployments_workflow_version_id_workflow_versions_id_fk",
"tableFrom": "deployments",
"tableTo": "workflow_versions",
"columnsFrom": [
"workflow_version_id"
],
"columnsTo": [
"id"
],
"onDelete": "no action",
"onUpdate": "no action"
},
"deployments_workflow_id_workflows_id_fk": {
"name": "deployments_workflow_id_workflows_id_fk",
"tableFrom": "deployments",
"tableTo": "workflows",
"columnsFrom": [
"workflow_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
},
"deployments_machine_id_machines_id_fk": {
"name": "deployments_machine_id_machines_id_fk",
"tableFrom": "deployments",
"tableTo": "machines",
"columnsFrom": [
"machine_id"
],
"columnsTo": [
"id"
],
"onDelete": "no action",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
},
"machines": {
"name": "machines",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"user_id": {
"name": "user_id",
"type": "text",
"primaryKey": false,
"notNull": true
},
"name": {
"name": "name",
"type": "text",
"primaryKey": false,
"notNull": true
},
"org_id": {
"name": "org_id",
"type": "text",
"primaryKey": false,
"notNull": false
},
"endpoint": {
"name": "endpoint",
"type": "text",
"primaryKey": false,
"notNull": true
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"disabled": {
"name": "disabled",
"type": "boolean",
"primaryKey": false,
"notNull": true,
"default": false
},
"auth_token": {
"name": "auth_token",
"type": "text",
"primaryKey": false,
"notNull": false
},
"type": {
"name": "type",
"type": "machine_type",
"primaryKey": false,
"notNull": true,
"default": "'classic'"
},
"status": {
"name": "status",
"type": "machine_status",
"primaryKey": false,
"notNull": true,
"default": "'ready'"
},
"snapshot": {
"name": "snapshot",
"type": "jsonb",
"primaryKey": false,
"notNull": false
},
"models": {
"name": "models",
"type": "jsonb",
"primaryKey": false,
"notNull": false
},
"gpu": {
"name": "gpu",
"type": "machine_gpu",
"primaryKey": false,
"notNull": false
},
"build_machine_instance_id": {
"name": "build_machine_instance_id",
"type": "text",
"primaryKey": false,
"notNull": false
},
"build_log": {
"name": "build_log",
"type": "text",
"primaryKey": false,
"notNull": false
}
},
"indexes": {},
"foreignKeys": {
"machines_user_id_users_id_fk": {
"name": "machines_user_id_users_id_fk",
"tableFrom": "machines",
"tableTo": "users",
"columnsFrom": [
"user_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
},
"users": {
"name": "users",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "text",
"primaryKey": true,
"notNull": true
},
"username": {
"name": "username",
"type": "text",
"primaryKey": false,
"notNull": true
},
"name": {
"name": "name",
"type": "text",
"primaryKey": false,
"notNull": true
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": false,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": false,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
},
"workflow_run_outputs": {
"name": "workflow_run_outputs",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"run_id": {
"name": "run_id",
"type": "uuid",
"primaryKey": false,
"notNull": true
},
"data": {
"name": "data",
"type": "jsonb",
"primaryKey": false,
"notNull": false
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"workflow_run_outputs_run_id_workflow_runs_id_fk": {
"name": "workflow_run_outputs_run_id_workflow_runs_id_fk",
"tableFrom": "workflow_run_outputs",
"tableTo": "workflow_runs",
"columnsFrom": [
"run_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
},
"workflow_runs": {
"name": "workflow_runs",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"workflow_version_id": {
"name": "workflow_version_id",
"type": "uuid",
"primaryKey": false,
"notNull": false
},
"workflow_inputs": {
"name": "workflow_inputs",
"type": "jsonb",
"primaryKey": false,
"notNull": false
},
"workflow_id": {
"name": "workflow_id",
"type": "uuid",
"primaryKey": false,
"notNull": true
},
"machine_id": {
"name": "machine_id",
"type": "uuid",
"primaryKey": false,
"notNull": false
},
"origin": {
"name": "origin",
"type": "workflow_run_origin",
"primaryKey": false,
"notNull": true,
"default": "'api'"
},
"status": {
"name": "status",
"type": "workflow_run_status",
"primaryKey": false,
"notNull": true,
"default": "'not-started'"
},
"ended_at": {
"name": "ended_at",
"type": "timestamp",
"primaryKey": false,
"notNull": false
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"workflow_runs_workflow_version_id_workflow_versions_id_fk": {
"name": "workflow_runs_workflow_version_id_workflow_versions_id_fk",
"tableFrom": "workflow_runs",
"tableTo": "workflow_versions",
"columnsFrom": [
"workflow_version_id"
],
"columnsTo": [
"id"
],
"onDelete": "set null",
"onUpdate": "no action"
},
"workflow_runs_workflow_id_workflows_id_fk": {
"name": "workflow_runs_workflow_id_workflows_id_fk",
"tableFrom": "workflow_runs",
"tableTo": "workflows",
"columnsFrom": [
"workflow_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
},
"workflow_runs_machine_id_machines_id_fk": {
"name": "workflow_runs_machine_id_machines_id_fk",
"tableFrom": "workflow_runs",
"tableTo": "machines",
"columnsFrom": [
"machine_id"
],
"columnsTo": [
"id"
],
"onDelete": "set null",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
},
"workflows": {
"name": "workflows",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"user_id": {
"name": "user_id",
"type": "text",
"primaryKey": false,
"notNull": true
},
"org_id": {
"name": "org_id",
"type": "text",
"primaryKey": false,
"notNull": false
},
"name": {
"name": "name",
"type": "text",
"primaryKey": false,
"notNull": true
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"workflows_user_id_users_id_fk": {
"name": "workflows_user_id_users_id_fk",
"tableFrom": "workflows",
"tableTo": "users",
"columnsFrom": [
"user_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
},
"workflow_versions": {
"name": "workflow_versions",
"schema": "comfyui_deploy",
"columns": {
"workflow_id": {
"name": "workflow_id",
"type": "uuid",
"primaryKey": false,
"notNull": true
},
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"workflow": {
"name": "workflow",
"type": "jsonb",
"primaryKey": false,
"notNull": false
},
"workflow_api": {
"name": "workflow_api",
"type": "jsonb",
"primaryKey": false,
"notNull": false
},
"version": {
"name": "version",
"type": "integer",
"primaryKey": false,
"notNull": true
},
"snapshot": {
"name": "snapshot",
"type": "jsonb",
"primaryKey": false,
"notNull": false
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"workflow_versions_workflow_id_workflows_id_fk": {
"name": "workflow_versions_workflow_id_workflows_id_fk",
"tableFrom": "workflow_versions",
"tableTo": "workflows",
"columnsFrom": [
"workflow_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
}
},
"enums": {
"deployment_environment": {
"name": "deployment_environment",
"values": {
"staging": "staging",
"production": "production"
}
},
"machine_gpu": {
"name": "machine_gpu",
"values": {
"T4": "T4",
"A10G": "A10G",
"A100": "A100"
}
},
"machine_status": {
"name": "machine_status",
"values": {
"ready": "ready",
"building": "building",
"error": "error"
}
},
"machine_type": {
"name": "machine_type",
"values": {
"classic": "classic",
"runpod-serverless": "runpod-serverless",
"modal-serverless": "modal-serverless",
"comfy-deploy-serverless": "comfy-deploy-serverless"
}
},
"workflow_run_origin": {
"name": "workflow_run_origin",
"values": {
"manual": "manual",
"api": "api"
}
},
"workflow_run_status": {
"name": "workflow_run_status",
"values": {
"not-started": "not-started",
"running": "running",
"uploading": "uploading",
"success": "success",
"failed": "failed"
}
}
},
"schemas": {
"comfyui_deploy": "comfyui_deploy"
},
"_meta": {
"schemas": {},
"tables": {},
"columns": {}
}
}

View File

@ -183,6 +183,13 @@
"when": 1704640269305,
"tag": "0025_hard_jean_grey",
"breakpoints": true
},
{
"idx": 26,
"version": "5",
"when": 1704979846175,
"tag": "0026_premium_rocket_raccoon",
"breakpoints": true
}
]
}

View File

@ -1,7 +1,9 @@
import { AccessType } from "../../../lib/AccessType";
import { MachineList } from "@/components/MachineList";
import { db } from "@/db/db";
import { machinesTable } from "@/db/schema";
import { auth } from "@clerk/nextjs";
import { clerkClient } from "@clerk/nextjs/server";
import { desc, eq, isNull, and } from "drizzle-orm";
export default function Page() {
@ -15,6 +17,8 @@ async function MachineListServer() {
return <div>No auth</div>;
}
const user = await clerkClient.users.getUser(userId);
const machines = await db.query.machinesTable.findMany({
orderBy: desc(machinesTable.updated_at),
where:
@ -26,7 +30,10 @@ async function MachineListServer() {
return (
<div className="w-full">
{/* <div>Machines</div> */}
<MachineList data={machines} />
<MachineList
data={machines}
userMetadata={AccessType.parse(user.privateMetadata ?? {})}
/>
</div>
);
}

View File

@ -23,6 +23,7 @@ import {
TableRow,
} from "@/components/ui/table";
import { type MachineType } from "@/db/schema";
import type { AccessType } from "@/lib/AccessType";
import {
addCustomMachineSchema,
addMachineSchema,
@ -53,6 +54,7 @@ import {
import { ArrowUpDown, MoreHorizontal } from "lucide-react";
import * as React from "react";
import { useState } from "react";
import type { z } from "zod";
export type Machine = MachineType;
@ -237,6 +239,9 @@ export const columns: ColumnDef<Machine>[] = [
models: {
fieldType: "models",
},
gpu: {
inputProps: {},
},
}}
/>
) : (
@ -263,7 +268,13 @@ export const columns: ColumnDef<Machine>[] = [
},
];
export function MachineList({ data }: { data: Machine[] }) {
export function MachineList({
data,
userMetadata,
}: {
data: Machine[];
userMetadata: z.infer<typeof AccessType>;
}) {
const [sorting, setSorting] = React.useState<SortingState>([]);
const [columnFilters, setColumnFilters] = React.useState<ColumnFiltersState>(
[]
@ -332,6 +343,14 @@ export function MachineList({ data }: { data: Machine[] }) {
models: {
fieldType: "models",
},
gpu: {
fieldType: !userMetadata.betaFeaturesAccess
? "fallback"
: "select",
inputProps: {
disabled: !userMetadata.betaFeaturesAccess,
},
},
}}
/>
<InsertModal

View File

@ -59,7 +59,7 @@ function SnapshotPresetPicker({
React.useEffect(() => {
findAllDeployments().then((a) => {
console.log(a);
// console.log(a);
const frameworks = a
.map((item) => {

View File

@ -104,6 +104,8 @@ export const workflowRunOrigin = pgEnum("workflow_run_origin", [
"api",
]);
export const machineGPUOptions = pgEnum("machine_gpu", ["T4", "A10G", "A100"]);
export const machinesType = pgEnum("machine_type", [
"classic",
"runpod-serverless",
@ -205,6 +207,7 @@ export const machinesTable = dbSchema.table("machines", {
status: machinesStatus("status").notNull().default("ready"),
snapshot: jsonb("snapshot").$type<any>(),
models: jsonb("models").$type<any>(),
gpu: machineGPUOptions("gpu"),
build_machine_instance_id: text("build_machine_instance_id"),
build_log: text("build_log"),
});

View File

@ -0,0 +1,5 @@
import { z } from "zod";
export const AccessType = z.object({
betaFeaturesAccess: z.boolean().default(false),
});

View File

@ -11,6 +11,7 @@ export const addMachineSchema = insertMachineSchema.pick({
export const insertCustomMachineSchema = createInsertSchema(machinesTable, {
name: (schema) => schema.name.default("My Machine"),
type: (schema) => schema.type.default("comfy-deploy-serverless"),
gpu: (schema) => schema.gpu.default("T4"),
snapshot: (schema) =>
schema.snapshot.default({
comfyui: "d0165d819afe76bd4e6bdd710eb5f3e571b6a804",
@ -37,4 +38,5 @@ export const addCustomMachineSchema = insertCustomMachineSchema.pick({
type: true,
snapshot: true,
models: true,
gpu: true,
});

View File

@ -93,6 +93,7 @@ export const updateCustomMachine = withServerPromise(
JSON.stringify(data.snapshot) !== JSON.stringify(currentMachine.snapshot);
const modelsChanged =
JSON.stringify(data.models) !== JSON.stringify(currentMachine.models);
const gpuChanged = data.gpu !== currentMachine.gpu;
// return {
// message: `snapshotChanged: ${snapshotChanged}, modelsChanged: ${modelsChanged}`,
@ -101,7 +102,7 @@ export const updateCustomMachine = withServerPromise(
await db.update(machinesTable).set(data).where(eq(machinesTable.id, id));
// If there are changes
if (snapshotChanged || modelsChanged) {
if (snapshotChanged || modelsChanged || gpuChanged) {
// Update status to building
await db
.update(machinesTable)
@ -174,7 +175,7 @@ async function buildMachine(
snapshot: data.snapshot, //JSON.parse( as string),
callback_url: `${protocol}://${domain}/api/machine-built`,
models: data.models, //JSON.parse(data.models as string),
gpu: "T4",
gpu: data.gpu ?? "T4",
}),
});