Adapt training with plotting

.gitignore

@@ -3,7 +3,7 @@ __pycache__/
 *.py[cod]
 *$py.class
 *.pkl
-
+results/*
 # C extensions
 *.so
 *.html

docker/benchmark_tests/Dockerfile.run_test

@@ -10,8 +10,8 @@ COPY src/base ./src/base
 COPY src/train ./src/train
 COPY config.yaml .
 COPY docker/benchmark_tests .
-COPY data ./data
+COPY data/dgta ./data/dgta
 
 RUN rm -rf /root/.cache
 
-CMD [ "python", "run_test.py"]
+CMD [ "python", "run_test.py" ]

docker/benchmark_tests/run_test.py

@@ -1,3 +1,4 @@
+import argparse
 import datetime
 import ipaddress
 import os
@@ -27,6 +28,7 @@ def __init__(self, data_base_path: str = "./data"):
         datasets = DatasetLoader(base_path=data_base_path, max_rows=10000)
 
         dataset = Dataset(
+            name="",
             data_path="",
             data=pl.concat(
                 [
@@ -262,7 +264,32 @@ def __init__(self, length_in_min: float | int, msg_per_sec: int = 500):
         super().__init__(full_length_in_min=length_in_min, msg_per_sec=msg_per_sec)
 
 
-def main(test_type_nr):
+def main():
+    # Get the environment variable, default to 1 if not set
+    env_test_type_nr = int(os.getenv("TEST_TYPE_NR", 1))
+
+    parser = argparse.ArgumentParser(
+        description="Example script with test_type_nr argument."
+    )
+
+    parser.add_argument(
+        "--test_type_nr",
+        type=int,
+        choices=[1, 2, 3, 4],
+        default=env_test_type_nr,
+        help="""
+        1 - Ramp-up test
+        2 - Burst test
+        3 - Maximum throughput test
+        4 - Long-term test
+        """,
+    )
+
+    args = parser.parse_args()
+
+    print(f"Selected test type number: {args.test_type_nr}")
+    test_type_nr = args.test_type_nr
+
     """Creates the test instance and executes the test."""
     match test_type_nr:
         case 1:
@@ -300,10 +327,5 @@ def main(test_type_nr):
 
 
 if __name__ == "__main__":
-    """
-    1 - Ramp-up test
-    2 - Burst test
-    3 - Maximum throughput test
-    4 - Long-term test
-    """
-    main(3)
+    """ """
+    main()

docker/docker-compose-swarm/docker-compose.swarm-run_test.yml

@@ -1,6 +1,8 @@
 services:
   benchmark_test_run:
-    image: benchmark_test_run
+    image: localhost:5000/benchmark_test_run
+    environment:
+      - TEST_TYPE_NR=3
     networks:
       heidgaf:
     deploy:

requirements/requirements.train.txt

@@ -7,3 +7,8 @@ pyarrow
 polars
 click
 optuna
+matplotlib
+te2rules
+seaborn
+lightgbm
+imblearn

src/train/__init__.py

@@ -0,0 +1,3 @@
+RESULT_FOLDER = "results"
+CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"], show_default=True)
+SEED = 108