EIR/docs/doc_modules/user_guides/a_streaming_hands_on.py at 0c48e1fe308eb80f0100e0857f929c1d335c2287 · arnor-sigurdsson/EIR · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
import os
import subprocess
import time
from collections.abc import Sequence
from pathlib import Path

from docs.doc_modules.experiments import AutoDocExperimentInfo, run_capture_and_save
from eir.setup.config_setup_modules.config_setup_utils import load_yaml_config
from eir.utils.logging import get_logger

logger = get_logger(name=__name__)

CONTENT_ROOT = "user_guides"
TUTORIAL_NAME = "01_streaming_data"


def run_with_server(command: list[str]) -> Path:
    server_path = f"docs.doc_modules.{CONTENT_ROOT}.simulation_streamer"

    globals_file = next(i for i in command if "globals" in i)
    globals_dict = load_yaml_config(config_path=globals_file)
    run_folder = Path(globals_dict["basic_experiment"]["output_folder"])

    output_folder_injected = tuple(i for i in command if ".output_folder=" in i)
    if output_folder_injected:
        assert len(output_folder_injected) == 1
        output_folder_inject_string = output_folder_injected[0]
        run_folder = Path(output_folder_inject_string.split(".output_folder=")[-1])

    if run_folder.exists():
        return run_folder

    cur_env = os.environ.copy()
    cur_env["SEQUENCE_LENGTH"] = "64"

    server_process = subprocess.Popen(
        ["python", "-m", server_path],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        env=cur_env,
    )
    logger.info("Server process starting...")

    try:
        time.sleep(10)
        subprocess.run(args=command, check=True)
    finally:
        server_process.terminate()
        server_process.wait(timeout=5)

    return run_folder


def get_streaming_generation_experiment() -> AutoDocExperimentInfo:
    base_path = f"docs/tutorials/tutorial_files/{CONTENT_ROOT}/{TUTORIAL_NAME}"
    conf_output_path = f"eir_tutorials/{CONTENT_ROOT}/{TUTORIAL_NAME}"

    command = [
        "eirtrain",
        "--global_configs",
        f"{conf_output_path}/globals.yaml",
        "--input_configs",
        f"{conf_output_path}/input.yaml",
        "--fusion_configs",
        f"{conf_output_path}/fusion.yaml",
        "--output_configs",
        f"{conf_output_path}/output.yaml",
    ]

    mapping = [
        (
            "training_curve_LOSS",
            "figures/training_curve_LOSS.pdf",
        ),
        (
            "samples/500/auto/0_generated.txt",
            "figures/auto_generated_iter_500.txt",
        ),
        (
            "samples/2500/auto/0_generated.txt",
            "figures/auto_generated_iter_2500.txt",
        ),
    ]

    get_tutorial_folder = (
        run_capture_and_save,
        {
            "command": [
                "tree",
                f"eir_tutorials/{CONTENT_ROOT}/{TUTORIAL_NAME}",
                "-L",
                "3",
                "--noreport",
            ],
            "output_path": Path(base_path) / "commands/tutorial_folder.txt",
        },
    )

    ade = AutoDocExperimentInfo(
        name="STREAMING_SEQUENCE_GENERATION",
        data_url=None,
        data_output_path=None,
        base_path=Path(base_path),
        conf_output_path=Path(conf_output_path),
        command=command,
        files_to_copy_mapping=mapping,
        post_run_functions=(get_tutorial_folder,),
        run_command_wrapper=run_with_server,
    )

    return ade


def get_experiments() -> Sequence[AutoDocExperimentInfo]:
    exp_1 = get_streaming_generation_experiment()
    return [exp_1]