diff --git a/src/cost-calculator.js b/src/cost-calculator.js index c7300a5..ad8f3c9 100644 --- a/src/cost-calculator.js +++ b/src/cost-calculator.js @@ -29,8 +29,9 @@ function loadRatesConfig() { /** * Calculate charges from usage records applying rates and overrides. * - * @param {Array} usage - Array of {account, date, core_hours} - * @param {Object} config - Configuration with defaultRate, historicalRates, overrides + * @param {Array} usage - Array of {account, date, core_hours, gpu_hours} + * @param {Object} config - Configuration with defaultRate, defaultGpuRate, + * historicalRates, historicalGpuRates, overrides * @returns {Object} charges grouped by month then account */ function calculateCharges(usage, config) { @@ -41,23 +42,37 @@ function calculateCharges(usage, config) { typeof config.defaultRate === 'number' && config.defaultRate > 0 ? config.defaultRate : 0; + const defaultGpuRate = + typeof config.defaultGpuRate === 'number' && config.defaultGpuRate > 0 + ? config.defaultGpuRate + : 0; const historical = config.historicalRates || {}; + const gpuHistorical = config.historicalGpuRates || {}; const overrides = config.overrides || {}; const charges = {}; for (const record of usage) { - if (!record || typeof record.core_hours !== 'number' || record.core_hours <= 0) { + if (!record) { continue; } const account = record.account || 'unknown'; const month = (record.date || '').slice(0, 7); // YYYY-MM const ovr = overrides[account] || {}; + const coreHours = typeof record.core_hours === 'number' && record.core_hours > 0 ? record.core_hours : 0; + const gpuHours = typeof record.gpu_hours === 'number' && record.gpu_hours > 0 ? record.gpu_hours : 0; + if (coreHours <= 0 && gpuHours <= 0) { + continue; + } const rate = typeof ovr.rate === 'number' ? ovr.rate : (typeof historical[month] === 'number' ? historical[month] : defaultRate); + const gpuRate = typeof ovr.gpuRate === 'number' + ? ovr.gpuRate + : (typeof gpuHistorical[month] === 'number' ? gpuHistorical[month] : defaultGpuRate); const validRate = rate > 0 ? rate : 0; - let cost = record.core_hours * validRate; + const validGpuRate = gpuRate > 0 ? gpuRate : 0; + let cost = coreHours * validRate + gpuHours * validGpuRate; const rawDiscount = typeof ovr.discount === 'number' ? ovr.discount : 0; const discount = Math.min(1, Math.max(0, rawDiscount)); if (discount > 0) { @@ -66,9 +81,10 @@ function calculateCharges(usage, config) { if (!charges[month]) charges[month] = {}; if (!charges[month][account]) { - charges[month][account] = { core_hours: 0, cost: 0 }; + charges[month][account] = { core_hours: 0, gpu_hours: 0, cost: 0 }; } - charges[month][account].core_hours += record.core_hours; + charges[month][account].core_hours += coreHours; + charges[month][account].gpu_hours += gpuHours; charges[month][account].cost += cost; } @@ -76,6 +92,7 @@ function calculateCharges(usage, config) { for (const account of Object.keys(charges[month])) { const entry = charges[month][account]; entry.core_hours = Number(entry.core_hours.toFixed(2)); + entry.gpu_hours = Number(entry.gpu_hours.toFixed(2)); entry.cost = Number(entry.cost.toFixed(2)); } } diff --git a/src/rates-schema.json b/src/rates-schema.json index 9c06cbb..0a4458f 100644 --- a/src/rates-schema.json +++ b/src/rates-schema.json @@ -1,13 +1,17 @@ { "$schema": "http://json-schema.org/draft-07/schema#", "title": "RatesConfig", - "description": "Schema for core hour rate configuration", + "description": "Schema for core and GPU hour rate configuration", "type": "object", "properties": { "defaultRate": { "type": "number", "minimum": 0 }, + "defaultGpuRate": { + "type": "number", + "minimum": 0 + }, "historicalRates": { "type": "object", "patternProperties": { @@ -18,6 +22,16 @@ }, "additionalProperties": false }, + "historicalGpuRates": { + "type": "object", + "patternProperties": { + "^[0-9]{4}-(0[1-9]|1[0-2])$": { + "type": "number", + "minimum": 0 + } + }, + "additionalProperties": false + }, "overrides": { "type": "object", "patternProperties": { @@ -28,6 +42,10 @@ "type": "number", "minimum": 0 }, + "gpuRate": { + "type": "number", + "minimum": 0 + }, "discount": { "type": "number", "minimum": 0, diff --git a/src/rates.json b/src/rates.json index c1873e2..48665f5 100644 --- a/src/rates.json +++ b/src/rates.json @@ -1,17 +1,23 @@ { "defaultRate": 0.02, + "defaultGpuRate": 0.2, "historicalRates": { "2024-01": 0.015 }, + "historicalGpuRates": { + "2024-01": 0.15 + }, "overrides": { "research": { - "rate": 0.01 + "rate": 0.01, + "gpuRate": 0.1 }, "education": { "discount": 0.5 }, "special": { "rate": 0.025, + "gpuRate": 0.25, "discount": 0.1 } } diff --git a/src/slurmdb.py b/src/slurmdb.py index 5398768..9377483 100644 --- a/src/slurmdb.py +++ b/src/slurmdb.py @@ -381,24 +381,32 @@ def export_summary(self, start_time, end_time): logging.error("Failed to parse rates file %s: %s", rates_path, e) raise default_rate = rates_cfg.get('defaultRate', 0.01) + default_gpu_rate = rates_cfg.get('defaultGpuRate', 0.0) overrides = rates_cfg.get('overrides', {}) historical = rates_cfg.get('historicalRates', {}) + gpu_historical = rates_cfg.get('historicalGpuRates', {}) for month, accounts in usage.items(): base_rate = historical.get(month, default_rate) + base_gpu_rate = gpu_historical.get(month, default_gpu_rate) for account, vals in accounts.items(): ovr = overrides.get(account, {}) rate = ovr.get('rate', base_rate) + gpu_rate = ovr.get('gpuRate', base_gpu_rate) discount = ovr.get('discount', 0) if rate < 0: raise ValueError(f"Invalid rate {rate} for account {account}") + if gpu_rate < 0: + raise ValueError( + f"Invalid GPU rate {gpu_rate} for account {account}" + ) if not 0 <= discount <= 1: raise ValueError( f"Invalid discount {discount} for account {account}" ) - acct_cost = vals['core_hours'] * rate + acct_cost = vals['core_hours'] * rate + vals.get('gpu_hours', 0.0) * gpu_rate if 0 < discount < 1: acct_cost *= 1 - discount users = [] diff --git a/test/unit/billing_summary.test.py b/test/unit/billing_summary.test.py index 5319d75..6b16937 100644 --- a/test/unit/billing_summary.test.py +++ b/test/unit/billing_summary.test.py @@ -33,11 +33,11 @@ def test_export_summary_aggregates_costs(self): with mock.patch.object(SlurmDB, 'fetch_invoices', return_value=invoices): db = SlurmDB() summary = db.export_summary('2023-10-01', '2023-10-31') - self.assertEqual(summary['summary']['total'], 0.2) + self.assertEqual(summary['summary']['total'], 1.2) self.assertEqual(summary['details'][0]['account'], 'acct') self.assertEqual(summary['details'][0]['core_hours'], 10.0) self.assertEqual(summary['details'][0]['gpu_hours'], 5.0) - self.assertEqual(summary['details'][0]['cost'], 0.2) + self.assertEqual(summary['details'][0]['cost'], 1.2) self.assertEqual(summary['summary']['gpu_hours'], 5.0) self.assertEqual(summary['invoices'][0]['file'], 'inv1.pdf') diff --git a/test/unit/calculator.test.js b/test/unit/calculator.test.js index 0847fdf..2f308da 100644 --- a/test/unit/calculator.test.js +++ b/test/unit/calculator.test.js @@ -5,31 +5,38 @@ const { calculateCharges, loadRatesConfig } = require('../../src/cost-calculator function testFileConfig() { const usage = [ - { account: 'education', date: '2024-01-15', core_hours: 100 }, - { account: 'research', date: '2024-02-01', core_hours: 50 }, - { account: 'special', date: '2024-02-01', core_hours: 100 }, + { account: 'education', date: '2024-01-15', core_hours: 100, gpu_hours: 10 }, + { account: 'research', date: '2024-02-01', core_hours: 50, gpu_hours: 5 }, + { account: 'special', date: '2024-02-01', core_hours: 100, gpu_hours: 20 }, { account: 'other', date: '2024-02-01', core_hours: 10 } ]; const config = loadRatesConfig(); const charges = calculateCharges(usage, config); - assert.strictEqual(charges['2024-01'].education.cost, 100 * 0.015 * 0.5); - assert.strictEqual(charges['2024-02'].research.cost, 50 * 0.01); - assert.strictEqual(charges['2024-02'].special.cost, 100 * 0.025 * 0.9); + assert.strictEqual(charges['2024-01'].education.cost, (100 * 0.015 + 10 * 0.15) * 0.5); + assert.strictEqual(charges['2024-02'].research.cost, 50 * 0.01 + 5 * 0.1); + assert.strictEqual(charges['2024-02'].special.cost, (100 * 0.025 + 20 * 0.25) * 0.9); assert.strictEqual(charges['2024-02'].other.cost, 10 * 0.02); + assert.strictEqual(charges['2024-01'].education.gpu_hours, 10); } function testPassedConfig() { const usage = [ - { account: 'acct', date: '2024-03-01', core_hours: 100 } + { account: 'acct', date: '2024-03-01', core_hours: 100, gpu_hours: 10 } ]; - const config = { defaultRate: 0.02, historicalRates: { '2024-03': 0.03 }, overrides: { acct: { discount: 0.25 } } }; + const config = { + defaultRate: 0.02, + defaultGpuRate: 0.2, + historicalRates: { '2024-03': 0.03 }, + historicalGpuRates: { '2024-03': 0.3 }, + overrides: { acct: { discount: 0.25 } } + }; const charges = calculateCharges(usage, config); - assert.strictEqual(charges['2024-03'].acct.cost, 100 * 0.03 * 0.75); + assert.strictEqual(charges['2024-03'].acct.cost, (100 * 0.03 + 10 * 0.3) * 0.75); } function testInvalidUsageIgnored() { const usage = [ - { account: 'bad', date: '2024-04-01', core_hours: 'NaN' } + { account: 'bad', date: '2024-04-01', core_hours: 'NaN', gpu_hours: 'NaN' } ]; const charges = calculateCharges(usage, { defaultRate: 0.01 }); assert.deepStrictEqual(charges, {}); @@ -77,10 +84,10 @@ function testNegativeInputs() { const charges = calculateCharges(usage, config); const may = charges['2024-05']; assert.ok(!('negHours' in may)); - assert.deepStrictEqual(may.negRate, { core_hours: 5, cost: 0 }); - assert.deepStrictEqual(may.discLow, { core_hours: 5, cost: 0.5 }); - assert.deepStrictEqual(may.discHigh, { core_hours: 5, cost: 0 }); - assert.deepStrictEqual(may.def, { core_hours: 5, cost: 0 }); + assert.deepStrictEqual(may.negRate, { core_hours: 5, gpu_hours: 0, cost: 0 }); + assert.deepStrictEqual(may.discLow, { core_hours: 5, gpu_hours: 0, cost: 0.5 }); + assert.deepStrictEqual(may.discHigh, { core_hours: 5, gpu_hours: 0, cost: 0 }); + assert.deepStrictEqual(may.def, { core_hours: 5, gpu_hours: 0, cost: 0 }); } function testRoundingTotals() { @@ -91,6 +98,7 @@ function testRoundingTotals() { const charges = calculateCharges(usage, { defaultRate: 0.333 }); const june = charges['2024-06']; assert.strictEqual(june.round.core_hours, 0.67); + assert.strictEqual(june.round.gpu_hours, 0); assert.strictEqual(june.round.cost, 0.22); }