import { useMemo } from 'react'; import { components, MultiValue } from 'react-select'; import { MultiValueRemoveProps } from 'react-select/dist/declarations/src/components/MultiValue'; import { ActionMeta, OnChangeValue, } from 'react-select/dist/declarations/src/types'; import { OptionProps } from 'react-select/dist/declarations/src/components/Option'; import { Select } from '@@/form-components/ReactSelect'; import { Switch } from '@@/form-components/SwitchField/Switch'; import { Tooltip } from '@@/Tip/Tooltip'; interface Values { enabled: boolean; useSpecific: boolean; selectedGPUs: string[]; capabilities: string[]; } interface GpuOption { value: string; label: string; description?: string; } export interface GPU { value: string; name: string; } export interface Props { values: Values; onChange(values: Values): void; gpus: GPU[]; usedGpus: string[]; usedAllGpus: boolean; } const NvidiaCapabilitiesOptions = [ // Taken from https://github.com/containerd/containerd/blob/master/contrib/nvidia/nvidia.go#L40 { value: 'compute', label: 'compute', description: 'required for CUDA and OpenCL applications', }, { value: 'compat32', label: 'compat32', description: 'required for running 32-bit applications', }, { value: 'graphics', label: 'graphics', description: 'required for running OpenGL and Vulkan applications', }, { value: 'utility', label: 'utility', description: 'required for using nvidia-smi and NVML', }, { value: 'video', label: 'video', description: 'required for using the Video Codec SDK', }, { value: 'display', label: 'display', description: 'required for leveraging X11 display', }, ]; function Option(props: OptionProps) { const { data: { value, description }, } = props; return (
{/* eslint-disable-next-line react/jsx-props-no-spreading */} {`${value} - ${description}`}
); } function MultiValueRemove(props: MultiValueRemoveProps) { const { selectProps: { value }, } = props; if (value && (value as MultiValue).length === 1) { return null; } // eslint-disable-next-line react/jsx-props-no-spreading return ; } export function Gpu({ values, onChange, gpus = [], usedGpus = [], usedAllGpus, }: Props) { const options = useMemo(() => { const options = (gpus || []).map((gpu) => ({ value: gpu.value, label: usedGpus.includes(gpu.value) || usedAllGpus ? `${gpu.name} (in use)` : gpu.name, })); options.unshift({ value: 'all', label: 'Use All GPUs', }); return options; }, [gpus, usedGpus, usedAllGpus]); function onChangeValues(key: string, newValue: boolean | string[]) { const newValues = { ...values, [key]: newValue, }; onChange(newValues); } function toggleEnableGpu() { onChangeValues('enabled', !values.enabled); } function onChangeSelectedGpus( newValue: OnChangeValue, actionMeta: ActionMeta ) { let { useSpecific } = values; let selectedGPUs = newValue.map((option) => option.value); if (actionMeta.action === 'select-option') { useSpecific = actionMeta.option?.value !== 'all'; selectedGPUs = selectedGPUs.filter((value) => useSpecific ? value !== 'all' : value === 'all' ); } const newValues = { ...values, selectedGPUs, useSpecific }; onChange(newValues); } function onChangeSelectedCaps(newValue: OnChangeValue) { onChangeValues( 'capabilities', newValue.map((option) => option.value) ); } const gpuCmd = useMemo(() => { const devices = values.selectedGPUs.join(','); const deviceStr = devices === 'all' ? 'all,' : `device=${devices},`; const caps = values.capabilities.join(','); return `--gpus '${deviceStr}"capabilities=${caps}"'`; }, [values.selectedGPUs, values.capabilities]); const gpuValue = useMemo( () => options.filter((option) => values.selectedGPUs.includes(option.value)), [values.selectedGPUs, options] ); const capValue = useMemo( () => NvidiaCapabilitiesOptions.filter((option) => values.capabilities.includes(option.value) ), [values.capabilities] ); return (
Enable GPU
isMulti closeMenuOnSelect value={gpuValue} isClearable={false} backspaceRemovesValue={false} isDisabled={!values.enabled} onChange={onChangeSelectedGpus} options={options} components={{ MultiValueRemove }} />
{values.enabled && ( <>
Capabilities
isMulti closeMenuOnSelect value={capValue} options={NvidiaCapabilitiesOptions} components={{ Option }} onChange={onChangeSelectedCaps} />
Control
{gpuCmd}
)}
); }