132 lines
3.6 KiB
TypeScript
132 lines
3.6 KiB
TypeScript
import { DeviceRequest } from 'docker-types/generated/1.44';
|
|
|
|
import { computeDockerGPUCommand } from './GpuRow';
|
|
|
|
describe('computeDockerGPUCommand', () => {
|
|
it('should return "No GPU config found" when deviceRequests is empty', () => {
|
|
const result = computeDockerGPUCommand([]);
|
|
expect(result).toBe('No GPU config found');
|
|
});
|
|
|
|
it('should return "No GPU config found" when no GPU request is found', () => {
|
|
const deviceRequests: Array<DeviceRequest> = [
|
|
{
|
|
Driver: 'other',
|
|
Count: 1,
|
|
DeviceIDs: [],
|
|
Capabilities: [],
|
|
Options: {},
|
|
},
|
|
];
|
|
const result = computeDockerGPUCommand(deviceRequests);
|
|
expect(result).toBe('No GPU config found');
|
|
});
|
|
|
|
it('should return "all" with capabilities for nvidia driver with Count=-1', () => {
|
|
const deviceRequests: Array<DeviceRequest> = [
|
|
{
|
|
Driver: 'nvidia',
|
|
Count: -1,
|
|
DeviceIDs: [],
|
|
Capabilities: [['gpu', 'utility']],
|
|
Options: {},
|
|
},
|
|
];
|
|
const result = computeDockerGPUCommand(deviceRequests);
|
|
expect(result).toBe('all,"capabilities=gpu,utility"');
|
|
});
|
|
|
|
it('should return device IDs with capabilities when Count is not -1', () => {
|
|
const deviceRequests: Array<DeviceRequest> = [
|
|
{
|
|
Driver: 'nvidia',
|
|
Count: 2,
|
|
DeviceIDs: ['0', '1'],
|
|
Capabilities: [['gpu', 'compute']],
|
|
Options: {},
|
|
},
|
|
];
|
|
const result = computeDockerGPUCommand(deviceRequests);
|
|
expect(result).toBe('"device=0,1","capabilities=gpu,compute"');
|
|
});
|
|
|
|
it('should work with gpu capability detection instead of nvidia driver', () => {
|
|
const deviceRequests: Array<DeviceRequest> = [
|
|
{
|
|
Driver: '',
|
|
Count: -1,
|
|
DeviceIDs: [],
|
|
Capabilities: [['gpu']],
|
|
Options: {},
|
|
},
|
|
];
|
|
const result = computeDockerGPUCommand(deviceRequests);
|
|
expect(result).toBe('all,"capabilities=gpu"');
|
|
});
|
|
|
|
it('should return "all" without capabilities when Capabilities is undefined', () => {
|
|
const deviceRequests: Array<DeviceRequest> = [
|
|
{
|
|
Driver: 'nvidia',
|
|
Count: -1,
|
|
DeviceIDs: [],
|
|
Capabilities: undefined,
|
|
Options: {},
|
|
},
|
|
];
|
|
const result = computeDockerGPUCommand(deviceRequests);
|
|
expect(result).toBe('all,');
|
|
});
|
|
|
|
it('should handle multiple device IDs correctly', () => {
|
|
const deviceRequests: Array<DeviceRequest> = [
|
|
{
|
|
Driver: 'nvidia',
|
|
Count: 3,
|
|
DeviceIDs: ['GPU-0', 'GPU-1', 'GPU-2'],
|
|
Capabilities: [['gpu', 'utility', 'compute']],
|
|
Options: {},
|
|
},
|
|
];
|
|
const result = computeDockerGPUCommand(deviceRequests);
|
|
expect(result).toBe(
|
|
'"device=GPU-0,GPU-1,GPU-2","capabilities=gpu,utility,compute"'
|
|
);
|
|
});
|
|
|
|
it('should find GPU request when it is not the first element', () => {
|
|
const deviceRequests: Array<DeviceRequest> = [
|
|
{
|
|
Driver: 'other',
|
|
Count: 1,
|
|
DeviceIDs: [],
|
|
Capabilities: [],
|
|
Options: {},
|
|
},
|
|
{
|
|
Driver: 'nvidia',
|
|
Count: -1,
|
|
DeviceIDs: [],
|
|
Capabilities: [['gpu']],
|
|
Options: {},
|
|
},
|
|
];
|
|
const result = computeDockerGPUCommand(deviceRequests);
|
|
expect(result).toBe('all,"capabilities=gpu"');
|
|
});
|
|
|
|
it('should handle empty capabilities array', () => {
|
|
const deviceRequests: Array<DeviceRequest> = [
|
|
{
|
|
Driver: 'nvidia',
|
|
Count: -1,
|
|
DeviceIDs: [],
|
|
Capabilities: [[]],
|
|
Options: {},
|
|
},
|
|
];
|
|
const result = computeDockerGPUCommand(deviceRequests);
|
|
expect(result).toBe('all,"capabilities="');
|
|
});
|
|
});
|