fix(slurm): fail over after GCE capacity errors#5909
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a circuit breaker mechanism for GCE nodesets to mitigate issues caused by capacity stockouts. By automatically draining and probing nodesets that encounter specific capacity errors, the controller prevents Slurm from repeatedly attempting to schedule work on unavailable machine families, improving overall cluster reliability and scheduling efficiency. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an automatic capacity circuit breaker feature to the Slurm GCP v6 controller, which drains affected nodesets after GCE capacity errors and uses a small probe set to retry after an exponential cooldown. The implementation includes a new capacity_circuit.py script, integration with resume.py and slurmsync.py, and corresponding unit tests. The review feedback highlights critical performance and robustness improvements: implementing a process-level lazy cache for owns_node to avoid O(N^2) disk I/O operations, invalidating this cache on writes, clearing it in test fixtures to prevent test pollution, and defensively guarding against potential None values in partition configurations and API error responses.
| def owns_node(node: str, lkp: Optional[util.Lookup] = None) -> bool: | ||
| lkp = lkp or util.lookup() | ||
| if not STATE_FILE.exists(): | ||
| return False | ||
| try: | ||
| nodeset_name = lkp.node_nodeset_name(node) | ||
| except Exception: | ||
| return False | ||
| with _locked_state(write=False) as state: | ||
| record = state["nodesets"].get(nodeset_name) | ||
| return bool(record and node in record.get("target_nodes", [])) |
There was a problem hiding this comment.
owns_node is called for every node in slurmsync.py during sync_instances(). Without caching, this results in
def owns_node(node: str, lkp: Optional[util.Lookup] = None) -> bool:
global _owned_nodes_cache
if _owned_nodes_cache is None:
_owned_nodes_cache = set()
if STATE_FILE.exists():
with _locked_state(write=False) as state:
for record in state.get("nodesets", {}).values():
_owned_nodes_cache.update(record.get("target_nodes", []))
return node in _owned_nodes_cacheThere was a problem hiding this comment.
[codex] Fixed in 5dbbb59. owns_node now parses the circuit state once per file version and does O(1) membership checks thereafter. The cache is partitioned by nodeset to preserve the existing lookup semantics and is keyed by the atomically replaced file identity/metadata so writes from a separate resume process are detected.
| LOCK_FILE = util.slurmdirs.state / "capacity_circuits.lock" | ||
| TRANSIENT_INSTANCE_STATES = frozenset({"PROVISIONING", "STAGING"}) | ||
| TRANSIENT_SLURM_BASE_STATES = frozenset({"ALLOCATED", "MIXED"}) | ||
| TRANSIENT_SLURM_FLAGS = frozenset({"CONFIGURING", "POWERING_UP"}) |
There was a problem hiding this comment.
Define _owned_nodes_cache at the module level to support lazy caching of owned nodes and prevent redundant disk I/O.
| TRANSIENT_SLURM_FLAGS = frozenset({"CONFIGURING", "POWERING_UP"}) | |
| TRANSIENT_SLURM_FLAGS = frozenset({"CONFIGURING", "POWERING_UP"}) | |
| _owned_nodes_cache: Optional[set[str]] = None |
There was a problem hiding this comment.
[codex] Fixed in 5dbbb59. Added a typed module-level cache containing the state-file version and the per-nodeset owned-node sets.
| def _write_state_unlocked(state: dict[str, Any]) -> None: | ||
| STATE_FILE.parent.mkdir(parents=True, exist_ok=True) | ||
| if not state.get("nodesets"): | ||
| STATE_FILE.unlink(missing_ok=True) | ||
| return | ||
| temporary = STATE_FILE.with_name(f"{STATE_FILE.name}.{os.getpid()}.tmp") | ||
| with temporary.open("w", encoding="utf-8") as stream: | ||
| json.dump(state, stream, indent=2, sort_keys=True) | ||
| stream.flush() | ||
| os.fsync(stream.fileno()) | ||
| os.replace(temporary, STATE_FILE) | ||
|
|
There was a problem hiding this comment.
Invalidate the _owned_nodes_cache cache on any write operation to ensure consistency with the state file.
| def _write_state_unlocked(state: dict[str, Any]) -> None: | |
| STATE_FILE.parent.mkdir(parents=True, exist_ok=True) | |
| if not state.get("nodesets"): | |
| STATE_FILE.unlink(missing_ok=True) | |
| return | |
| temporary = STATE_FILE.with_name(f"{STATE_FILE.name}.{os.getpid()}.tmp") | |
| with temporary.open("w", encoding="utf-8") as stream: | |
| json.dump(state, stream, indent=2, sort_keys=True) | |
| stream.flush() | |
| os.fsync(stream.fileno()) | |
| os.replace(temporary, STATE_FILE) | |
| def _write_state_unlocked(state: dict[str, Any]) -> None: | |
| global _owned_nodes_cache | |
| _owned_nodes_cache = None | |
| STATE_FILE.parent.mkdir(parents=True, exist_ok=True) | |
| if not state.get("nodesets"): | |
| STATE_FILE.unlink(missing_ok=True) | |
| return | |
| temporary = STATE_FILE.with_name(f"{STATE_FILE.name}.{os.getpid()}.tmp") | |
| with temporary.open("w", encoding="utf-8") as stream: | |
| json.dump(state, stream, indent=2, sort_keys=True) | |
| stream.flush() | |
| os.fsync(stream.fileno()) | |
| os.replace(temporary, STATE_FILE) |
There was a problem hiding this comment.
[codex] Fixed in 5dbbb59. Writes invalidate the process-local cache. The read path also fingerprints the atomically replaced state file, which prevents stale cache entries when another controller process writes the file.
| for partition in getattr(lkp.cfg, "partitions", {}).values(): | ||
| partition_nodesets = set(getattr(partition, "partition_nodeset", [])) | ||
| partition_nodesets.update(getattr(partition, "partition_nodeset_dyn", [])) |
There was a problem hiding this comment.
Use defensive programming to guard against None values for partition_nodeset and partition_nodeset_dyn attributes, which can be None if explicitly set to null or left unpopulated.
| for partition in getattr(lkp.cfg, "partitions", {}).values(): | |
| partition_nodesets = set(getattr(partition, "partition_nodeset", [])) | |
| partition_nodesets.update(getattr(partition, "partition_nodeset_dyn", [])) | |
| for partition in getattr(lkp.cfg, "partitions", {}).values(): | |
| partition_nodesets = set(getattr(partition, "partition_nodeset", []) or []) | |
| partition_nodesets.update(getattr(partition, "partition_nodeset_dyn", []) or []) |
There was a problem hiding this comment.
[codex] Fixed in 5dbbb59. Both partition nodeset lists now use an empty-list fallback for explicit null values, with a regression test covering a null static list and populated dynamic list.
| # across all bulkInserts (goes one level above this function) | ||
| failed = _get_failed_instance_inserts(op, util.lookup()) | ||
| if not failed and created == 0: | ||
| errors = op.get("error", {}).get("errors", []) |
There was a problem hiding this comment.
op.get("error", {}) can return None if the "error" key is explicitly set to null in the API response. Calling .get("errors", []) on None will raise an AttributeError. Use (op.get("error") or {}) to safely handle None values.
| errors = op.get("error", {}).get("errors", []) | |
| errors = (op.get("error") or {}).get("errors", []) |
There was a problem hiding this comment.
[codex] Fixed in 5dbbb59. The handler now normalises top-level operation errors once with (op.get("error") or {}), reuses that list for logging and circuit detection, and has a regression test for an explicit null error.
| def state_files(tmp_path, monkeypatch): | ||
| monkeypatch.setattr(capacity_circuit, "STATE_FILE", tmp_path / "state.json") | ||
| monkeypatch.setattr(capacity_circuit, "LOCK_FILE", tmp_path / "state.lock") |
There was a problem hiding this comment.
Clear the _owned_nodes_cache in the state_files fixture to prevent test pollution and ensure tests remain independent.
| def state_files(tmp_path, monkeypatch): | |
| monkeypatch.setattr(capacity_circuit, "STATE_FILE", tmp_path / "state.json") | |
| monkeypatch.setattr(capacity_circuit, "LOCK_FILE", tmp_path / "state.lock") | |
| def state_files(tmp_path, monkeypatch): | |
| monkeypatch.setattr(capacity_circuit, "STATE_FILE", tmp_path / "state.json") | |
| monkeypatch.setattr(capacity_circuit, "LOCK_FILE", tmp_path / "state.lock") | |
| monkeypatch.setattr(capacity_circuit, "_owned_nodes_cache", None) |
There was a problem hiding this comment.
[codex] Fixed in 5dbbb59. The state-file fixture now resets the module cache for every test; additional tests verify one parse per unchanged version and reload after an external atomic replacement.
f019958 to
46d94c6
Compare
|
Hi @reidlevesque Thanks. |
2 participants
Why
A GCE capacity stockout currently marks the requested Slurm nodes DOWN, but the rest of that nodeset remains schedulable. In a partition containing several machine families, Slurm can therefore keep selecting the unavailable family instead of shifting work to another nodeset. Static node weights cannot adapt when any family may stock out.
What changed
capacity_circuit_breakercontroller setting (disabled by default)ResumeTimeout, plus a bounded hard deadline, so in-flight probes are not reset prematurely and stale probes cannot block recovery foreverscontrolcalls and safely ignore targets removed by a nodeset resizeValidation
pre-commit run --files <changed files>pytest: 202 passedmypy .../scripts --check-untyped-defs: 27 files cleanmake gcluster