Spaces:
Running
Running
GuillaumeSalouHF HF Staff
feat: highlight capacity observed within the past ten days
900a11b verified | import json | |
| import queue | |
| import tempfile | |
| import unittest | |
| from pathlib import Path | |
| from unittest.mock import patch | |
| from gpu_capacity import history, main | |
| class RecentAvailabilityTests(unittest.TestCase): | |
| def setUp(self): | |
| self.now = 1791417600 # 2026-10-08 00:00 UTC | |
| self.cutoff = self.now - 10 * 86400 | |
| for name in ("_hist", "_last_available"): | |
| mock = patch.dict(getattr(history, name), {}, clear=True) | |
| mock.start() | |
| self.addCleanup(mock.stop) | |
| for name, value in (("_queue", queue.Queue()), ("_today", []), ("_today_day", [""])): | |
| mock = patch.object(history, name, value) | |
| mock.start() | |
| self.addCleanup(mock.stop) | |
| clock = patch.object(history.time, "time", return_value=self.now) | |
| clock.start() | |
| self.addCleanup(clock.stop) | |
| def result(self, ts, status="available", available=True): | |
| return { | |
| "checked_at": ts, | |
| "status": status, | |
| "azs": {"us-east-1a": {"status": "available" if available else "unavailable"}}, | |
| } | |
| def test_positive_observation_survives_failures_and_out_of_order_results(self): | |
| for result in ( | |
| self.result(self.now - 100), | |
| self.result(self.now - 50, "unavailable", False), | |
| self.result(self.now - 200), | |
| self.result(self.now - 25, "quota", False), | |
| self.result(self.now, "available", False), | |
| ): | |
| history.record("g7e.2xlarge", {"us-east-1": result}) | |
| self.assertEqual( | |
| history.recent_availability(10), | |
| {"g7e.2xlarge": {"us-east-1": self.now - 100}}, | |
| ) | |
| def test_window_boundary_and_expiration_without_new_probes(self): | |
| history._last_available.update({ | |
| ("g7e.2xlarge", "us-east-1"): self.cutoff, | |
| ("g7e.2xlarge", "us-east-2"): self.cutoff - 1, | |
| ("g7e.48xlarge", "us-east-1"): self.now + 1, | |
| }) | |
| self.assertEqual( | |
| history.recent_availability(10), | |
| {"g7e.2xlarge": {"us-east-1": self.cutoff}}, | |
| ) | |
| with patch.object(history.time, "time", return_value=self.now + 0.5): | |
| self.assertEqual(history.recent_availability(10), {}) | |
| def test_restore_history_finds_latest_success_not_latest_probe(self): | |
| observations = [ | |
| (self.now - 86400, 1, "available"), | |
| (self.now - 2 * 86400, 2, "available"), | |
| (self.now - 100, 0, "unavailable"), | |
| (self.now - 50, 0, "available"), | |
| (self.now, 1, "quota"), | |
| ] | |
| with tempfile.TemporaryDirectory() as directory: | |
| path = Path(directory) / "2026-10-07.jsonl" | |
| path.write_text("\n".join(json.dumps({ | |
| "t": ts, "a": available, "n": 2, "s": status, | |
| "type": "g7e.2xlarge", "region": "us-east-1", | |
| }) for ts, available, status in observations)) | |
| with patch.object(history.config, "HISTORY_DIR", directory): | |
| history.load() | |
| self.assertEqual( | |
| history.recent_availability(10), | |
| {"g7e.2xlarge": {"us-east-1": self.now - 86400}}, | |
| ) | |
| def test_state_api_exposes_history_without_changing_current_status(self): | |
| history.record("g7e.2xlarge", {"us-east-1": self.result(self.now - 86400)}) | |
| current = {"g7e.2xlarge": {"us-east-1": self.result(self.now, "unavailable", False)}} | |
| with patch.object(main.state, "get", return_value=current), patch.object(main.prober, "gpu_info", return_value={}): | |
| response = main.get_state() | |
| self.assertEqual(response["recent_availability_days"], 10) | |
| self.assertEqual(response["recent_availability"]["g7e.2xlarge"]["us-east-1"], self.now - 86400) | |
| self.assertEqual(response["state"]["g7e.2xlarge"]["us-east-1"]["status"], "unavailable") | |
| if __name__ == "__main__": | |
| unittest.main() | |