@@ -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+
475503def 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