From c33e58aef58f3259818e64e6391674d258b82a5b Mon Sep 17 00:00:00 2001 From: Robert Romero Date: Thu, 7 Aug 2025 13:33:17 -0700 Subject: [PATCH] Add cluster core count setting for cost recovery --- README.md | 2 +- src/rates.json | 1 + src/slurmcostmanager.js | 31 ++++++++++++++++++++++++++++++- src/slurmdb.py | 23 +++++++++++++++++++++++ test/unit/billing_summary.test.py | 24 ++++++++++++++++++++++++ 5 files changed, 79 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 8c0d081..87faf82 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ This repository now includes a responsive Cockpit UI built with React. The inte - **Detailed cost drill-downs** (core‑hours, GPU-hours) for per‑account transparency. - **Historical billing data** accessible from account inception for auditing and trend analysis. - **Organization-wide views** consolidating charges across all member Slurm accounts. -- **Configurable rate table** with per-account overrides, editable from a dedicated Settings tab. +- **Configurable rate table** with per-account overrides and cluster core count, editable from a dedicated Settings tab. ## 📁 Project Structure diff --git a/src/rates.json b/src/rates.json index 48665f5..1ebfac3 100644 --- a/src/rates.json +++ b/src/rates.json @@ -1,6 +1,7 @@ { "defaultRate": 0.02, "defaultGpuRate": 0.2, + "clusterCores": 100, "historicalRates": { "2024-01": 0.015 }, diff --git a/src/slurmcostmanager.js b/src/slurmcostmanager.js index 62b2331..553748a 100644 --- a/src/slurmcostmanager.js +++ b/src/slurmcostmanager.js @@ -866,7 +866,10 @@ function Rates({ onRatesUpdated }) { } if (cancelled) return; const json = JSON.parse(text); - setConfig({ defaultRate: json.defaultRate }); + setConfig({ + defaultRate: json.defaultRate, + clusterCores: json.clusterCores || '' + }); const ovrs = json.overrides ? Object.entries(json.overrides).map(([account, cfg]) => ({ account, @@ -946,6 +949,16 @@ function Rates({ onRatesUpdated }) { const json = { defaultRate }; + if (config.clusterCores !== undefined && config.clusterCores !== '') { + const cores = parseInt(config.clusterCores, 10); + if (!Number.isFinite(cores) || cores < 0) { + console.warn('Invalid cluster core count:', config.clusterCores); + setError('Invalid cluster core count'); + return; + } + json.clusterCores = cores; + } + if (overrides.length) { const overridesJson = {}; overrides.forEach(o => { @@ -1020,6 +1033,22 @@ function Rates({ onRatesUpdated }) { }) ) ), + React.createElement( + 'div', + null, + React.createElement( + 'label', + null, + 'Total Cluster Cores: ', + React.createElement('input', { + type: 'number', + step: '1', + value: config.clusterCores, + onChange: e => + setConfig({ ...config, clusterCores: e.target.value }) + }) + ) + ), React.createElement('h3', null, 'Account Overrides'), React.createElement( 'table', diff --git a/src/slurmdb.py b/src/slurmdb.py index b9a3b6a..cd2e16e 100644 --- a/src/slurmdb.py +++ b/src/slurmdb.py @@ -5,6 +5,7 @@ import logging import sys from datetime import date, datetime, timedelta +from calendar import monthrange try: import pymysql @@ -552,6 +553,7 @@ def export_summary(self, start_time, end_time): overrides = rates_cfg.get('overrides', {}) historical = rates_cfg.get('historicalRates', {}) gpu_historical = rates_cfg.get('historicalGpuRates', {}) + cluster_cores = rates_cfg.get('clusterCores') for month, accounts in usage.items(): base_rate = historical.get(month, default_rate) @@ -637,6 +639,27 @@ def export_summary(self, start_time, end_time): 'core_hours': round(total_ch, 2), 'gpu_hours': round(total_gpu, 2), } + if cluster_cores: + start_date = start_dt.date() + end_date = end_dt.date() + current = date(start_date.year, start_date.month, 1) + end_marker = date(end_date.year, end_date.month, 1) + projected_revenue = 0.0 + while current <= end_marker: + days_in_month = monthrange(current.year, current.month)[1] + month_start = date(current.year, current.month, 1) + month_end = date(current.year, current.month, days_in_month) + overlap_start = max(month_start, start_date) + overlap_end = min(month_end, end_date) + if overlap_start <= overlap_end: + days = (overlap_end - overlap_start).days + 1 + rate = historical.get(current.strftime('%Y-%m'), default_rate) + projected_revenue += cluster_cores * 24 * days * rate + if current.month == 12: + current = date(current.year + 1, 1, 1) + else: + current = date(current.year, current.month + 1, 1) + summary['summary']['projected_revenue'] = round(projected_revenue, 2) summary['daily'] = [ { 'date': d, diff --git a/test/unit/billing_summary.test.py b/test/unit/billing_summary.test.py index c1f8e55..799bb07 100644 --- a/test/unit/billing_summary.test.py +++ b/test/unit/billing_summary.test.py @@ -186,6 +186,30 @@ def fake_open(path, *args, **kwargs): with self.assertRaises(ValueError): db.export_summary('2023-10-01', '2023-10-31') + def test_export_summary_projected_revenue(self): + usage = { + '2024-02': { + 'acct': {'core_hours': 10.0, 'users': {}} + } + } + + def fake_open(path, *args, **kwargs): + if path.endswith('rates.json'): + return io.StringIO('{"defaultRate": 0.02, "clusterCores": 100}') + return open_orig(path, *args, **kwargs) + + open_orig = open + with mock.patch.object( + SlurmDB, + 'aggregate_usage', + return_value=(usage, {'daily': {}, 'monthly': {}, 'yearly': {}}), + ), mock.patch.object(SlurmDB, 'fetch_invoices', return_value=[]), mock.patch( + 'builtins.open', side_effect=fake_open + ): + db = SlurmDB() + summary = db.export_summary('2024-02-01', '2024-02-29') + self.assertAlmostEqual(summary['summary']['projected_revenue'], 1392.0) + if __name__ == '__main__': unittest.main()