fix: Display properly rounded subset counts (#6748)

Author: dwesolow

application/ui/src/features/models/model-listing/components/three-section-range/three-section-range.component.tsx

@@ -3,6 +3,8 @@
 
 import { Flex, Text } from '@geti/ui';
 
+import { distributeByLargestRemainder } from '../../../utils';
+
 import classes from './three-section-range.module.scss';
 
 type ThreeSectionRangeProps = {
@@ -12,38 +14,13 @@ type ThreeSectionRangeProps = {
     testingValue: number;
 };
 
-// Distribute rounded percentages so that they always sum to 100%
-const computeRoundedPercentages = (values: number[]): number[] => {
-    const total = values.reduce((sum, value) => sum + value, 0);
-
-    if (total <= 0) {
-        return values.map(() => 0);
-    }
-
-    const exactPercentages = values.map((value) => (value * 100) / total);
-    const flooredPercentages = exactPercentages.map((percentage) => Math.floor(percentage));
-    let remainder = 100 - flooredPercentages.reduce((sum, value) => sum + value, 0);
-
-    const indicesByRemainder = exactPercentages
-        .map((percentage, index) => ({ index, fractional: percentage - Math.floor(percentage) }))
-        .sort((a, b) => b.fractional - a.fractional);
-
-    const result = [...flooredPercentages];
-    for (const { index } of indicesByRemainder) {
-        if (remainder <= 0) break;
-        result[index] += 1;
-        remainder -= 1;
-    }
-
-    return result;
-};
+const MAX_VALUE = 100;
 
 export const ThreeSectionRange = ({ id, trainingValue, validationValue, testingValue }: ThreeSectionRangeProps) => {
-    const [trainingPercentage, validationPercentage, testingPercentage] = computeRoundedPercentages([
-        trainingValue,
-        validationValue,
-        testingValue,
-    ]);
+    const [trainingPercentage, validationPercentage, testingPercentage] = distributeByLargestRemainder(
+        [trainingValue, validationValue, testingValue],
+        MAX_VALUE
+    );
 
     const labelledPercentages = [
         { label: 'Training', percentage: trainingPercentage, color: 'var(--training-subset)' },

application/ui/src/features/models/train-model/advanced-settings/data-management/training-subsets/resulting-dataset-distribution.component.tsx

@@ -3,15 +3,16 @@
 
 import { Flex, Grid, repeat, Text, View } from '@geti/ui';
 
+import { distributeByLargestRemainder } from '../../../../utils';
 import { LABEL_COLOR_MAPPING, SubsetTile } from './subset-distribution-stats.component';
 
 import classes from './training-subsets.module.scss';
 
 type SubsetDistributionRowProps = {
     existingSize: number;
     newSize: number;
-    totalSize: number;
     label: string;
+    percentage: number;
 };
 
 const SubsetLabel = ({ label, color }: { color: string; label: string }) => {
@@ -23,9 +24,8 @@ const SubsetLabel = ({ label, color }: { color: string; label: string }) => {
     );
 };
 
-const SubsetDistributionRow = ({ existingSize, newSize, totalSize, label }: SubsetDistributionRowProps) => {
+const SubsetDistributionRow = ({ existingSize, newSize, label, percentage }: SubsetDistributionRowProps) => {
     const resultingSize = existingSize + newSize;
-    const percentage = totalSize > 0 ? Math.round((resultingSize * 100) / totalSize) : 0;
 
     return (
         <>
@@ -42,23 +42,28 @@ const SubsetDistributionRow = ({ existingSize, newSize, totalSize, label }: Subs
 type ResultingDatasetDistributionSubsetProps = {
     color: string;
     label: string;
-    totalSize: number;
     newSize: number;
     existingSize: number;
+    percentage: number;
 };
 
 const ResultingDatasetDistributionSubset = ({
     color,
     label,
     existingSize,
     newSize,
-    totalSize,
+    percentage,
 }: ResultingDatasetDistributionSubsetProps) => {
     return (
         <>
             <SubsetLabel label={label} color={color} />
 
-            <SubsetDistributionRow label={label} totalSize={totalSize} newSize={newSize} existingSize={existingSize} />
+            <SubsetDistributionRow
+                label={label}
+                newSize={newSize}
+                existingSize={existingSize}
+                percentage={percentage}
+            />
         </>
     );
 };
@@ -70,18 +75,27 @@ type ResultingDatasetDistributionProps = {
     newTrainingSubsetSize: number;
     newValidationSubsetSize: number;
     newTestingSubsetSize: number;
-    totalDatasetItemsSize: number;
 };
 
+const MAX_VALUE = 100;
+
 export const ResultingDatasetDistribution = ({
     trainingSubsetSize,
     validationSubsetSize,
     testingSubsetSize,
     newTrainingSubsetSize,
     newValidationSubsetSize,
     newTestingSubsetSize,
-    totalDatasetItemsSize,
 }: ResultingDatasetDistributionProps) => {
+    const [trainingPercentage, validationPercentage, testingPercentage] = distributeByLargestRemainder(
+        [
+            trainingSubsetSize + newTrainingSubsetSize,
+            validationSubsetSize + newValidationSubsetSize,
+            testingSubsetSize + newTestingSubsetSize,
+        ],
+        MAX_VALUE
+    );
+
     return (
         <Flex direction={'column'} gap={'size-50'}>
             <Text>Resulting dataset distribution:</Text>
@@ -96,25 +110,25 @@ export const ResultingDatasetDistribution = ({
                     <ResultingDatasetDistributionSubset
                         label={'Training'}
                         color={LABEL_COLOR_MAPPING.training}
-                        totalSize={totalDatasetItemsSize}
                         newSize={newTrainingSubsetSize}
                         existingSize={trainingSubsetSize}
+                        percentage={trainingPercentage}
                     />
 
                     <ResultingDatasetDistributionSubset
                         label={'Validation'}
                         color={LABEL_COLOR_MAPPING.validation}
-                        totalSize={totalDatasetItemsSize}
                         newSize={newValidationSubsetSize}
                         existingSize={validationSubsetSize}
+                        percentage={validationPercentage}
                     />
 
                     <ResultingDatasetDistributionSubset
                         label={'Test'}
                         color={LABEL_COLOR_MAPPING.test}
-                        totalSize={totalDatasetItemsSize}
                         newSize={newTestingSubsetSize}
                         existingSize={testingSubsetSize}
+                        percentage={testingPercentage}
                     />
                 </Grid>
             </View>

application/ui/src/features/models/train-model/advanced-settings/data-management/training-subsets/training-subsets.component.tsx

@@ -9,6 +9,7 @@ import { isEqual } from 'lodash-es';
 
 import type { ConfigurableParameter, TrainingConfiguration } from '../../../../../../constants/shared-types';
 import { isParameterGroup } from '../../../../model-listing/model-training-parameters/utils';
+import { distributeByLargestRemainder } from '../../../../utils';
 import { Accordion } from '../../components/accordion/accordion.component';
 import { ResultingDatasetDistribution } from './resulting-dataset-distribution.component';
 import { SubsetsDistribution } from './subset-distribution.component';
@@ -163,9 +164,10 @@ export const TrainingSubsets = ({
         });
     };
 
-    const newValidationSubsetSize = Math.floor((validationSubsetRatio / 100) * unassignedSubsetSize);
-    const newTestingSubsetSize = Math.floor((testSubsetRatio / 100) * unassignedSubsetSize);
-    const newTrainingSubsetSize = unassignedSubsetSize - newValidationSubsetSize - newTestingSubsetSize;
+    const [newTrainingSubsetSize, newValidationSubsetSize, newTestingSubsetSize] = distributeByLargestRemainder(
+        [trainingSubsetRatio, validationSubsetRatio, testSubsetRatio],
+        unassignedSubsetSize
+    );
 
     const areSubsetsSizesValid = () => {
         const resultingTrainingSubsetSize = trainingSubsetSize + newTrainingSubsetSize;
@@ -217,7 +219,6 @@ export const TrainingSubsets = ({
                         newTrainingSubsetSize={newTrainingSubsetSize}
                         newValidationSubsetSize={newValidationSubsetSize}
                         newTestingSubsetSize={newTestingSubsetSize}
-                        totalDatasetItemsSize={totalDatasetItemsSize}
                     />
                 </View>
 

application/ui/src/features/models/train-model/advanced-settings/data-management/training-subsets/training-subsets.test.tsx

@@ -14,6 +14,7 @@ import { render } from 'test-utils/render';
 import { http } from '../../../../../../api/utils';
 import { TrainingConfiguration } from '../../../../../../constants/shared-types';
 import { server } from '../../../../../../msw-node-setup';
+import { distributeByLargestRemainder } from '../../../../utils';
 import { TrainingSubsets } from './training-subsets.component';
 import { getSubsetSplitParameters, SubsetSplitParameters } from './utils';
 
@@ -108,9 +109,10 @@ const expectTrainingSubsetsDistribution = async ({
     validationSubset: number;
     testSubset: number;
 }) => {
-    const newTestingSize = Math.floor(unassignedSize * (testSubset / 100));
-    const newValidationSize = Math.floor(unassignedSize * (validationSubset / 100));
-    const newTrainingSize = unassignedSize - newValidationSize - newTestingSize;
+    const [newTrainingSize, newValidationSize, newTestingSize] = distributeByLargestRemainder(
+        [trainingSubset, validationSubset, testSubset],
+        unassignedSize
+    );
 
     expectTrainingSubsetsDistributionProportion({
         validationSubset,
@@ -393,4 +395,50 @@ describe('TrainingSubsets', () => {
         expect(alert).toBeInTheDocument();
         expect(within(alert).getByRole('heading')).toHaveTextContent('Invalid training subsets configuration');
     });
+
+    it('45/28/27 split on 5 unassigned items distributes as Training=2, Validation=2, Test=1', async () => {
+        const trainingSize = 0;
+        const validationSize = 0;
+        const testSize = 0;
+        const unassignedSize = 5;
+
+        mockSubsetsNetworkRequest({ trainingSize, validationSize, unassignedSize, testSize });
+
+        const regressionSubsetParameters: SubsetSplitParameters = [
+            getMockedConfigurationParameter({
+                key: 'training',
+                value_type: 'int',
+                name: 'Training percentage',
+                value: 45,
+                description: 'Percentage of data to use for training',
+                default_value: 45,
+                max_value: 100,
+                min_value: 1,
+            }),
+            getMockedConfigurationParameter({
+                key: 'validation',
+                value_type: 'int',
+                name: 'Validation percentage',
+                value: 28,
+                description: 'Percentage of data to use for validation',
+                default_value: 28,
+                max_value: 100,
+                min_value: 1,
+            }),
+            getMockedConfigurationParameter({
+                key: 'test',
+                value_type: 'int',
+                name: 'Test percentage',
+                value: 27,
+                description: 'Percentage of data to use for testing',
+                default_value: 27,
+                max_value: 100,
+                min_value: 1,
+            }),
+        ];
+
+        render(<App subsetParameters={regressionSubsetParameters} />);
+
+        await expectSubsetSizes({ trainingSize: 2, validationSize: 2, testSize: 1 });
+    });
 });

application/ui/src/features/models/utils.test.ts

@@ -4,7 +4,52 @@
 import { getMockedModel } from 'mocks/mock-model';
 import { getMockedVariant } from 'mocks/mock-model-variant';
 
-import { getAllModelsWithOpenVINOVariants, getModelIdentifierPayload, SelectableModel } from './utils';
+import {
+    distributeByLargestRemainder,
+    getAllModelsWithOpenVINOVariants,
+    getModelIdentifierPayload,
+    SelectableModel,
+} from './utils';
+
+describe('distributeByLargestRemainder', () => {
+    it('returns all zeros when total is zero', () => {
+        expect(distributeByLargestRemainder([70, 20, 10], 0)).toEqual([0, 0, 0]);
+    });
+
+    it('returns all zeros when sum of values is zero', () => {
+        expect(distributeByLargestRemainder([0, 0, 0], 100)).toEqual([0, 0, 0]);
+    });
+
+    it('returns empty array for empty input with non-zero total', () => {
+        expect(distributeByLargestRemainder([], 100)).toEqual([]);
+    });
+
+    it('result always sums to total', () => {
+        const result = distributeByLargestRemainder([750, 125, 125], 100);
+        expect(result.reduce((a, b) => a + b, 0)).toBe(100);
+        expect(result).toEqual([75, 13, 12]);
+    });
+
+    it('tie-breaking: largest fractional part gets the extra unit', () => {
+        // [1, 1, 1] → each gets 33.33%, floors to 33, remainder = 1
+        // all fractionals equal → first in sorted order wins
+        const result = distributeByLargestRemainder([1, 1, 1], 100);
+        expect(result.reduce((a, b) => a + b, 0)).toBe(100);
+        expect(result).toEqual([34, 33, 33]);
+    });
+
+    it('count case: 45/28/27 split on 5 items → [2, 2, 1]', () => {
+        expect(distributeByLargestRemainder([45, 28, 27], 5)).toEqual([2, 2, 1]);
+    });
+
+    it('standard percentage case: 70/20/10 on 100 → [70, 20, 10]', () => {
+        expect(distributeByLargestRemainder([70, 20, 10], 100)).toEqual([70, 20, 10]);
+    });
+
+    it('handles a single value', () => {
+        expect(distributeByLargestRemainder([5], 100)).toEqual([100]);
+    });
+});
 
 describe('getAllModelsWithOpenVINOVariants', () => {
     it('returns empty array for empty models array', () => {

application/ui/src/features/models/utils.ts

@@ -10,6 +10,31 @@ export const getModelIdentifierPayload = (model: SelectableModel): { model_id: s
     model_variant_id: model.modelVariantId,
 });
 
+export const distributeByLargestRemainder = (values: number[], total: number): number[] => {
+    const sum = values.reduce((acc, value) => acc + value, 0);
+
+    if (sum <= 0 || total <= 0) {
+        return values.map(() => 0);
+    }
+
+    const exactShares = values.map((value) => (value / sum) * total);
+    const flooredShares = exactShares.map((share) => Math.floor(share));
+    let remainder = total - flooredShares.reduce((acc, value) => acc + value, 0);
+
+    const indicesByRemainder = exactShares
+        .map((share, index) => ({ index, fractional: share - Math.floor(share) }))
+        .sort((a, b) => b.fractional - a.fractional);
+
+    const result = [...flooredShares];
+    for (const { index } of indicesByRemainder) {
+        if (remainder <= 0) break;
+        result[index] += 1;
+        remainder -= 1;
+    }
+
+    return result;
+};
+
 export const getAllModelsWithOpenVINOVariants = (models: Model[]): SelectableModel[] => {
     return models.flatMap((model) =>
         model.variants