fix the parallel number of CI nodes when it is smaller than number of tests (#33276)

* fix the parallel number

* this?

* keep it simple

* woups

* nit

* style

* fix param name

* fix

* fix dtype

* yups

* ???

* ??

* this?

* ????

* no default flow style

* ??

* print config

* ????

* there we go!

* documentation

* update

* remove unwanted file

.circleci/config.yml

@@ -55,9 +55,9 @@ jobs:
                     url="https://circleci.com/api/v2/project/${project_slug}/${job_number}/artifacts"
                     curl -o  test_preparation/artifacts.json ${url}
             - run:
-                name: "Show Artifacts"
+                name: "Prepare pipeline parameters"
                 command: |
-                    cat test_preparation/artifacts.json | jq '.items | map({(.path | split("/")[-1][:-4]): .url}) | add | del(.["generated_config"])' > test_preparation/transformed_artifacts.json
+                    python utils/process_test_artifacts.py 
             
             # To avoid too long generated_config.yaml on the continuation orb, we pass the links to the artifacts as parameters.
             # Otherwise the list of tests was just too big. Explicit is good but for that it was a limitation.
@@ -68,6 +68,8 @@ jobs:
                 
             - store_artifacts:
                 path: test_preparation/transformed_artifacts.json
+            - store_artifacts:
+                path: test_preparation/artifacts.json
             - continuation/continue:
                 parameters:  test_preparation/transformed_artifacts.json
                 configuration_path: test_preparation/generated_config.yml

.circleci/create_circleci_config.py

@@ -113,6 +113,7 @@ class CircleCIJob:
         timeout_cmd = f"timeout {self.command_timeout} " if self.command_timeout else ""
         marker_cmd = f"-m '{self.marker}'" if self.marker is not None else ""
         additional_flags = f" -p no:warning -o junit_family=xunit1 --junitxml=test-results/junit.xml"
+        parallel = f' << pipeline.parameters.{self.job_name}_parallelism >> '
         steps = [
             "checkout",
             {"attach_workspace": {"at": "test_preparation"}},
@@ -152,7 +153,7 @@ class CircleCIJob:
             {"store_artifacts": {"path": "installed.txt"}},
         ]
         if self.parallelism:
-            job["parallelism"] = self.parallelism
+            job["parallelism"] = parallel
         job["steps"] = steps
         return job
 
@@ -359,12 +360,13 @@ def create_circleci_config(folder=None):
             "nightly": {"type": "boolean", "default": False},
             "tests_to_run": {"type": "string", "default": ''},
             **{j.job_name + "_test_list":{"type":"string", "default":''} for j in jobs},
+            **{j.job_name + "_parallelism":{"type":"integer", "default":1} for j in jobs},
         },
         "jobs" : {j.job_name: j.to_dict() for j in jobs},
         "workflows": {"version": 2, "run_tests": {"jobs": [j.job_name for j in jobs]}}
     }
     with open(os.path.join(folder, "generated_config.yml"), "w") as f:
-        f.write(yaml.dump(config, indent=2, width=1000000, sort_keys=False))
+        f.write(yaml.dump(config, sort_keys=False, default_flow_style=False).replace("' << pipeline", " << pipeline").replace(">> '", " >>"))
 
 
 if __name__ == "__main__":

utils/process_test_artifacts.py

@@ -0,0 +1,75 @@
+# coding=utf-8
+# Copyright 2024 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+
+This helper computes the "ideal" number of nodes to use in circle CI.
+For each job, we compute this parameter and pass it to the `generated_config.yaml`.
+"""
+
+import json
+import math
+import os
+
+
+MAX_PARALLEL_NODES = 8  # TODO create a mapping!
+AVERAGE_TESTS_PER_NODES = 5
+
+
+def count_lines(filepath):
+    """Count the number of lines in a file."""
+    try:
+        with open(filepath, "r") as f:
+            return len(f.read().split("\n"))
+    except FileNotFoundError:
+        return 0
+
+
+def compute_parallel_nodes(line_count, max_tests_per_node=10):
+    """Compute the number of parallel nodes required."""
+    num_nodes = math.ceil(line_count / AVERAGE_TESTS_PER_NODES)
+    if line_count < 4:
+        return 1
+    return min(MAX_PARALLEL_NODES, num_nodes)
+
+
+def process_artifacts(input_file, output_file):
+    # Read the JSON data from the input file
+    with open(input_file, "r") as f:
+        data = json.load(f)
+
+    # Process items and build the new JSON structure
+    transformed_data = {}
+    for item in data.get("items", []):
+        if "test_list" in item["path"]:
+            key = os.path.splitext(os.path.basename(item["path"]))[0]
+            transformed_data[key] = item["url"]
+            parallel_key = key.split("_test")[0] + "_parallelism"
+            file_path = os.path.join("test_preparation", f"{key}.txt")
+            line_count = count_lines(file_path)
+            transformed_data[parallel_key] = compute_parallel_nodes(line_count)
+
+    # Remove the "generated_config" key if it exists
+    if "generated_config" in transformed_data:
+        del transformed_data["generated_config"]
+
+    # Write the transformed data to the output file
+    with open(output_file, "w") as f:
+        json.dump(transformed_data, f, indent=2)
+
+
+if __name__ == "__main__":
+    input_file = "test_preparation/artifacts.json"
+    output_file = "test_preparation/transformed_artifacts.json"
+    process_artifacts(input_file, output_file)