Merge pull request #137 from arnor-sigurdsson/more-streaming-docs

More streaming docs

Author: arnor-sigurdsson

docs/automake_docs.py

@@ -44,6 +44,7 @@
     AutoDocServingInfo,
     make_serving_tutorial_data,
 )
+from docs.doc_modules.user_guides import a_streaming_hands_on
 from docs.doc_scripts.doc_performance_report import (
     collect_performance_data,
     generate_report,
@@ -159,6 +160,14 @@ def get_i_scaling_experiments() -> Iterable[AutoDocExperimentInfo]:
     )
 
 
+def get_user_guide_experiments() -> Iterable[AutoDocExperimentInfo]:
+    a_experiments = a_streaming_hands_on.get_experiments()
+
+    return chain(
+        a_experiments,
+    )
+
+
 def generate_performance_report(root_dir, output_dir):
     if not os.path.exists(root_dir):
         logger.error(f"Error: Root directory '{root_dir}' does not exist.")
@@ -223,7 +232,9 @@ def parse_args():
     args = parse_args()
 
     if args.experiments is None:
-        run_a = run_b = run_c = run_d = run_e = run_f = run_g = run_h = run_i = True
+        run_a = run_b = run_c = run_d = run_e = run_f = run_g = run_h = run_i = (
+            run_ug
+        ) = True
     else:
         experiment_ids = args.experiments.split(",")
         run_a = "a" in experiment_ids
@@ -235,6 +246,7 @@ def parse_args():
         run_g = "g" in experiment_ids
         run_h = "h" in experiment_ids
         run_i = "i" in experiment_ids
+        run_ug = "ug" in experiment_ids
 
     experiments = []
 
@@ -274,6 +286,10 @@ def parse_args():
         i_scaling_experiments = get_i_scaling_experiments()
         experiments.append(i_scaling_experiments)
 
+    if run_ug:
+        user_guide_experiments = get_user_guide_experiments()
+        experiments.append(user_guide_experiments)
+
     experiment_iter = chain.from_iterable(experiments)
     for experiment in experiment_iter:
         match experiment:

docs/doc_modules/user_guides/__init__.py

@@ -0,0 +1,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]

docs/doc_modules/user_guides/a_streaming_hands_on.py

@@ -0,0 +1,85 @@
+import os
+from typing import Any
+
+from docs.doc_modules.user_guides.single_sample_simulation import simulate_health_sample
+from eir.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+class DataSimulator:
+    def __init__(
+        self,
+        sequence_length: int = 64,
+        max_iterations: int | None = None,
+    ):
+        self.sequence_length = sequence_length
+        self.max_iterations = max_iterations
+        self.position = 0
+
+        logger.info(f"Using sequence_length={self.sequence_length}")
+        if self.max_iterations is not None:
+            logger.info(f"Will terminate after {self.max_iterations} iterations")
+
+    def reset(self):
+        self.position = 0
+        if self.max_iterations is not None:
+            logger.info(f"Reset: Will terminate after {self.max_iterations} iterations")
+
+    def get_batch(self, batch_size: int) -> list[dict[str, Any]]:
+        if self.max_iterations is not None and self.position >= self.max_iterations:
+            logger.info(f"Reached max iterations ({self.max_iterations}), terminating")
+            return []
+
+        batch = []
+        for _ in range(batch_size):
+            if self.max_iterations is not None and self.position >= self.max_iterations:
+                break
+
+            text_input, text_sequence = simulate_health_sample(sequence_length=64)
+            sample_id = f"sample_{self.position}"
+
+            batch.append(
+                {
+                    "inputs": {
+                        "text_output": text_sequence,
+                        "text_input": text_input,
+                    },
+                    "target_labels": {"text_output": {"text_output": text_sequence}},
+                    "sample_id": sample_id,
+                }
+            )
+
+            self.position += 1
+
+        return batch
+
+    def get_status(self) -> dict[str, Any]:
+        status_data = {
+            "current_position": self.position,
+        }
+
+        if self.max_iterations is not None:
+            status_data["max_iterations"] = self.max_iterations
+            status_data["remaining_iterations"] = max(
+                0, self.max_iterations - self.position
+            )
+
+        return status_data
+
+
+def create_simulator() -> DataSimulator:
+    sequence_length = int(os.getenv("SEQUENCE_LENGTH", "512"))
+
+    max_iterations_str = os.getenv("MAX_ITERATIONS")
+    max_iterations = int(max_iterations_str) if max_iterations_str else None
+
+    logger.info(
+        f"Creating DataSimulator with sequence_length={sequence_length}, "
+        f"max_iterations={max_iterations}"
+    )
+
+    return DataSimulator(
+        sequence_length=sequence_length,
+        max_iterations=max_iterations,
+    )