Skip to content

Commit 5fa17ad

Browse files
NessieCanCodeNessieCanCode
authored
Merge pull request #97 from NessieCanCode/extend-cost-calculator-for-gpu-hour-billing
Add GPU-hour billing support
2 parents d4ea2aa + b73f089 commit 5fa17ad

File tree

6 files changed

+82
-25
lines changed

6 files changed

+82
-25
lines changed

‎src/cost-calculator.js‎

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,9 @@ function loadRatesConfig() {
2929
/**
3030
* Calculate charges from usage records applying rates and overrides.
3131
*
32-
* @param {Array} usage - Array of {account, date, core_hours}
33-
* @param {Object} config - Configuration with defaultRate, historicalRates, overrides
32+
* @param {Array} usage - Array of {account, date, core_hours, gpu_hours}
33+
* @param {Object} config - Configuration with defaultRate, defaultGpuRate,
34+
* historicalRates, historicalGpuRates, overrides
3435
* @returns {Object} charges grouped by month then account
3536
*/
3637
function calculateCharges(usage, config) {
@@ -41,23 +42,37 @@ function calculateCharges(usage, config) {
4142
typeof config.defaultRate === 'number' && config.defaultRate > 0
4243
? config.defaultRate
4344
: 0;
45+
const defaultGpuRate =
46+
typeof config.defaultGpuRate === 'number' && config.defaultGpuRate > 0
47+
? config.defaultGpuRate
48+
: 0;
4449
const historical = config.historicalRates || {};
50+
const gpuHistorical = config.historicalGpuRates || {};
4551
const overrides = config.overrides || {};
4652

4753
const charges = {};
4854

4955
for (const record of usage) {
50-
if (!record || typeof record.core_hours !== 'number' || record.core_hours <= 0) {
56+
if (!record) {
5157
continue;
5258
}
5359
const account = record.account || 'unknown';
5460
const month = (record.date || '').slice(0, 7); // YYYY-MM
5561
const ovr = overrides[account] || {};
62+
const coreHours = typeof record.core_hours === 'number' && record.core_hours > 0 ? record.core_hours : 0;
63+
const gpuHours = typeof record.gpu_hours === 'number' && record.gpu_hours > 0 ? record.gpu_hours : 0;
64+
if (coreHours <= 0 && gpuHours <= 0) {
65+
continue;
66+
}
5667
const rate = typeof ovr.rate === 'number'
5768
? ovr.rate
5869
: (typeof historical[month] === 'number' ? historical[month] : defaultRate);
70+
const gpuRate = typeof ovr.gpuRate === 'number'
71+
? ovr.gpuRate
72+
: (typeof gpuHistorical[month] === 'number' ? gpuHistorical[month] : defaultGpuRate);
5973
const validRate = rate > 0 ? rate : 0;
60-
let cost = record.core_hours * validRate;
74+
const validGpuRate = gpuRate > 0 ? gpuRate : 0;
75+
let cost = coreHours * validRate + gpuHours * validGpuRate;
6176
const rawDiscount = typeof ovr.discount === 'number' ? ovr.discount : 0;
6277
const discount = Math.min(1, Math.max(0, rawDiscount));
6378
if (discount > 0) {
@@ -66,16 +81,18 @@ function calculateCharges(usage, config) {
6681

6782
if (!charges[month]) charges[month] = {};
6883
if (!charges[month][account]) {
69-
charges[month][account] = { core_hours: 0, cost: 0 };
84+
charges[month][account] = { core_hours: 0, gpu_hours: 0, cost: 0 };
7085
}
71-
charges[month][account].core_hours += record.core_hours;
86+
charges[month][account].core_hours += coreHours;
87+
charges[month][account].gpu_hours += gpuHours;
7288
charges[month][account].cost += cost;
7389
}
7490

7591
for (const month of Object.keys(charges)) {
7692
for (const account of Object.keys(charges[month])) {
7793
const entry = charges[month][account];
7894
entry.core_hours = Number(entry.core_hours.toFixed(2));
95+
entry.gpu_hours = Number(entry.gpu_hours.toFixed(2));
7996
entry.cost = Number(entry.cost.toFixed(2));
8097
}
8198
}

‎src/rates-schema.json‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,17 @@
11
{
22
"$schema": "http://json-schema.org/draft-07/schema#",
33
"title": "RatesConfig",
4-
"description": "Schema for core hour rate configuration",
4+
"description": "Schema for core and GPU hour rate configuration",
55
"type": "object",
66
"properties": {
77
"defaultRate": {
88
"type": "number",
99
"minimum": 0
1010
},
11+
"defaultGpuRate": {
12+
"type": "number",
13+
"minimum": 0
14+
},
1115
"historicalRates": {
1216
"type": "object",
1317
"patternProperties": {
@@ -18,6 +22,16 @@
1822
},
1923
"additionalProperties": false
2024
},
25+
"historicalGpuRates": {
26+
"type": "object",
27+
"patternProperties": {
28+
"^[0-9]{4}-(0[1-9]|1[0-2])$": {
29+
"type": "number",
30+
"minimum": 0
31+
}
32+
},
33+
"additionalProperties": false
34+
},
2135
"overrides": {
2236
"type": "object",
2337
"patternProperties": {
@@ -28,6 +42,10 @@
2842
"type": "number",
2943
"minimum": 0
3044
},
45+
"gpuRate": {
46+
"type": "number",
47+
"minimum": 0
48+
},
3149
"discount": {
3250
"type": "number",
3351
"minimum": 0,

‎src/rates.json‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,23 @@
11
{
22
"defaultRate": 0.02,
3+
"defaultGpuRate": 0.2,
34
"historicalRates": {
45
"2024-01": 0.015
56
},
7+
"historicalGpuRates": {
8+
"2024-01": 0.15
9+
},
610
"overrides": {
711
"research": {
8-
"rate": 0.01
12+
"rate": 0.01,
13+
"gpuRate": 0.1
914
},
1015
"education": {
1116
"discount": 0.5
1217
},
1318
"special": {
1419
"rate": 0.025,
20+
"gpuRate": 0.25,
1521
"discount": 0.1
1622
}
1723
}

‎src/slurmdb.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -393,24 +393,32 @@ def export_summary(self, start_time, end_time):
393393
logging.error("Failed to parse rates file %s: %s", rates_path, e)
394394
raise
395395
default_rate = rates_cfg.get('defaultRate', 0.01)
396+
default_gpu_rate = rates_cfg.get('defaultGpuRate', 0.0)
396397
overrides = rates_cfg.get('overrides', {})
397398
historical = rates_cfg.get('historicalRates', {})
399+
gpu_historical = rates_cfg.get('historicalGpuRates', {})
398400

399401
for month, accounts in usage.items():
400402
base_rate = historical.get(month, default_rate)
403+
base_gpu_rate = gpu_historical.get(month, default_gpu_rate)
401404
for account, vals in accounts.items():
402405
ovr = overrides.get(account, {})
403406
rate = ovr.get('rate', base_rate)
407+
gpu_rate = ovr.get('gpuRate', base_gpu_rate)
404408
discount = ovr.get('discount', 0)
405409

406410
if rate < 0:
407411
raise ValueError(f"Invalid rate {rate} for account {account}")
412+
if gpu_rate < 0:
413+
raise ValueError(
414+
f"Invalid GPU rate {gpu_rate} for account {account}"
415+
)
408416
if not 0 <= discount <= 1:
409417
raise ValueError(
410418
f"Invalid discount {discount} for account {account}"
411419
)
412420

413-
acct_cost = vals['core_hours'] * rate
421+
acct_cost = vals['core_hours'] * rate + vals.get('gpu_hours', 0.0) * gpu_rate
414422
if 0 < discount < 1:
415423
acct_cost *= 1 - discount
416424
users = []

‎test/unit/billing_summary.test.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,11 +33,11 @@ def test_export_summary_aggregates_costs(self):
3333
with mock.patch.object(SlurmDB, 'fetch_invoices', return_value=invoices):
3434
db = SlurmDB()
3535
summary = db.export_summary('2023-10-01', '2023-10-31')
36-
self.assertEqual(summary['summary']['total'], 0.2)
36+
self.assertEqual(summary['summary']['total'], 1.2)
3737
self.assertEqual(summary['details'][0]['account'], 'acct')
3838
self.assertEqual(summary['details'][0]['core_hours'], 10.0)
3939
self.assertEqual(summary['details'][0]['gpu_hours'], 5.0)
40-
self.assertEqual(summary['details'][0]['cost'], 0.2)
40+
self.assertEqual(summary['details'][0]['cost'], 1.2)
4141
self.assertEqual(summary['summary']['gpu_hours'], 5.0)
4242
self.assertEqual(summary['invoices'][0]['file'], 'inv1.pdf')
4343

‎test/unit/calculator.test.js‎

Lines changed: 22 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -5,31 +5,38 @@ const { calculateCharges, loadRatesConfig } = require('../../src/cost-calculator
55

66
function testFileConfig() {
77
const usage = [
8-
{ account: 'education', date: '2024-01-15', core_hours: 100 },
9-
{ account: 'research', date: '2024-02-01', core_hours: 50 },
10-
{ account: 'special', date: '2024-02-01', core_hours: 100 },
8+
{ account: 'education', date: '2024-01-15', core_hours: 100, gpu_hours: 10 },
9+
{ account: 'research', date: '2024-02-01', core_hours: 50, gpu_hours: 5 },
10+
{ account: 'special', date: '2024-02-01', core_hours: 100, gpu_hours: 20 },
1111
{ account: 'other', date: '2024-02-01', core_hours: 10 }
1212
];
1313
const config = loadRatesConfig();
1414
const charges = calculateCharges(usage, config);
15-
assert.strictEqual(charges['2024-01'].education.cost, 100 * 0.015 * 0.5);
16-
assert.strictEqual(charges['2024-02'].research.cost, 50 * 0.01);
17-
assert.strictEqual(charges['2024-02'].special.cost, 100 * 0.025 * 0.9);
15+
assert.strictEqual(charges['2024-01'].education.cost, (100 * 0.015 + 10 * 0.15) * 0.5);
16+
assert.strictEqual(charges['2024-02'].research.cost, 50 * 0.01 + 5 * 0.1);
17+
assert.strictEqual(charges['2024-02'].special.cost, (100 * 0.025 + 20 * 0.25) * 0.9);
1818
assert.strictEqual(charges['2024-02'].other.cost, 10 * 0.02);
19+
assert.strictEqual(charges['2024-01'].education.gpu_hours, 10);
1920
}
2021

2122
function testPassedConfig() {
2223
const usage = [
23-
{ account: 'acct', date: '2024-03-01', core_hours: 100 }
24+
{ account: 'acct', date: '2024-03-01', core_hours: 100, gpu_hours: 10 }
2425
];
25-
const config = { defaultRate: 0.02, historicalRates: { '2024-03': 0.03 }, overrides: { acct: { discount: 0.25 } } };
26+
const config = {
27+
defaultRate: 0.02,
28+
defaultGpuRate: 0.2,
29+
historicalRates: { '2024-03': 0.03 },
30+
historicalGpuRates: { '2024-03': 0.3 },
31+
overrides: { acct: { discount: 0.25 } }
32+
};
2633
const charges = calculateCharges(usage, config);
27-
assert.strictEqual(charges['2024-03'].acct.cost, 100 * 0.03 * 0.75);
34+
assert.strictEqual(charges['2024-03'].acct.cost, (100 * 0.03 + 10 * 0.3) * 0.75);
2835
}
2936

3037
function testInvalidUsageIgnored() {
3138
const usage = [
32-
{ account: 'bad', date: '2024-04-01', core_hours: 'NaN' }
39+
{ account: 'bad', date: '2024-04-01', core_hours: 'NaN', gpu_hours: 'NaN' }
3340
];
3441
const charges = calculateCharges(usage, { defaultRate: 0.01 });
3542
assert.deepStrictEqual(charges, {});
@@ -77,10 +84,10 @@ function testNegativeInputs() {
7784
const charges = calculateCharges(usage, config);
7885
const may = charges['2024-05'];
7986
assert.ok(!('negHours' in may));
80-
assert.deepStrictEqual(may.negRate, { core_hours: 5, cost: 0 });
81-
assert.deepStrictEqual(may.discLow, { core_hours: 5, cost: 0.5 });
82-
assert.deepStrictEqual(may.discHigh, { core_hours: 5, cost: 0 });
83-
assert.deepStrictEqual(may.def, { core_hours: 5, cost: 0 });
87+
assert.deepStrictEqual(may.negRate, { core_hours: 5, gpu_hours: 0, cost: 0 });
88+
assert.deepStrictEqual(may.discLow, { core_hours: 5, gpu_hours: 0, cost: 0.5 });
89+
assert.deepStrictEqual(may.discHigh, { core_hours: 5, gpu_hours: 0, cost: 0 });
90+
assert.deepStrictEqual(may.def, { core_hours: 5, gpu_hours: 0, cost: 0 });
8491
}
8592

8693
function testRoundingTotals() {
@@ -91,6 +98,7 @@ function testRoundingTotals() {
9198
const charges = calculateCharges(usage, { defaultRate: 0.333 });
9299
const june = charges['2024-06'];
93100
assert.strictEqual(june.round.core_hours, 0.67);
101+
assert.strictEqual(june.round.gpu_hours, 0);
94102
assert.strictEqual(june.round.cost, 0.22);
95103
}
96104

0 commit comments

Comments
 (0)