Skip to content

Commit

Permalink
feat: display remote workers
Browse files Browse the repository at this point in the history
  • Loading branch information
liangfung committed Nov 20, 2023
1 parent d2281c7 commit e7c00c7
Show file tree
Hide file tree
Showing 15 changed files with 2,849 additions and 77 deletions.
44 changes: 19 additions & 25 deletions ee/tabby-ui/app/(dashboard)/components/worker-card.tsx
Original file line number Diff line number Diff line change
@@ -1,33 +1,26 @@
import {
CardTitle,
CardHeader,
CardContent,
Card,
CardDescription
} from '@/components/ui/card'
import { HealthInfo } from '@/lib/hooks/use-health'
import { CardTitle, CardHeader, CardContent, Card } from '@/components/ui/card'
import { Worker, WorkerKind } from '@/lib/gql/generates/graphql'

type RunnerType = 'completion' | 'chat' | 'index'
type RunnerType = WorkerKind | 'INDEX'

interface RunnerCardProps {
source: string
name: string
type: RunnerType
health: HealthInfo
interface RunnerCardProps extends Partial<Omit<Worker, '__typename' | 'kind'>> {
kind: RunnerType
}

export default function RunnerCard({
source,
addr,
name,
type,
health
kind,
device,
cudaDevices,
cpuCount,
cpuInfo
}: RunnerCardProps) {
const { device, cuda_devices } = health
return (
<Card className="rounded-xl p-2 shadow-md">
<CardHeader className="p-0 px-4 pb-2 pt-4">
<CardTitle className="text-md flex items-center font-normal">
<ModelIcon type={type} />
<ModelIcon type={kind} />
<p className="ml-2">{name}</p>
</CardTitle>
</CardHeader>
Expand All @@ -51,7 +44,7 @@ export default function RunnerCard({
<path d="M5 16v-3a1 1 0 0 1 1-1h12a1 1 0 0 1 1 1v3" />
<path d="M12 12V8" />
</svg>
<p className="ml-2">{source}</p>
<p className="ml-2">{addr}</p>
</Info>
<Info>
<svg
Expand All @@ -78,11 +71,12 @@ export default function RunnerCard({
<path d="M9 20v2" />
</svg>
<p className="ml-2">
{health.cpu_info} ({health.cpu_count} cores)
{cpuInfo} ({cpuCount} cores)
</p>
</Info>
{device == 'cuda' &&
cuda_devices.map((x, i) => (
cudaDevices?.length &&
cudaDevices.map((x, i) => (
<Info key={i}>
<svg
className=" h-5 w-5 text-gray-400"
Expand Down Expand Up @@ -128,7 +122,7 @@ function Info({ children }: InfoProps) {

function ModelIcon({ type }: { type: RunnerType }) {
const className = 'h-5 w-5'
if (type == 'completion') {
if (type == WorkerKind.Completion) {
return (
<svg
xmlns="http://www.w3.org/2000/svg"
Expand All @@ -148,7 +142,7 @@ function ModelIcon({ type }: { type: RunnerType }) {
<path d="m14 17 2-2-2-2" />
</svg>
)
} else if (type == 'chat') {
} else if (type == WorkerKind.Chat) {
return (
<svg
xmlns="http://www.w3.org/2000/svg"
Expand All @@ -166,7 +160,7 @@ function ModelIcon({ type }: { type: RunnerType }) {
<path d="M18 9h2a2 2 0 0 1 2 2v11l-4-4h-6a2 2 0 0 1-2-2v-1" />
</svg>
)
} else if (type == 'index') {
} else if (type == 'INDEX') {
return (
<svg
xmlns="http://www.w3.org/2000/svg"
Expand Down
40 changes: 23 additions & 17 deletions ee/tabby-ui/app/(dashboard)/page.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@ import { Separator } from '@/components/ui/separator'
import { useHealth } from '@/lib/hooks/use-health'
import { PropsWithChildren, useEffect, useState } from 'react'
import WorkerCard from './components/worker-card'
import { useMergedWorkers } from '@/lib/hooks/use-remote-worker'
import { WorkerKind } from '@/lib/gql/generates/graphql'

const COMMUNITY_DIALOG_SHOWN_KEY = 'community-dialog-shown'

Expand Down Expand Up @@ -72,6 +74,7 @@ function toBadgeString(str: string) {

function MainPanel() {
const { data: healthInfo } = useHealth()
const workers = useMergedWorkers(healthInfo)

if (!healthInfo) return

Expand All @@ -98,26 +101,29 @@ function MainPanel() {
<div className="mt-4 rounded-lg bg-zinc-100 p-4 dark:bg-zinc-800">
<span className="font-bold">Workers</span>
<div className="mt-4 flex flex-col gap-4 lg:flex-row lg:flex-wrap">
{healthInfo.model &&
<WorkerCard
source="localhost"
name={healthInfo.model}
type="completion"
health={healthInfo}
/>}
{healthInfo.chat_model && (
<WorkerCard
source="localhost"
name={healthInfo.chat_model}
type="chat"
health={healthInfo}
/>
{!!workers?.[WorkerKind.Completion] && (
<>
{workers[WorkerKind.Completion].map((worker, i) => {
return <WorkerCard key={i} {...worker} />
})}
</>
)}
{!!workers?.[WorkerKind.Chat] && (
<>
{workers[WorkerKind.Chat].map((worker, i) => {
return <WorkerCard key={i} {...worker} />
})}
</>
)}
<WorkerCard
source="localhost"
addr="localhost"
name="Code Search Index"
type="index"
health={healthInfo}
kind="INDEX"
arch=""
device={healthInfo.device}
cudaDevices={healthInfo.cuda_devices}
cpuCount={healthInfo.cpu_count}
cpuInfo={healthInfo.cpu_info}
/>
</div>
</div>
Expand Down
19 changes: 19 additions & 0 deletions ee/tabby-ui/codegen.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@

import type { CodegenConfig } from '@graphql-codegen/cli';

const config: CodegenConfig = {
overwrite: true,
schema: "../tabby-webserver/graphql/schema.graphql",
documents: "./**/*.tsx",
// documents: "./**/*.graphql",
ignoreNoDocuments: true,
generates: {
"lib/gql/generates/": {
preset: "client",
plugins: []
}
},
hooks: { afterAllFileWrite: ['prettier --write'] }
};

export default config;
7 changes: 5 additions & 2 deletions ee/tabby-ui/components/header.tsx
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
'use client'

import * as React from 'react'

import { cn } from '@/lib/utils'
import { buttonVariants } from '@/components/ui/button'
import { IconGitHub, IconNotice } from '@/components/ui/icons'
Expand All @@ -10,6 +9,9 @@ import Link from 'next/link'
import { useHealth } from '@/lib/hooks/use-health'
import { ReleaseInfo, useLatestRelease } from '@/lib/hooks/use-latest-release'
import { compare } from 'compare-versions'
import { useMergedWorkers } from '@/lib/hooks/use-remote-worker'
import { WorkerKind } from '@/lib/gql/generates/graphql'
import { has } from 'lodash-es'

const ThemeToggle = dynamic(
() => import('@/components/theme-toggle').then(x => x.ThemeToggle),
Expand All @@ -18,7 +20,8 @@ const ThemeToggle = dynamic(

export function Header() {
const { data } = useHealth()
const isChatEnabled = !!data?.chat_model
const workers = useMergedWorkers(data)
const isChatEnabled = has(workers, WorkerKind.Chat)
const version = data?.version?.git_describe
const { data: latestRelease } = useLatestRelease()
const newVersionAvailable = isNewVersionAvailable(version, latestRelease)
Expand Down
85 changes: 85 additions & 0 deletions ee/tabby-ui/lib/gql/generates/fragment-masking.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,85 @@
import {
ResultOf,
DocumentTypeDecoration,
TypedDocumentNode
} from '@graphql-typed-document-node/core'
import { FragmentDefinitionNode } from 'graphql'
import { Incremental } from './graphql'

export type FragmentType<
TDocumentType extends DocumentTypeDecoration<any, any>
> = TDocumentType extends DocumentTypeDecoration<infer TType, any>
? [TType] extends [{ ' $fragmentName'?: infer TKey }]
? TKey extends string
? { ' $fragmentRefs'?: { [key in TKey]: TType } }
: never
: never
: never

// return non-nullable if `fragmentType` is non-nullable
export function useFragment<TType>(
_documentNode: DocumentTypeDecoration<TType, any>,
fragmentType: FragmentType<DocumentTypeDecoration<TType, any>>
): TType
// return nullable if `fragmentType` is nullable
export function useFragment<TType>(
_documentNode: DocumentTypeDecoration<TType, any>,
fragmentType:
| FragmentType<DocumentTypeDecoration<TType, any>>
| null
| undefined
): TType | null | undefined
// return array of non-nullable if `fragmentType` is array of non-nullable
export function useFragment<TType>(
_documentNode: DocumentTypeDecoration<TType, any>,
fragmentType: ReadonlyArray<FragmentType<DocumentTypeDecoration<TType, any>>>
): ReadonlyArray<TType>
// return array of nullable if `fragmentType` is array of nullable
export function useFragment<TType>(
_documentNode: DocumentTypeDecoration<TType, any>,
fragmentType:
| ReadonlyArray<FragmentType<DocumentTypeDecoration<TType, any>>>
| null
| undefined
): ReadonlyArray<TType> | null | undefined
export function useFragment<TType>(
_documentNode: DocumentTypeDecoration<TType, any>,
fragmentType:
| FragmentType<DocumentTypeDecoration<TType, any>>
| ReadonlyArray<FragmentType<DocumentTypeDecoration<TType, any>>>
| null
| undefined
): TType | ReadonlyArray<TType> | null | undefined {
return fragmentType as any
}

export function makeFragmentData<
F extends DocumentTypeDecoration<any, any>,
FT extends ResultOf<F>
>(data: FT, _fragment: F): FragmentType<F> {
return data as FragmentType<F>
}
export function isFragmentReady<TQuery, TFrag>(
queryNode: DocumentTypeDecoration<TQuery, any>,
fragmentNode: TypedDocumentNode<TFrag>,
data:
| FragmentType<TypedDocumentNode<Incremental<TFrag>, any>>
| null
| undefined
): data is FragmentType<typeof fragmentNode> {
const deferredFields = (
queryNode as {
__meta__?: { deferredFields: Record<string, (keyof TFrag)[]> }
}
).__meta__?.deferredFields

if (!deferredFields) return true

const fragDef = fragmentNode.definitions[0] as
| FragmentDefinitionNode
| undefined
const fragName = fragDef?.name?.value

const fields = (fragName && deferredFields[fragName]) || []
return fields.length > 0 && fields.every(field => data && field in data)
}
46 changes: 46 additions & 0 deletions ee/tabby-ui/lib/gql/generates/gql.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
/* eslint-disable */
import * as types from './graphql'
import { TypedDocumentNode as DocumentNode } from '@graphql-typed-document-node/core'

/**
* Map of all GraphQL operations in the project.
*
* This map has several performance disadvantages:
* 1. It is not tree-shakeable, so it will include all operations in the project.
* 2. It is not minifiable, so the string of a GraphQL query will be multiple times inside the bundle.
* 3. It does not support dead code elimination, so it will add unused operations.
*
* Therefore it is highly recommended to use the babel or swc plugin for production.
*/
const documents = {
'\n query GetWorkers {\n workers {\n kind\n name\n addr\n device\n arch\n cpuInfo\n cpuCount\n cudaDevices\n }\n }\n':
types.GetWorkersDocument
}

/**
* The graphql function is used to parse GraphQL queries into a document that can be used by GraphQL clients.
*
*
* @example
* ```ts
* const query = graphql(`query GetUser($id: ID!) { user(id: $id) { name } }`);
* ```
*
* The query argument is unknown!
* Please regenerate the types.
*/
export function graphql(source: string): unknown

/**
* The graphql function is used to parse GraphQL queries into a document that can be used by GraphQL clients.
*/
export function graphql(
source: '\n query GetWorkers {\n workers {\n kind\n name\n addr\n device\n arch\n cpuInfo\n cpuCount\n cudaDevices\n }\n }\n'
): (typeof documents)['\n query GetWorkers {\n workers {\n kind\n name\n addr\n device\n arch\n cpuInfo\n cpuCount\n cudaDevices\n }\n }\n']

export function graphql(source: string) {
return (documents as any)[source] ?? {}
}

export type DocumentType<TDocumentNode extends DocumentNode<any, any>> =
TDocumentNode extends DocumentNode<infer TType, any> ? TType : never
Loading

0 comments on commit e7c00c7

Please sign in to comment.