Merge branch 'ko3n1g/ci/fixes-to-jet' into 'main'

tests: Minor improvements to JET

See merge request ADLR/megatron-lm!2133

Author: ko3n1g

.gitlab/stages/02.functional-tests.yml

@@ -41,13 +41,14 @@ jet-generate:
           "--run-name"
           $FUNCTIONAL_TEST_NAME
           "--wandb-experiment"
-          "test"
+          $(echo $FUNCTIONAL_TEST_NAME | tr '/' '-')
         )
       else
         RELEASE_ARGS=()
       fi
 
     - |
+      export PYTHONPATH=$(pwd)
       python tests/functional_tests/python_test_utils/jet/generate_jet_trigger_job.py \
         --scope $FUNCTIONAL_TEST_SCOPE \
         --a100-cluster $A100_CLUSTER \

tests/functional_tests/python_test_utils/jet/common.py

@@ -86,9 +86,26 @@ def filter_by_scope(
     return workload_manifests
 
 
+def filter_by_model(
+    workload_manifests: List[jetclient.JETWorkloadManifest], model: str
+) -> List[jetclient.JETWorkloadManifest]:
+    """Returns all workload with matching model."""
+    workload_manifests = list(
+        workload_manifest
+        for workload_manifest in workload_manifests
+        if workload_manifest.spec.model == model
+    )
+
+    if len(workload_manifests) == 0:
+        raise ValueError("No test_case found!")
+
+    return workload_manifests
+
+
 def load_workloads(
     container_tag: str,
     scope: Optional[str] = None,
+    model: Optional[str] = None,
     test_case: Optional[str] = None,
     container_image: Optional[str] = None,
 ) -> List[jetclient.JETWorkloadManifest]:
@@ -106,6 +123,9 @@ def load_workloads(
     if scope:
         workloads = filter_by_scope(workload_manifests=workloads, scope=scope)
 
+    if model:
+        workloads = filter_by_model(workload_manifests=workloads, model=model)
+
     if test_case:
         workloads = [filter_by_test_case(workload_manifests=workloads, test_case=test_case)]
 

tests/functional_tests/python_test_utils/jet/generate_jet_trigger_job.py

@@ -47,6 +47,7 @@ def main(
             raise ValueError(f"Platform {test_case.spec.platforms} unknown")
 
         script = [
+            "export PYTHONPATH=$(pwd); "
             "python tests/functional_tests/python_test_utils/jet/launch_jet_workload.py",
             f"--model {test_case.spec.model}",
             f"--test-case {test_case.spec.test_case}",

tests/functional_tests/python_test_utils/jet/generate_local_jobs.py

@@ -0,0 +1,49 @@
+import pathlib
+from typing import Optional
+
+import click
+import jetclient
+import yaml
+
+from tests.functional_tests.python_test_utils.jet import common
+
+
+def load_script(config_path: str) -> str:
+    with open(config_path) as stream:
+        try:
+            jetclient.JETWorkloadManifest(**yaml.safe_load(stream)).spec.script
+        except yaml.YAMLError as exc:
+            raise exc
+
+
+@click.command()
+@click.option("--model", required=False, type=str, help="Filters all tests by matching model")
+@click.option("--scope", required=False, type=str, help="Filters all tests by matching scope")
+@click.option(
+    "--test-case", required=False, type=str, help="Returns a single test-case with matching name."
+)
+@click.option("--output-path", required=True, type=str, help="Path to write jobs to")
+def main(model: Optional[str], scope: Optional[str], test_case: Optional[str], output_path: str):
+    workloads = common.load_workloads(
+        container_image='none', scope=scope, model=model, test_case=test_case, container_tag='none'
+    )
+
+    for workload in workloads:
+        if workload.type == "build":
+            continue
+        magic_values = dict(workload.spec)
+        magic_values["assets_dir"] = "."
+
+        file_path = (
+            pathlib.Path(output_path)
+            / "test_cases"
+            / workload.spec.model
+            / f"{workload.spec.test_case}.sh"
+        )
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(file_path, "w", encoding="utf-8") as fh:
+            fh.write(workload.spec.script.format(**magic_values))
+
+
+if __name__ == "__main__":
+    main()