Repository navigation
flex mig support for static nodes - #6371
Conversation
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 adds support for machine-type fallback (flexibility policy) for static MIG NodeSets in the Slurm on GCP toolkit. By allowing users to specify a ranked list of machine types, the system can automatically fall back to alternative shapes when primary capacity is unavailable. The implementation includes robust validation to prevent incompatible configurations and updates to the Slurm configuration generation to maintain consistency across the different machine shapes. 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
|
There was a problem hiding this comment.
Code Review
This pull request introduces support for instance_flexibility_policy (machine-type fallback) in MIG NodeSets across the Slurm GCP v6 modules (nodeset, partition, and controller). It adds the necessary Terraform variables, outputs, and validation rules, updates the Python configuration scripts to floor node definitions to the smallest shape in the flexibility policy, and extends the Go configuration package to parse and include fallback machine types. As there are no review comments provided, I have no additional feedback to offer.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces an opt-in machine-type fallback policy (instance_flexibility_policy) for MIG NodeSets across the Slurm GCP v6 modules, allowing MIGs to fall back to alternative machine types when capacity is unavailable. The changes span Terraform configurations, Go configuration generation, and Python controller scripts, including new validation rules and unit tests. Feedback on the changes suggests: (1) refining error handling in pkg/config/machine_configs.go to avoid silently ignoring non-404 errors from GetMachineType, (2) avoiding autovivification on cached NSDict templates in util.py by using .get(), and (3) ensuring consistent defensive handling of the nodeset variable in mig_flex.py to prevent potential AttributeError crashes.
6c9948e to
6090909
Compare
|
Few other Edge cases that might arise:
|
6090909 to
e814eac
Compare
|
Replying to above top level comment:
|
Verification & Test Summary1. Unit Tests & Static Validation
2. Terraform Plan & Negative Guardrail Matrix (
|
| Partition | Provisioning Engine / Mode | Configuration | Live Verification Result |
|---|---|---|---|
mig |
MIG (BALANCED) |
2 static (n2-standard-2) |
Booted idle; srun -N 2 -p mig PASS |
migzone |
MIG (ANY_SINGLE_ZONE) |
2 static + enable_spot_vm |
Co-located in single zone; srun -N 2 -p migzone PASS |
flex |
MIG + instance_flexibility_policy (BALANCED) |
3 static; Rank 0: [n2-standard-2], Rank 1: [n2d-standard-2, e2-standard-2], Rank 2: [n1-standard-2] |
Organic GCE zone stockout triggered live fallback (2× n2-standard-2 + 1× n1-standard-2); cloud.conf floored RealMemory=7070 CPUs=1; all 3 registered idle (0 Low RealMemory drain) and ran srun -N 3 -p flex PASS |
flexzone |
MIG + instance_flexibility_policy (ANY_SINGLE_ZONE) |
2 static; Rank 0: [n2-standard-2], Rank 1: [n2d-standard-2] |
Both nodes fell back to Rank 1 (2× n2d-standard-2) co-located in a single zone; srun -N 2 -p flexzone PASS |
flexspot |
MIG + instance_flexibility_policy (ANY + SPOT) |
2 static; Rank 0: [n2-standard-2, n2d-standard-2], Rank 1: [e2-standard-2, n1-standard-2] |
Provisioned as SPOT VMs; srun -N 2 -p flexspot PASS |
bulk |
BULK_INSERT (static) |
2 static (n2-standard-2) |
Booted idle; srun -N 2 -p bulk PASS |
burst |
BULK_INSERT (static + dynamic) |
1 static + 2 dynamic |
Static + dynamic burst scale-up; srun -N 2 -p burst PASS |
debug |
Default AUTO (dynamic) |
0 static + 2 dynamic |
Dynamic scale-up via BulkInsert & idle scale-down; srun -N 2 -p debug PASS |
auto |
Explicit AUTO (dynamic) |
0 static + 2 dynamic |
Dynamic scale-up via BulkInsert & idle scale-down; srun -N 1 -p auto PASS |
Post-Deployment Lifecycle & Corner Cases Verified on Live Cluster
-
Heterogeneous Fallback Scheduling & Queueing: Submitted 5
--exclusivejobs and a 12-job burst flood against the 3-node heterogeneousflexpartition (2× n2-standard-2 + 1× n1-standard-2); all jobs queued and completed (ExitCode 0:0) with0node drains. -
Fault Handling: Verified rapid
scancel, multi-node (-N 3) mid-flightscancel, task crash (kill -9), walltime enforcement (--time=00:01:00$\rightarrow$ TIMEOUT), andscontrolDRAIN/RESUMErouting onflexnodes. -
MIG Auto-Healing &
slurmsync.pyRecovery: Triggeredgcloud compute instance-groups managed recreate-instanceson a liveflexfallback VM; verifiedslurmsync.pytransitioned the node toDOWN+CLOUD (MIG Auto-Healing instance repair in progress)and automatically resumed it toIDLEonce the recreated VM rejoined. -
Scale-Down (
suspend.py) & Clean Teardown: Verified automaticSuspendTimescale-down on dynamic partitions, manualPOWER_DOWNinstance deletion across all 5 Compute MIGs, and cleangcluster destroy(137resources destroyed,0orphaned MIGs, instance templates, resource policies, or VMs).
CI Triage: Failing / Stuck Checks (Unrelated to PR Changes)None of the 5 failing or stuck CI checks below reached the Slurm modules (
|
feat(slurm): MIG Phase 3 — Flexible MIGs (Machine-Type Fallback & Multi-Zone Scavenging) for Static Nodes
Overview
This PR implements MIG Phase 3, adding support for Flexible Managed Instance Groups (
instance_flexibility_policy) and Multi-Zone Capacity Scavenging & Cross-Zone Auto-Repair (zone_target_shape) on static Slurm MIG NodeSets (provisioning_engine = "MIG",node_count_dynamic_max = 0), along with lifecycle hardening for MIG auto-repair inslurmsync.pyand teardown incleanup_compute.sh.Key Changes
1. Ranked Machine-Type Fallback (
instance_flexibility_policy)instance_flexibility_policy.instance_selections(rank,machine_types, and optionalnamedefaulting to"selection-<idx+1>") acrossschedmd-slurm-gcp-v6-nodeset,schedmd-slurm-gcp-v6-partition, andschedmd-slurm-gcp-v6-controller(partition.tf)."primary-default"selection containingvar.machine_typewhenvar.machine_typeis not already listed ininstance_selections.instance_flexibility_policyand setsinstance_lifecycle_policy.force_update_on_repair = "YES"ongoogle_compute_region_instance_group_manager.nodeset_mig.2. Slurm Resource Flooring (
cloud.conf) & Blueprint Machine Configs (pkg/config)DRAINon Smaller Fallback VMs:Lookup.nodeset_machine_conf()inutil.py(used byconf.py) computes the floor acrossvar.machine_typeand all candidate fallbackmachine_typesforcpus,memory,sockets, andcores_per_socket(including heterogeneous SMT / non-SMT families liken4a,c4a,t2d). If any fallback shape cannot be resolved, it fails loudly with aRuntimeErrorrather than emitting an unflooredcloud.conf.pkg/config/machine_configs.go): ExtendsGetOutputConfigto query GCEMachineTypes.Getfor all fallback shapes ininstance_flexibility_policyand merge their CPU/GPU entries intomachine_configs. Skips404 Not Found(when a fallback shape is only available in another zone of the region) while failing fast on non-404 errors (403,429, network failures).3. Multi-Zone Capacity Scavenging & Cross-Zone Auto-Repair (
zone_target_shape)zone_target_shape(ANY_SINGLE_ZONEdefault,ANY,BALANCED) intodistribution_policy_target_shapeonnodeset_mig. When a multi-zone shape (ANYorBALANCED) is set on a static MIG without explicitzonesor reservations,zones_allowautomatically expands to all available zones in the region.on_repair.allow_changing_zone = "YES"andforce_update_on_repair = "YES"whenzone_target_shape != "ANY_SINGLE_ZONE", allowing preempted or failed VMs to be recreated in another zone with available capacity.zoneswithdws_flex.enabled = true, while enforcingzone_target_shape == "ANY_SINGLE_ZONE"for DWS Flex MIGs (mig_flex.py,metadata.yaml,outputs.tf) because GCEresizeRequestsrequireANY_SINGLE_ZONE.4. MIG Lifecycle & Teardown Hardening (
slurmsync.py&cleanup_compute.sh)slurmsync.pyAuto-Repair Alignment: When a MIG-backed VM entersTERMINATED,slurmsync.pymarks the Slurm nodeDOWN(Reason="MIG Auto-Healing instance repair in progress") and waits for GCE MIG auto-repair (delete+insert) instead of callingstart_instances.cleanup_compute.shDestroy Race Fix: Skips Terraform-managed<cluster>-<nodeset>-mig-<N>MIGs and their member VMs (metadata.items.created-by.basename()) socleanup_compute.shdoes not racegoogle_compute_region_instance_group_manager.nodeset_migdeletion duringterraform destroy.5. Fail-Fast Guardrails (GCE Flex MIG Platform Constraints)
Enforced at both
gcluster create(metadata.yaml) andterraform plan(outputs.tf/variables.tf):provisioning_engine = "MIG"andnode_count_dynamic_max = 0.instance_flexibility_policywithdws_flex.enabled = true,accelerator_topology, specific reservations (reservation_name/future_reservation), ora3-ultragpu/a4/a4xfamilies (unsupported by GCE'sinstanceFlexibilityPolicyAPI).x86_64vsArm64), matching attached GPU model/count, disk-type compatibility across all fallback shapes, unique selection names, and GCE's limit of<= 10distinct machine types.zone_target_shape == "ANY_SINGLE_ZONE"whenenable_placement = trueoraccelerator_topologyis configured on a MIG.Testing & Verification Matrix
1. Unit Tests & Static Pre-Flight Analysis
245 passed(test_util.py,test_resume.py,test_suspend.py,test_conf.py,test_repair.py,test_topology.py).0 issues foundacross28source files (util.py,slurmsync.py,mig_flex.py,conf.py)../pkg/config/...and./pkg/validators/...(includingTestExtractFallbackMachineTypes,TestAddMachineTypeToOutputConfig_MergesFallback, andTestGetOutputConfig_FallbackLookupErrors).terraform_fmt,terraform_tflint,terraform_validate,terraform-readme,go-cyclo,golangci-lint,shfmt,shellcheck,yamllint,PyMarkdown,codespell).2. 66-Case Terraform Plan & Validation Matrix
instance_flexibility_policyrank ordering, auto-prependedprimary-default,ANY/BALANCEDmulti-zone expansion,allow_changing_zone = "YES", DWS Flex withzones, and rejection of all invalid combinations (specific reservations,a3-ultragpu/a4/a4x, mismatched GPU/arch/disk types, >10 shapes).BULK_INSERT,AUTO, and Phase 2a4x1x72slice MIG plan cases pass unchanged.3. Live GCP End-to-End Verification & Teardown Audit
us-central1withBALANCEDmulti-zone Flex MIGs (n2-standard-4rank 1,n2d-standard-2rank 2),BULK_INSERTnodesets, and multi-zone DWS Flex nodesets.BALANCEDmulti-zone instance distribution acrossus-central1-aandus-central1-b.cloud.conffloored node definitions (CPUs=1 RealMemory=7552) and cleansrunjob execution across all partitions.slurmsync.pymarked the nodeDOWNwithout callingstart_instances, GCE MIG auto-repaired the VM, andslurmsync.pyautomatically restored the node toIDLEin ~28s.gcluster destroy(52 destroyed); verified0residual MIGs,0instances, and0resource policies.Impact & Backward Compatibility
BULK_INSERTWorkloads:instance_flexibility_policydefaults tonullandzone_target_shapedefaults to"ANY_SINGLE_ZONE". StandardAUTOandBULK_INSERTnodesets behave identically to prior releases.1x72slice MIG blueprints require no configuration changes.