Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Enable support for Kueue for GKETPUJob #623

Merged
merged 1 commit into from
Aug 13, 2024
Merged

Enable support for Kueue for GKETPUJob #623

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Enable support for Kueue for GKETPUJob
  • Loading branch information
@samos123
samos123 committed Aug 3, 2024
commit 24aa251f62f2fabdc756368ef94a76981370be21
32 changes: 18 additions & 14 deletions axlearn/cloud/gcp/job.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -313,13 +313,15 @@ class Config(GCPJob.Config):
gcsfuse_mount: Optional configs for the GCS FUSE sidecar and volume mount.
See `GCSFuseMount` for details.
enable_pre_provisioner: Whether to enable pre-provisioner.
queue: The Kueue LocalQueue to use. If not set, no queue is used.
"""

env_vars: Dict[str, str] = {}
namespace: str = "default"
gcsfuse_mount: Optional[GCSFuseMount] = None
# This config is made Optional for backwards compatibility. Default is False.
enable_pre_provisioner: Optional[bool] = None
queue: Optional[str] = None

@classmethod
def define_flags(cls, fv: flags.FlagValues):
Expand All @@ -333,6 +335,12 @@ def define_flags(cls, fv: flags.FlagValues):
"GCS FUSE mount spec in the format key=value.",
flag_values=fv,
)
flags.DEFINE_string(
"queue",
None,
"The name of the Kueue LocalQueue to use. If not set, no queue is used.",
flag_values=fv,
)

@classmethod
def from_flags(cls, fv: flags.FlagValues, **kwargs) -> Config:
Expand Down Expand Up @@ -610,15 +618,19 @@ def _build_jobset(self) -> Nested[Any]:
"""
cfg: TPUGKEJob.Config = self.config

annotations = {
# The exclusive topology annotation will ensure that all Pods will have affinity
# rules added that will ensure that they are fully scheduled on the same
# pod-slice node-pools.
"alpha.jobset.sigs.k8s.io/exclusive-topology": "cloud.google.com/gke-nodepool",
}
if cfg.queue:
annotations["kueue.x-k8s.io/queue-name"] = cfg.queue

return dict(
metadata=dict(
name=cfg.name,
annotations={
# The exclusive topology annotation will ensure that all Pods will have affinity
# rules added that will ensure that they are fully scheduled on the same
# pod-slice node-pools.
"alpha.jobset.sigs.k8s.io/exclusive-topology": "cloud.google.com/gke-nodepool",
},
annotations=annotations,
),
spec=dict(
failurePolicy=dict(maxRestarts=cfg.max_tries - 1),
Expand Down Expand Up @@ -668,23 +680,15 @@ class Config(GKEJob.Config):

Attributes:
accelerator: GPU configuration.
queue: The Kueue LocalQueue to use. If not set, no queue is used.
"""

accelerator: AcceleratorConfig = AcceleratorConfig()
queue: Optional[str] = None

@classmethod
def define_flags(cls, fv: flags.FlagValues):
super().define_flags(fv)
common_kwargs = dict(flag_values=fv, allow_override=True)
accelerator_flags(**common_kwargs)
flags.DEFINE_string(
"queue",
None,
"The name of the Kueue LocalQueue to use. If not set, no queue is used.",
**common_kwargs,
)

@classmethod
def from_flags(cls, fv: flags.FlagValues, **kwargs) -> Config:
Expand Down