Merge pull request #788 from douglasjacobsen/py-nemo-custom-string

py-nemo: Add custom string injection into workloads
GoogleCloudPlatform · Dec 5, 2024 · 1737954 · 1737954
2 parents 0dd2f36 + 57ec33f
commit 1737954
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/var/ramble/repos/builtin/applications/py-nemo/application.py b/var/ramble/repos/builtin/applications/py-nemo/application.py
@@ -43,6 +43,7 @@ class PyNemo(ExecutableApplication):
     executable(
         "pretraining_exec",
         'bash -c "cd /opt/NeMo; git rev-parse HEAD; '
+        "{custom_injected_string}; "
         "python3 -u /opt/NeMo/examples/nlp/language_modeling/megatron_gpt_pretraining.py "
         '--config-path={nemo_generated_config_path} --config-name={nemo_generated_config_name}"',
         use_mpi=True,
@@ -76,6 +77,13 @@ class PyNemo(ExecutableApplication):
     workload_group("pretraining", workloads=["pretraining"])
     all_workloads = ["pretraining"]
 
+    workload_variable(
+        "custom_injected_string",
+        default="",
+        description="Custom string to inject before execution NeMo workload",
+        workload_group="all_workloads",
+    )
+
     workload_variable(
         "model_inputs",
         default="{workload_input_dir}/{nemo_stage}/{nemo_model}",