Skip to content

Commit 6f2dbc0

Browse files
Support For NVCRE Executor (#5254)
Signed-off-by: Smitha Jayaram <sjayaram@nvidia.com> Signed-off-by: Alex Filby <afilby@nvidia.com> Co-authored-by: Alex Filby <afilby@nvidia.com>
1 parent a2ee52c commit 6f2dbc0

7 files changed

Lines changed: 477 additions & 22 deletions

File tree

‎scripts/performance/argument_parser.py‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -681,6 +681,107 @@ def parse_cli_args():
681681
default=None,
682682
)
683683

684+
# Nvcre
685+
nvcre_args = parser.add_argument_group("Nvcre arguments")
686+
nvcre_args.add_argument(
687+
"--nvcre_namespace",
688+
type=str,
689+
help="Kubernetes namespace for Nvcre WorkloadRun. When set, uses the Nvcre executor instead of Slurm.",
690+
required=False,
691+
)
692+
nvcre_args.add_argument(
693+
"--nvcre_image_pull_secret",
694+
type=str,
695+
help="Kubernetes image pull secret name for pulling the container image.",
696+
required=False,
697+
default=None,
698+
)
699+
nvcre_args.add_argument(
700+
"--nvcre_workdir_pvc",
701+
type=str,
702+
help="PVC name for syncing job workdir to the cluster before launch.",
703+
required=False,
704+
default=None,
705+
)
706+
nvcre_args.add_argument(
707+
"--nvcre_workdir_pvc_path",
708+
type=str,
709+
help="Mount path for the workdir PVC inside the training pod.",
710+
default="/nemo_run",
711+
required=False,
712+
)
713+
nvcre_args.add_argument(
714+
"--nvcre_workdir_local_path",
715+
type=str,
716+
help="Local directory rsynced into the workdir PVC before launch.",
717+
required=False,
718+
default=None,
719+
)
720+
nvcre_args.add_argument(
721+
"--nvcre_node_selector_json",
722+
type=str,
723+
help="JSON-encoded dict of node selector labels for targeting specific nodes.",
724+
required=False,
725+
default=None,
726+
)
727+
nvcre_args.add_argument(
728+
"--nvcre_volumes_json",
729+
type=str,
730+
help="JSON-encoded list of Kubernetes Volume dicts.",
731+
required=False,
732+
default=None,
733+
)
734+
nvcre_args.add_argument(
735+
"--nvcre_volume_mounts_json",
736+
type=str,
737+
help="JSON-encoded list of Kubernetes VolumeMount dicts.",
738+
required=False,
739+
default=None,
740+
)
741+
nvcre_args.add_argument(
742+
"--nvcre_timeout_per_job",
743+
type=str,
744+
help="Per-job timeout passed to Nvcre orchestration (e.g. '24h').",
745+
required=False,
746+
default="24h",
747+
)
748+
749+
nvcre_args.add_argument(
750+
"--nvcre_kubeconfig",
751+
type=str,
752+
help="Path to kubeconfig file for nvcrectl/kubectl.",
753+
required=False,
754+
default=None,
755+
)
756+
nvcre_args.add_argument(
757+
"--nvcre_kube_context",
758+
type=str,
759+
help="Kubernetes context to use with nvcrectl/kubectl.",
760+
required=False,
761+
default=None,
762+
)
763+
nvcre_args.add_argument(
764+
"--nvcre_gang_scheduler_name",
765+
type=str,
766+
help="Gang scheduler name for the WorkloadRun (e.g. 'kai-scheduler').",
767+
required=False,
768+
default=None,
769+
)
770+
nvcre_args.add_argument(
771+
"--nvcre_hf_token_secret_name",
772+
type=str,
773+
help="Name of the Kubernetes Secret holding the HuggingFace token.",
774+
required=False,
775+
default=None,
776+
)
777+
nvcre_args.add_argument(
778+
"--nvcre_hf_token_secret_key",
779+
type=str,
780+
help="Key inside the Kubernetes Secret that holds the HuggingFace token value.",
781+
required=False,
782+
default="token",
783+
)
784+
684785
# For performance
685786
performance_args = parser.add_argument_group("Performance arguments")
686787
performance_args.add_argument(

‎scripts/performance/perf_plugins.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,12 @@
3535
from nemo_run import Plugin, Script, SlurmExecutor
3636

3737

38+
try:
39+
from nemo_run.core.execution.nvcre import NvcreExecutor
40+
except ImportError:
41+
NvcreExecutor = None # type: ignore[assignment,misc]
42+
43+
3844
logger: logging.Logger = logging.getLogger(__name__)
3945
NSYS_SQLITE_EXPORT_ARG = "--export=sqlite"
4046

@@ -185,6 +191,9 @@ def setup(self, task: Union["run.Partial", "run.Script"], executor: "run.Executo
185191
if isinstance(executor, SlurmExecutor):
186192
# NOTE: DO NOT change to f-string, `%q{}` is Slurm placeholder
187193
launcher.nsys_filename = "profile_%p_%q{SLURM_JOB_ID}_node%q{SLURM_NODEID}_rank%q{SLURM_PROCID}"
194+
elif NvcreExecutor is not None and isinstance(executor, NvcreExecutor):
195+
# Nvcre pods use PET_* env vars for distributed rank info.
196+
launcher.nsys_filename = "profile_node%q{PET_NODE_RANK}_rank%p"
188197

189198
if self.nsys_gpu_metrics:
190199
if hasattr(launcher, "nsys_gpu_metrics"):

‎scripts/performance/setup_experiment.py‎

Lines changed: 110 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -88,22 +88,31 @@ def _filter_run_script_args(argv: List[str]) -> List[str]:
8888
carry JSON values whose ``{}`` / ``[]`` are brace/glob-expanded by the
8989
shell in the generated launch command, corrupting argv and leaking tokens
9090
into the training entrypoint's Hydra override parser.
91+
* ``--offline`` — controls HF Hub access on the launcher node only; offline
92+
behaviour in the rank-local script is governed by the ``HF_HUB_OFFLINE``
93+
environment variable, not by this flag.
9194
92-
All of these take a value, passed either as ``--flag value`` (two tokens) or
93-
``--flag=value`` (one token).
95+
Value-taking flags are passed as ``--flag value`` or ``--flag=value``.
96+
Boolean flags (no following value) are listed in ``_LAUNCHER_ONLY_BOOL``.
9497
"""
98+
_LAUNCHER_ONLY_BOOL = {"--offline", "--dryrun", "-d"}
9599

96100
def _is_launcher_only(flag: str) -> bool:
97-
return flag in (
98-
"-lgc",
99-
"-lmc",
100-
"-vb",
101-
"--additional_slurm_params",
102-
"--csp",
103-
"--enable_vboost",
104-
"--lock_gpu_freq",
105-
"--peak_mem_clk",
106-
) or flag.startswith("--kubeflow_")
101+
return (
102+
flag
103+
in (
104+
"-lgc",
105+
"-lmc",
106+
"-vb",
107+
"--additional_slurm_params",
108+
"--csp",
109+
"--enable_vboost",
110+
"--lock_gpu_freq",
111+
"--peak_mem_clk",
112+
)
113+
or flag.startswith("--kubeflow_")
114+
or flag.startswith("--nvcre_")
115+
)
107116

108117
filtered_args = []
109118
skip_next = False
@@ -112,6 +121,8 @@ def _is_launcher_only(flag: str) -> bool:
112121
if skip_next:
113122
skip_next = False
114123
continue
124+
if arg in _LAUNCHER_ONLY_BOOL:
125+
continue
115126
if _is_launcher_only(arg.split("=", 1)[0]):
116127
skip_next = "=" not in arg
117128
continue
@@ -472,6 +483,23 @@ def maybe_increase_n_attempts_on_flaky_failure(
472483
return n_attempts
473484

474485

486+
def _nvcre_env_vars(custom_env_vars: Dict[str, str], hf_token_secret_name: Optional[str]) -> Dict[str, str]:
487+
"""Build the plain env-var dict for an NVCRE executor.
488+
489+
When an HF token secret is configured, HF_TOKEN is stripped from the plain
490+
env (it may have arrived via cluster_config environment: or --env flags) and
491+
injected via secret_env_vars instead, so the token never appears as a plain
492+
value in the WorkloadRun manifest. HF connectivity is always enabled so the
493+
pod can reach HuggingFace regardless of any --offline flag passed for
494+
launcher-side setup.
495+
"""
496+
env = custom_env_vars.copy()
497+
if hf_token_secret_name:
498+
env.pop("HF_TOKEN", None)
499+
env.update({"HF_HUB_OFFLINE": "0", "TRANSFORMERS_OFFLINE": "0"})
500+
return env
501+
502+
475503
def main(
476504
use_recipes: bool,
477505
model_family_name: str,
@@ -542,6 +570,21 @@ def main(
542570
kubeflow_container_kwargs_json: Optional[str],
543571
kubeflow_labels_json: Optional[str],
544572
kubeflow_pod_annotations_json: Optional[str],
573+
nvcre_namespace: Optional[str] = None,
574+
nvcre_image_pull_secret: Optional[str] = None,
575+
nvcre_workdir_pvc: Optional[str] = None,
576+
nvcre_workdir_pvc_path: str = "/nemo_run",
577+
nvcre_workdir_local_path: Optional[str] = None,
578+
nvcre_node_selector_json: Optional[str] = None,
579+
nvcre_volumes_json: Optional[str] = None,
580+
nvcre_volume_mounts_json: Optional[str] = None,
581+
nvcre_timeout_per_job: str = "24h",
582+
nvcre_kubeconfig: Optional[str] = None,
583+
nvcre_kube_context: Optional[str] = None,
584+
nvcre_gang_scheduler_name: Optional[str] = None,
585+
nvcre_hf_token_secret_name: Optional[str] = None,
586+
nvcre_hf_token_secret_key: str = "token",
587+
deterministic: bool = False,
545588
config_variant: str | None = None,
546589
gres: Optional[str] = None,
547590
packager: str = "git",
@@ -594,7 +637,7 @@ def main(
594637
# runs. Creating the dir from the launcher would either fail (PVC
595638
# not present) or create a useless dir on the launcher's local FS.
596639
# Let the trainer container create its own dirs on first write.
597-
if kubeflow_namespace is None:
640+
if kubeflow_namespace is None and nvcre_namespace is None:
598641
save_dir_path.mkdir(parents=True, exist_ok=True)
599642
save_dir_mount = f"{save_dir_path}:{save_dir_path}"
600643
if save_dir_mount not in custom_mounts:
@@ -613,7 +656,7 @@ def main(
613656
# Kubeflow the trainer pod runs the image — which ships Megatron-Bridge at
614657
# /opt/Megatron-Bridge — and custom_mounts do not apply, so the launcher's
615658
# /tmp path does not exist in the pod; use the image's script path instead.
616-
if kubeflow_namespace:
659+
if kubeflow_namespace or nvcre_namespace:
617660
in_container_script_dir = "/opt/Megatron-Bridge/scripts/performance"
618661
in_container_script_path = f"{in_container_script_dir}/{script_name}"
619662
else:
@@ -634,11 +677,13 @@ def main(
634677
if peak_mem_clk == -1:
635678
peak_mem_clk = None
636679

637-
if kubeflow_namespace:
680+
if kubeflow_namespace or nvcre_namespace:
638681
if enable_vboost or lock_gpu_freq is not None or peak_mem_clk is not None:
639682
logger.warning(
640-
"--enable_vboost, --lock_gpu_freq, and --peak_mem_clk are Slurm-only and will be ignored on Kubeflow."
683+
"--enable_vboost, --lock_gpu_freq, and --peak_mem_clk are Slurm-only and will be ignored on Kubeflow/NVCRE."
641684
)
685+
686+
if kubeflow_namespace:
642687
executor = kubeflow_executor(
643688
namespace=kubeflow_namespace,
644689
nodes=-(num_gpus // -gpus_per_node),
@@ -670,6 +715,31 @@ def main(
670715
labels=json.loads(kubeflow_labels_json) if kubeflow_labels_json else None,
671716
pod_annotations=(json.loads(kubeflow_pod_annotations_json) if kubeflow_pod_annotations_json else None),
672717
)
718+
elif nvcre_namespace is not None:
719+
try:
720+
from utils.executors import nvcre_executor
721+
except ImportError:
722+
from .utils.executors import nvcre_executor
723+
executor = nvcre_executor(
724+
namespace=nvcre_namespace,
725+
container_image=container_image,
726+
nodes=-(num_gpus // -gpus_per_node),
727+
num_gpus_per_node=gpus_per_node,
728+
image_pull_secret=nvcre_image_pull_secret,
729+
workdir_pvc=nvcre_workdir_pvc,
730+
workdir_pvc_path=nvcre_workdir_pvc_path,
731+
workdir_local_path=nvcre_workdir_local_path,
732+
node_selector=json.loads(nvcre_node_selector_json) if nvcre_node_selector_json else None,
733+
volumes=json.loads(nvcre_volumes_json) if nvcre_volumes_json else None,
734+
volume_mounts=json.loads(nvcre_volume_mounts_json) if nvcre_volume_mounts_json else None,
735+
timeout_per_job=nvcre_timeout_per_job,
736+
kubeconfig=nvcre_kubeconfig,
737+
kube_context=nvcre_kube_context,
738+
gang_scheduler_name=nvcre_gang_scheduler_name,
739+
)
740+
executor.env_vars = _nvcre_env_vars(custom_env_vars, nvcre_hf_token_secret_name)
741+
if nvcre_hf_token_secret_name:
742+
executor.secret_env_vars["HF_TOKEN"] = (nvcre_hf_token_secret_name, nvcre_hf_token_secret_key)
673743
else:
674744
executor = slurm_executor(
675745
gpu=gpu,
@@ -788,15 +858,17 @@ def main(
788858
logger.info("dryrun requested: exiting")
789859
return
790860

791-
job_dir, job_status = get_job_dir_and_status_from_run(exp_name)
792-
793-
terminal_failure = job_status not in ["SUCCEEDED", "SUBMITTED", "PENDING", "RUNNING"]
794-
795861
if detach:
862+
# For detached runs (e.g. Nvcre), skip status polling — the
863+
# caller (llmb-run) polls for completion via nvcrectl.
796864
is_finished_experiment = True
797865
is_testing_passed = True
798866
break
799867

868+
job_dir, job_status = get_job_dir_and_status_from_run(exp_name)
869+
870+
terminal_failure = job_status not in ["SUCCEEDED", "SUBMITTED", "PENDING", "RUNNING"]
871+
800872
log_file_paths = list(Path(f"{job_dir}").glob("log*.out"))
801873
ensure_logs_where_written(log_file_paths)
802874

@@ -990,9 +1062,10 @@ def main(
9901062
task=args.task,
9911063
compute_dtype=args.compute_dtype,
9921064
gpu=args.gpu,
993-
hf_token=args.hf_token,
1065+
hf_token=args.hf_token or os.environ.get("HF_TOKEN"),
9941066
offline=args.offline,
995-
detach=args.detach,
1067+
# Force detach for Nvcre — llmb-run polls for completion via nvcrectl
1068+
detach=True if args.nvcre_namespace else args.detach,
9961069
dryrun=args.dryrun,
9971070
enable_vboost=args.enable_vboost,
9981071
lock_gpu_freq=args.lock_gpu_freq,
@@ -1069,6 +1142,21 @@ def main(
10691142
kubeflow_container_kwargs_json=args.kubeflow_container_kwargs_json,
10701143
kubeflow_labels_json=args.kubeflow_labels_json,
10711144
kubeflow_pod_annotations_json=args.kubeflow_pod_annotations_json,
1145+
nvcre_namespace=args.nvcre_namespace,
1146+
nvcre_image_pull_secret=args.nvcre_image_pull_secret,
1147+
nvcre_workdir_pvc=args.nvcre_workdir_pvc,
1148+
nvcre_workdir_pvc_path=args.nvcre_workdir_pvc_path,
1149+
nvcre_workdir_local_path=args.nvcre_workdir_local_path,
1150+
nvcre_node_selector_json=args.nvcre_node_selector_json,
1151+
nvcre_volumes_json=args.nvcre_volumes_json,
1152+
nvcre_volume_mounts_json=args.nvcre_volume_mounts_json,
1153+
nvcre_timeout_per_job=args.nvcre_timeout_per_job,
1154+
nvcre_kubeconfig=args.nvcre_kubeconfig,
1155+
nvcre_kube_context=args.nvcre_kube_context,
1156+
nvcre_gang_scheduler_name=args.nvcre_gang_scheduler_name,
1157+
nvcre_hf_token_secret_name=args.nvcre_hf_token_secret_name,
1158+
nvcre_hf_token_secret_key=args.nvcre_hf_token_secret_key,
1159+
deterministic=args.deterministic,
10721160
config_variant=config_variant,
10731161
gres=args.gres,
10741162
packager=args.packager,

0 commit comments

Comments
 (0)