skills/civitai__civitai
1
0
Fork 0
mirror of https://github.com/civitai/civitai.git synced 2026-09-20 22:08:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix issues with moderator review page

Browse Source
This commit is contained in:
Briant Diehl
2025-11-20 17:52:12 -07:00
parent e724dbef65
commit bd7e9d5309
6 changed files with 164 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.eslintrc.js
+1 -1
View File
@@ -25,7 +25,7 @@ module.exports = {
// aligns closing brackets for tags
'react/jsx-closing-bracket-location': ['error', 'line-aligned'],
'import/no-cycle': ['error'],
// 'import/no-cycle': ['error'],
// prettier overrides
'prettier/prettier': [
src/pages/api/webhooks/resource-training-v2/[modelVersionId].ts
+42 -37
View File
@@ -266,43 +266,48 @@ export async function updateRecords(workflow: Workflow, status: WorkflowStatus)
} catch {}
}
try {
const bodyData: TrainingUpdateSignalSchema = {
modelId: model.id,
modelVersionId: modelVersion.id,
status: trainingStatus,
fileMetadata: newMetadata,
};
await fetch(
`${env.SIGNALS_ENDPOINT}/users/${model.user.id}/signals/${SignalMessages.TrainingUpdate}`,
{
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify(bodyData),
}
);
} catch (e: unknown) {
logWebhook({
message: 'Failed to send signal for update',
data: { error: (e as Error)?.message, cause: (e as Error)?.cause, workflowId: workflow.id },
});
}
if (last?.status !== trainingStatus) {
try {
const bodyData: TrainingUpdateSignalSchema = {
modelId: model.id,
modelVersionId: modelVersion.id,
status: trainingStatus,
fileMetadata: newMetadata,
};
await fetch(
`${env.SIGNALS_ENDPOINT}/users/${model.user.id}/signals/${SignalMessages.TrainingUpdate}`,
{
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify(bodyData),
}
);
} catch (e: unknown) {
logWebhook({
message: 'Failed to send signal for update',
data: { error: (e as Error)?.message, cause: (e as Error)?.cause, workflowId: workflow.id },
});
}
if (trainingStatus === TrainingStatus.InReview) {
trainingCompleteEmail
.send({
model,
mName: modelVersion.name,
user: model.user,
})
.catch((error) => logWebhook({ message: 'Failed to send training complete email', error }));
} else if (trainingStatus === TrainingStatus.Failed || trainingStatus === TrainingStatus.Denied) {
trainingFailEmail
.send({
model,
mName: modelVersion.name,
user: model.user,
})
.catch((error) => logWebhook({ message: 'Failed to send training fail email', error }));
if (trainingStatus === TrainingStatus.InReview) {
trainingCompleteEmail
.send({
model,
mName: modelVersion.name,
user: model.user,
})
.catch((error) => logWebhook({ message: 'Failed to send training complete email', error }));
} else if (
trainingStatus === TrainingStatus.Failed ||
trainingStatus === TrainingStatus.Denied
) {
trainingFailEmail
.send({
model,
mName: modelVersion.name,
user: model.user,
})
.catch((error) => logWebhook({ message: 'Failed to send training fail email', error }));
}
}
}
src/pages/moderator/review/training-data/index.tsx
+38 -36
View File
@@ -1,4 +1,4 @@
import { Button, Center, Group, Loader, Text, Title, Tooltip } from '@mantine/core';
import { Button, Center, Group, Loader, Text, Title, Tooltip, Badge } from '@mantine/core';
import { IconRefresh } from '@tabler/icons-react';
import { useRouter } from 'next/router';
import { useMemo } from 'react';
@@ -38,7 +38,7 @@ export default function ReviewTrainingDataPage() {
});
const handleRecheckTrainingStatus = (modelVersionId: number) => {
recheckTrainingStatusMutation.mutate({ id: modelVersionId });
queryUtils.moderator.modelVersions.invalidate();
};
const flatData = useMemo(() => data?.pages.flatMap((x) => x.items), [data]);
@@ -60,41 +60,43 @@ export default function ReviewTrainingDataPage() {
</Center>
) : (
<div className="flex flex-col gap-3">
{flatData?.map((item) => (
<div key={item.id} className="flex items-center justify-between gap-3 p-3 card">
<div className="flex flex-col items-center">
<Text lineClamp={1}>
{item.model.name} - {item.name}
</Text>
<Text c="dimmed" size="xs">
Created: {formatDate(item.createdAt)}
</Text>
<Text c="dimmed" size="xs">
WorkflowId: {item.workflowId}
</Text>
</div>
<Group>
<Button size="compact-sm" component={Link} href={`${router.asPath}/${item.id}`}>
Review
</Button>
{flatData?.map((item) => {
return (
<div key={item.id} className="flex items-center justify-between gap-3 p-3 card">
<div className="flex flex-col items-center">
<Text lineClamp={1}>
{item.model.name} - {item.name}
</Text>
<Text c="dimmed" size="xs">
Created: {formatDate(item.createdAt)}
</Text>
<Text c="dimmed" size="xs">
WorkflowId: {item.workflowId}
</Text>
</div>
<Group>
<Button size="compact-sm" component={Link} href={`${router.asPath}/${item.id}`}>
Review
</Button>
<Tooltip label="Recheck Training Status" withArrow>
<LegacyActionIcon
variant="light"
size="md"
radius="xl"
loading={
recheckTrainingStatusMutation.isLoading &&
recheckTrainingStatusMutation.variables?.id === item.id
}
onClick={() => handleRecheckTrainingStatus(item.id)}
>
<IconRefresh size={16} />
</LegacyActionIcon>
</Tooltip>
</Group>
</div>
))}
<Tooltip label="Recheck Training Status" withArrow>
<LegacyActionIcon
variant="light"
size="md"
radius="xl"
loading={
recheckTrainingStatusMutation.isLoading &&
recheckTrainingStatusMutation.variables?.id === item.id
}
onClick={() => handleRecheckTrainingStatus(item.id)}
>
<IconRefresh size={16} />
</LegacyActionIcon>
</Tooltip>
</Group>
</div>
);
})}
</div>
)}
{hasNextPage && (
src/server/controllers/model-version.controller.ts
+33 -11
View File
@@ -75,6 +75,7 @@ import { env } from '~/env/server';
import { getWorkflow } from '~/server/services/orchestrator/workflows';
import { WorkflowStatus } from '@civitai/client';
import { getAllowedAccountTypes } from '~/server/utils/buzz-helpers';
import { isDefined } from '~/utils/type-guards';
export const getModelVersionRunStrategiesHandler = ({ input: { id } }: { input: GetByIdInput }) => {
try {
@@ -745,19 +746,40 @@ export async function queryModelVersionsForModeratorHandler({
},
});
const workflowIds: string[] = [];
const mappedItems = items.map(({ files, meta, ...version }) => {
const trainingFile = files[0];
const trainingResults = (trainingFile?.metadata as FileMetadata)
?.trainingResults as TrainingResultsV2;
if (trainingResults?.workflowId) workflowIds.push(trainingResults.workflowId);
return {
...version,
meta: meta as ModelVersionMeta | null,
workflowId: trainingResults?.workflowId,
};
});
/*
querying the workflows here may seem pointless, but querying the workflow can cause the orchestrator to take action on a workflow with failed/expired jobs.
Perhaps we need to move this to a method that can be called from the client to refresh the list as needed
*/
const workflows = await Promise.all(
workflowIds.map((workflowId) =>
getWorkflow({ token: env.ORCHESTRATOR_ACCESS_TOKEN, path: { workflowId } }).catch(() => null)
)
);
return {
nextCursor,
items: items.map(({ files, meta, ...version }) => {
const trainingFile = files[0];
const trainingResults = (trainingFile?.metadata as FileMetadata)
?.trainingResults as TrainingResultsV2;
return {
...version,
meta: meta as ModelVersionMeta | null,
workflowId: trainingResults?.workflowId,
};
}),
items: mappedItems
.map((item) => ({
...item,
workflow: workflows.find((x) => x && x.id === item.workflowId),
}))
.filter((x) => x.workflow),
};
}
src/server/controllers/training.controller.ts
+49 -30
View File
@@ -1,3 +1,4 @@
import type { WorkflowStatus } from '@civitai/client';
import { TRPCError } from '@trpc/server';
import { env } from '~/env/server';
import type { CustomImageResourceTrainingStep } from '~/pages/api/webhooks/resource-training-v2/[modelVersionId]';
@@ -106,6 +107,8 @@ const getJobIdFromVersion = async (modelVersionId: number) => {
path: { workflowId },
});
if (!workflow) throw new Error(`Could not find workflow with id: ${workflowId}`);
const step = workflow.steps?.[0] as CustomImageResourceTrainingStep | undefined;
// nb: get exactly the second job
const gateId = step?.jobs?.[1]?.id;
@@ -117,58 +120,74 @@ const getJobIdFromVersion = async (modelVersionId: number) => {
throw throwNotFoundError('Could not find jobId for gate job');
}
return gateId;
return { workflowId: workflow.id, status: workflow.status, gateId };
};
const moderateTrainingData = async ({
modelVersionId,
gateJobId,
gateId,
approve,
workflowId,
status,
}: {
modelVersionId: number;
gateJobId: string;
gateId: string;
approve: boolean;
workflowId?: string | null;
status?: WorkflowStatus;
}) => {
if (!env.ORCHESTRATOR_ENDPOINT) throw throwInternalServerError('No orchestrator endpoint');
try {
const gateResp = await fetch(
`${env.ORCHESTRATOR_ENDPOINT}/v1/manager/ambientjobs/${gateJobId}`,
{
method: 'PUT',
body: JSON.stringify({
approved: approve,
// message: ''
}),
headers: {
'Content-Type': 'application/json',
Authorization: `Bearer ${env.ORCHESTRATOR_ACCESS_TOKEN}`,
},
}
);
const response = await fetch(`${env.ORCHESTRATOR_ENDPOINT}/v1/manager/ambientjobs/${gateId}`, {
method: 'PUT',
body: JSON.stringify({
approved: approve,
// message: ''
}),
headers: {
'Content-Type': 'application/json',
Authorization: `Bearer ${env.ORCHESTRATOR_ACCESS_TOKEN}`,
},
});
if (!gateResp.ok) {
if (response.ok) {
if (workflowId && status) {
// handle calling the webhook endpoint. Resolves an issue with the orchestrator that has been plaguing us
await fetch(
`https://api.civitai.com/webhooks/resource-training-v2/${modelVersionId}?token=${env.WEBHOOK_TOKEN}`,
{
method: 'POST',
body: JSON.stringify({ workflowId, status }),
headers: {
'Content-Type': 'application/json',
},
}
);
}
logWebhook({
message: `${approve ? 'Approved' : 'Denied'} training dataset`,
type: 'info',
data: { modelVersionId },
});
} else {
logWebhook({
message: 'Could not connect to orchestrator',
data: {
modelVersionId,
important: true,
status: gateResp.status,
gateJobId,
status: response.status,
gateJobId: gateId,
},
});
if (gateResp.status === 429) {
if (response.status === 429) {
throw throwRateLimitError('Could not connect to orchestrator');
} else {
throw throwBadRequestError('Could not connect to orchestrator');
}
}
logWebhook({
message: `${approve ? 'Approved' : 'Denied'} training dataset`,
type: 'info',
data: { modelVersionId },
});
return 'ok';
} catch (e) {
logWebhook({
@@ -186,12 +205,12 @@ const moderateTrainingData = async ({
export async function handleApproveTrainingData({ input }: { input: GetByIdInput }) {
const modelVersionId = input.id;
const gateJobId = await getJobIdFromVersion(modelVersionId);
return await moderateTrainingData({ modelVersionId, gateJobId, approve: true });
const { gateId, workflowId, status } = await getJobIdFromVersion(modelVersionId);
return await moderateTrainingData({ modelVersionId, gateId, workflowId, status, approve: true });
}
export async function handleDenyTrainingData({ input }: { input: GetByIdInput }) {
const modelVersionId = input.id;
const gateJobId = await getJobIdFromVersion(modelVersionId);
return await moderateTrainingData({ modelVersionId, gateJobId, approve: false });
const { gateId, workflowId, status } = await getJobIdFromVersion(modelVersionId);
return await moderateTrainingData({ modelVersionId, gateId, workflowId, status, approve: false });
}
src/server/metrics/model.metrics.ts
+1 -2
View File
@@ -226,8 +226,7 @@ async function getModelMetrics(ctx: ModelMetricContext, sql: string, params: any
if (key === 'modelId') continue;
const value = row[key];
if (value == null) continue;
(ctx.modelUpdates[modelId] as any)[key] =
typeof value === 'string' ? parseInt(value) : value;
(ctx.modelUpdates[modelId] as any)[key] = typeof value === 'string' ? parseInt(value) : value;
}
}
}